Give args to prims.

1 files changed, 22 insertions, 3 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index ae74014..d1652d8 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -19,7 +19,7 @@ type expr =
   | App of expr index * expr index
   | Int of int
   | Lam of string * expr index
-  | Prim of prim
+  | Prim of prim * expr index array
   | Var of string
 [@@deriving show { with_path = false }]
 
@@ -39,6 +39,10 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
   in
   let rec loop (expr : Parsetree.expression) : expr index =
     match expr.pexp_desc with
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident "+"; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (`Add, [| loop l; loop r |]))
     | Pexp_apply (f, xs) ->
         List.fold_left
           (fun f -> function
@@ -50,7 +54,15 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
         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 "+"; _ } ->
+        add
+          (Lam
+             ( "$0",
+               add
+                 (Lam
+                    ( "$1",
+                      add (Prim (`Add, [| add (Var "$0"); add (Var "$1") |])) ))
+             ))
     | Pexp_ident { txt = Lident name; _ } -> add (Var name)
     | _ -> raise_unsupported_expr expr
   in
@@ -114,7 +126,14 @@ let parsetree_of_subexpr (ast : 'a ast) : expr -> Parsetree.expression =
                None,
                Wrap.pattern (Ppat_var (Wrap.var x)),
                loop (subexpr b) ))
-    | Prim `Add -> Wrap.expression (Pexp_ident (Wrap.ident "+"))
+    | Prim (`Add, [| l; r |]) ->
+        Wrap.expression
+          (Pexp_apply
+             ( Wrap.expression (Pexp_ident (Wrap.ident "+")),
+               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
+    | Prim (p, xs) ->
+        failwith
+          ("illegal Prim: " ^ [%derive.show: prim * expr index array] (p, xs))
     | Var name -> Wrap.expression (Pexp_ident (Wrap.ident name))
   in
   loop