Use GADTs for prims, actual redex finding.

5 files changed, 75 insertions, 27 deletions

diff --git a/discocaml/arraylist.ml b/discocaml/arraylist.ml
index d1a9ecd..110cd17 100644
--- a/discocaml/arraylist.ml
+++ b/discocaml/arraylist.ml
@@ -28,6 +28,7 @@ let%test _ = next_pow2 12345 = 16384
 let make len dummy =
   { array = Array.make (max (next_pow2 len) 4) dummy; dummy; len }
 
+let copy { array; dummy; len } = { array = Array.copy array; dummy; len }
 let length { len; _ } = len
 
 let get { array; len; _ } i =
diff --git a/discocaml/arraylist.mli b/discocaml/arraylist.mli
index e42d194..6b40ce0 100644
--- a/discocaml/arraylist.mli
+++ b/discocaml/arraylist.mli
@@ -1,6 +1,7 @@
 type 'a t
 
 val make : int -> 'a -> 'a t
+val copy : 'a t -> 'a t
 val length : 'a t -> int
 val get : 'a t -> int -> 'a
 val set : 'a t -> int -> 'a -> unit
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
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
index 6ea4ff2..e04a59c 100644
--- a/discocaml/draw_tree.ml
+++ b/discocaml/draw_tree.ml
@@ -9,9 +9,9 @@ let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
         i.index
   | Int n -> Format.fprintf fmt "  expr%d [label=\"%d\"];\n" i.index n
   | Lam (_, _) -> Format.fprintf fmt "  expr%d [label=\"λ\"];\n" i.index
-  | Prim (`Add, _) -> Format.fprintf fmt "  expr%d [label=\"+\"];\n" i.index
-  | Prim (`Sub, _) -> Format.fprintf fmt "  expr%d [label=\"-\"];\n" i.index
-  | Prim (`Mul, _) -> Format.fprintf fmt "  expr%d [label=\"*\"];\n" i.index
+  | Prim (Add, _) -> Format.fprintf fmt "  expr%d [label=\"+\"];\n" i.index
+  | Prim (Sub, _) -> Format.fprintf fmt "  expr%d [label=\"-\"];\n" i.index
+  | Prim (Mul, _) -> Format.fprintf fmt "  expr%d [label=\"*\"];\n" i.index
   | Var n -> Format.fprintf fmt "  expr%d [label=%S];\n" i.index n
 
 let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
@@ -31,7 +31,15 @@ let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
         Format.fprintf fmt "  expr%d -> expr%d_var;\n" i.index i.index;
         Format.fprintf fmt "  expr%d_var [label=%S];\n" i.index x;
         edge_to b
-    | Prim (_, xs) -> Array.iter edge_to xs
+    | Prim (Add, (l, r)) ->
+        edge_to l;
+        edge_to r
+    | Prim (Sub, (l, r)) ->
+        edge_to l;
+        edge_to r
+    | Prim (Mul, (l, r)) ->
+        edge_to l;
+        edge_to r
     | Var _ -> ()
   in
   loop
diff --git a/discocaml/eval.ml b/discocaml/eval.ml
index 816d269..9c12003 100644
--- a/discocaml/eval.ml
+++ b/discocaml/eval.ml
@@ -1,5 +1,43 @@
 open Ast
 
-let find_redex_cbv : expr ast -> expr index option = function _ -> None
-let find_redex_cbn : expr ast -> expr index option = function _ -> None
-let reduce : expr ast -> expr index -> expr ast = fun e _ -> e
+let or_else (type a) (l : a option) (r : a option Lazy.t) : a option =
+  match l with Some x -> Some x | None -> Lazy.force r
+
+let is_redex (ast : expr ast) (i : expr index) : bool =
+  match get_subexpr ast i with
+  | App (f, _) -> (
+      match get_subexpr ast f with Lam (_, _) -> true | _ -> false)
+  | Int _ -> false
+  | Lam (_, _) -> false
+  | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) -> (
+      match (get_subexpr ast l, get_subexpr ast r) with
+      | Int _, Int _ -> true
+      | _ -> false)
+  | Var _ -> false
+
+let find_redex_cbv_in (_ast : expr ast) (_i : expr index) : expr index option =
+  None
+
+let find_redex_cbn_in (ast : expr ast) : expr index -> expr index option =
+  let rec loop (i : expr index) : expr index option =
+    if is_redex ast i then Some i
+    else
+      match get_subexpr ast i with
+      | App (f, x) -> or_else (loop f) (lazy (loop x))
+      | Int _ -> None
+      | Lam (_, _) -> None
+      | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) ->
+          or_else (loop l) (lazy (loop r))
+      | Var _ -> None
+  in
+  loop
+
+let find_redex_cbv (ast : expr ast) : expr index option =
+  find_redex_cbv_in ast ast.root
+
+let find_redex_cbn (ast : expr ast) : expr index option =
+  find_redex_cbn_in ast ast.root
+
+let reduce (ast : expr ast) (_i : expr index) : expr ast =
+  let ast = copy ast in
+  ast