A bunch of the AST, but as a tree; we're gonna flatten it!

1 files changed, 72 insertions, 4 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index 22f3790..d01be66 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -1,8 +1,42 @@
-type expr = Foo [@@deriving show { with_path = false }]
+exception UnsupportedAst of string
 
-let expr_of_parsetree : Parsetree.expression -> expr = function _ -> Foo
+let raise_unsupported_ast (type a) (f : Format.formatter -> a -> unit) (x : a) =
+  let buf = Buffer.create 16 in
+  let fmt = Format.formatter_of_buffer buf in
+  f fmt x;
+  Format.pp_print_flush fmt ();
+  raise (UnsupportedAst (Buffer.contents buf))
 
-let parsetree_of_expr : expr -> Parsetree.expression =
+let raise_unsupported_expr = raise_unsupported_ast (Printast.expression 0)
+
+type prim = [ `Add ] [@@deriving show { with_path = false }]
+
+type expr =
+  | App of expr * expr
+  | Int of int
+  | Lam of string * expr
+  | 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 =
   let expression (pexp_desc : Parsetree.expression_desc) : Parsetree.expression
       =
     {
@@ -11,5 +45,39 @@ let parsetree_of_expr : expr -> Parsetree.expression =
       pexp_loc_stack = [];
       pexp_attributes = [];
     }
+  and ident (name : string) : Longident.t Asttypes.loc =
+    { txt = Lident name; loc = Location.none }
+  and 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 =
+    { txt = name; loc = Location.none }
+  in
+
+  let list_of_apps expr =
+    let rec loop = function
+      | App (f, x) ->
+          let f, xs = loop f in
+          (f, x :: xs)
+      | f -> (f, [])
+    in
+    let f, xs = loop expr in
+    (f, List.rev xs)
   in
-  function Foo -> expression (Pexp_constant (Pconst_char '!'))
+
+  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))