Use GADTs for prims, actual redex finding.

1 files changed, 20 insertions, 20 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index a706a19..63893fa 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -13,18 +13,22 @@ 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 | `Mul | `Sub ] [@@deriving show { with_path = false }]
+type 'a prim =
+  | Add : (expr index * expr index) prim
+  | Mul : (expr index * expr index) prim
+  | Sub : (expr index * expr index) prim
 
-type expr =
+and expr =
   | App of expr index * expr index
   | Int of int
   | Lam of string * expr index
-  | Prim of prim * expr index array
+  | Prim : 'a prim * 'a -> expr
   | Var of string
-[@@deriving show { with_path = false }]
 
 type 'a ast = { subexprs : expr Arraylist.t; mutable root : 'a index }
-[@@deriving show { with_path = false }]
+
+let copy (ast : 'a ast) : 'a ast =
+  { subexprs = Arraylist.copy ast.subexprs; root = ast.root }
 
 let get_subexpr (ast : 'a ast) : expr index -> expr = get ast.subexprs
 
@@ -37,12 +41,11 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
     Arraylist.push ast.subexprs expr;
     { index }
   in
-  let binop (prim : prim) : expr index =
+  let binop (prim : (expr index * expr index) prim) : expr index =
     add
       (Lam
          ( "$0",
-           add
-             (Lam ("$1", add (Prim (prim, [| add (Var "$0"); add (Var "$1") |]))))
+           add (Lam ("$1", add (Prim (prim, (add (Var "$0"), add (Var "$1"))))))
          ))
   in
 
@@ -51,15 +54,15 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
     | Pexp_apply
         ( { pexp_desc = Pexp_ident { txt = Lident "+"; _ }; _ },
           [ (Nolabel, l); (Nolabel, r) ] ) ->
-        add (Prim (`Add, [| loop l; loop r |]))
+        add (Prim (Add, (loop l, loop r)))
     | Pexp_apply
         ( { pexp_desc = Pexp_ident { txt = Lident "-"; _ }; _ },
           [ (Nolabel, l); (Nolabel, r) ] ) ->
-        add (Prim (`Sub, [| loop l; loop r |]))
+        add (Prim (Sub, (loop l, loop r)))
     | Pexp_apply
         ( { pexp_desc = Pexp_ident { txt = Lident "*"; _ }; _ },
           [ (Nolabel, l); (Nolabel, r) ] ) ->
-        add (Prim (`Mul, [| loop l; loop r |]))
+        add (Prim (Mul, (loop l, loop r)))
     | Pexp_apply (f, xs) ->
         List.fold_left
           (fun f -> function
@@ -71,9 +74,9 @@ 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 "+"; _ } -> binop `Add
-    | Pexp_ident { txt = Lident "-"; _ } -> binop `Sub
-    | Pexp_ident { txt = Lident "*"; _ } -> binop `Mul
+    | Pexp_ident { txt = Lident "+"; _ } -> binop Add
+    | Pexp_ident { txt = Lident "-"; _ } -> binop Sub
+    | Pexp_ident { txt = Lident "*"; _ } -> binop Mul
     | Pexp_ident { txt = Lident name; _ } -> add (Var name)
     | _ -> raise_unsupported_expr expr
   in
@@ -137,24 +140,21 @@ let parsetree_of_subexpr (ast : 'a ast) : expr -> Parsetree.expression =
                None,
                Wrap.pattern (Ppat_var (Wrap.var x)),
                loop (subexpr b) ))
-    | Prim (`Add, [| l; r |]) ->
+    | Prim (Add, (l, r)) ->
         Wrap.expression
           (Pexp_apply
              ( Wrap.expression (Pexp_ident (Wrap.ident "+")),
                [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
-    | Prim (`Sub, [| l; r |]) ->
+    | Prim (Sub, (l, r)) ->
         Wrap.expression
           (Pexp_apply
              ( Wrap.expression (Pexp_ident (Wrap.ident "-")),
                [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
-    | Prim (`Mul, [| l; r |]) ->
+    | Prim (Mul, (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