1 files changed, 77 insertions, 37 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index d01be66..ae74014 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -9,34 +9,59 @@ let raise_unsupported_ast (type a) (f : Format.formatter -> a -> unit) (x : a) =
 
 let raise_unsupported_expr = raise_unsupported_ast (Printast.expression 0)
 
+type 'a index = { index : int } [@@deriving show { with_path = false }]
+
+let get (arr : 'a Arraylist.t) (i : 'a index) : 'a = Arraylist.get arr i.index
+
 type prim = [ `Add ] [@@deriving show { with_path = false }]
 
 type expr =
-  | App of expr * expr
+  | App of expr index * expr index
   | Int of int
-  | Lam of string * expr
+  | Lam of string * expr index
   | Prim of prim
   | Var of string
 [@@deriving show { with_path = false }]
 
-let rec expr_of_parsetree (expr : Parsetree.expression) : expr =
-  match expr.pexp_desc with
-  | Pexp_apply (f, xs) ->
-      List.fold_left
-        (fun f -> function
-          | Asttypes.Nolabel, x -> App (f, expr_of_parsetree x)
-          | _ -> raise_unsupported_expr expr)
-        (expr_of_parsetree f) xs
-  | Pexp_constant (Pconst_integer (n, None)) -> Int (int_of_string n)
-  | Pexp_fun (Nolabel, None, arg_pat, body) -> (
-      match arg_pat.ppat_desc with
-      | Ppat_var { txt = name; _ } -> Lam (name, expr_of_parsetree body)
-      | _ -> raise_unsupported_expr expr)
-  | Pexp_ident { txt = Lident "+"; _ } -> Prim `Add
-  | Pexp_ident { txt = Lident name; _ } -> Var name
-  | _ -> raise_unsupported_expr expr
-
-let rec parsetree_of_expr : expr -> Parsetree.expression =
+type 'a ast = { subexprs : expr Arraylist.t; mutable root : 'a index }
+[@@deriving show { with_path = false }]
+
+let get_subexpr (ast : 'a ast) : expr index -> expr = get ast.subexprs
+
+let new_ast () =
+  { subexprs = Arraylist.make 0 (Var "<placeholder>"); root = { index = 0 } }
+
+let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
+  let add (expr : expr) : expr index =
+    let index = Arraylist.length ast.subexprs in
+    Arraylist.push ast.subexprs expr;
+    { index }
+  in
+  let rec loop (expr : Parsetree.expression) : expr index =
+    match expr.pexp_desc with
+    | Pexp_apply (f, xs) ->
+        List.fold_left
+          (fun f -> function
+            | Asttypes.Nolabel, x -> add (App (f, loop x))
+            | _ -> raise_unsupported_expr expr)
+          (loop f) xs
+    | Pexp_constant (Pconst_integer (n, None)) -> add (Int (int_of_string n))
+    | Pexp_fun (Nolabel, None, arg_pat, body) -> (
+        match arg_pat.ppat_desc with
+        | Ppat_var { txt = name; _ } -> add (Lam (name, loop body))
+        | _ -> raise_unsupported_expr expr)
+    | Pexp_ident { txt = Lident "+"; _ } -> add (Prim `Add)
+    | Pexp_ident { txt = Lident name; _ } -> add (Var name)
+    | _ -> raise_unsupported_expr expr
+  in
+  loop
+
+let expr_of_parsetree (expr : Parsetree.expression) : expr ast =
+  let ast = new_ast () in
+  ast.root <- add_expr_to_ast ast expr;
+  ast
+
+module Wrap = struct
   let expression (pexp_desc : Parsetree.expression_desc) : Parsetree.expression
       =
     {
@@ -45,23 +70,29 @@ let rec parsetree_of_expr : expr -> Parsetree.expression =
       pexp_loc_stack = [];
       pexp_attributes = [];
     }
-  and ident (name : string) : Longident.t Asttypes.loc =
+
+  let ident (name : string) : Longident.t Asttypes.loc =
     { txt = Lident name; loc = Location.none }
-  and pattern (ppat_desc : Parsetree.pattern_desc) : Parsetree.pattern =
+
+  let pattern (ppat_desc : Parsetree.pattern_desc) : Parsetree.pattern =
     {
       ppat_desc;
       ppat_loc = Location.none;
       ppat_loc_stack = [];
       ppat_attributes = [];
     }
-  and var (name : string) : string Asttypes.loc =
+
+  let var (name : string) : string Asttypes.loc =
     { txt = name; loc = Location.none }
-  in
+end
+
+let parsetree_of_subexpr (ast : 'a ast) : expr -> Parsetree.expression =
+  let subexpr = get_subexpr ast in
 
   let list_of_apps expr =
     let rec loop = function
       | App (f, x) ->
-          let f, xs = loop f in
+          let f, xs = loop (subexpr f) in
           (f, x :: xs)
       | f -> (f, [])
     in
@@ -69,15 +100,24 @@ let rec parsetree_of_expr : expr -> Parsetree.expression =
     (f, List.rev xs)
   in
 
-  function
-  | App (_, _) as expr ->
-      let f, xs = list_of_apps expr in
-      let xs = List.map (fun x -> (Asttypes.Nolabel, parsetree_of_expr x)) xs in
-      expression (Pexp_apply (parsetree_of_expr f, xs))
-  | Int n -> expression (Pexp_constant (Pconst_integer (Int.to_string n, None)))
-  | Lam (x, b) ->
-      expression
-        (Pexp_fun
-           (Nolabel, None, pattern (Ppat_var (var x)), parsetree_of_expr b))
-  | Prim `Add -> expression (Pexp_ident (ident "+"))
-  | Var name -> expression (Pexp_ident (ident name))
+  let rec loop : expr -> Parsetree.expression = function
+    | App (_, _) as expr ->
+        let f, xs = list_of_apps expr in
+        let xs = List.map (fun x -> (Asttypes.Nolabel, loop (subexpr x))) xs in
+        Wrap.expression (Pexp_apply (loop f, xs))
+    | Int n ->
+        Wrap.expression (Pexp_constant (Pconst_integer (Int.to_string n, None)))
+    | Lam (x, b) ->
+        Wrap.expression
+          (Pexp_fun
+             ( Nolabel,
+               None,
+               Wrap.pattern (Ppat_var (Wrap.var x)),
+               loop (subexpr b) ))
+    | Prim `Add -> Wrap.expression (Pexp_ident (Wrap.ident "+"))
+    | Var name -> Wrap.expression (Pexp_ident (Wrap.ident name))
+  in
+  loop
+
+let parsetree_of_expr (ast : expr ast) : Parsetree.expression =
+  parsetree_of_subexpr ast (get_subexpr ast ast.root)