Flattens the AST.

6 files changed, 188 insertions, 50 deletions

diff --git a/discocaml/arraylist.ml b/discocaml/arraylist.ml
new file mode 100644
index 0000000..d1a9ecd
--- /dev/null
+++ b/discocaml/arraylist.ml
@@ -0,0 +1,90 @@
+type 'a t = { mutable array : 'a array; dummy : 'a; mutable len : int }
+
+(* Utilities *)
+
+let next_pow2 x =
+  (* From Bit Twiddling Hacks *)
+  let x = x - 1 in
+  let x = x lor (x lsr 1) in
+  let x = x lor (x lsr 2) in
+  let x = x lor (x lsr 4) in
+  let x = x lor (x lsr 8) in
+  let x = x lor (x lsr 16) in
+  let x = x lor (x lsr 32) in
+  x + 1
+
+let%test _ = next_pow2 0 = 0
+let%test _ = next_pow2 1 = 1
+let%test _ = next_pow2 2 = 2
+let%test _ = next_pow2 3 = 4
+let%test _ = next_pow2 4 = 4
+let%test _ = next_pow2 5 = 8
+let%test _ = next_pow2 15 = 16
+let%test _ = next_pow2 50 = 64
+let%test _ = next_pow2 12345 = 16384
+
+(* Arraylist operations *)
+
+let make len dummy =
+  { array = Array.make (max (next_pow2 len) 4) dummy; dummy; len }
+
+let length { len; _ } = len
+
+let get { array; len; _ } i =
+  if i >= len then
+    raise
+      (Invalid_argument
+         (Format.sprintf
+            "index %d out of bounds for an Arraylist.t of length %d" i len))
+  else array.(i)
+
+let set { array; len; _ } i x =
+  if i >= len then
+    raise
+      (Invalid_argument
+         (Format.sprintf
+            "index %d out of bounds for an Arraylist.t of length %d" i len))
+  else array.(i) <- x
+
+let push arraylist x =
+  if Array.length arraylist.array = arraylist.len then
+    arraylist.array <-
+      Array.init (arraylist.len * 2) (fun i ->
+          if i < arraylist.len then arraylist.array.(i) else arraylist.dummy);
+  arraylist.array.(arraylist.len) <- x;
+  arraylist.len <- arraylist.len + 1
+
+let to_array { array; len; _ } = Array.sub array 0 len
+
+module Array_for_pp = struct
+  type 'a t = 'a array [@@deriving show { with_path = false }]
+end
+
+let pp (fmt_elem : Format.formatter -> 'a -> unit) (fmt : Format.formatter)
+    (arraylist : 'a t) : unit =
+  Array_for_pp.pp fmt_elem fmt (to_array arraylist)
+
+let%expect_test _ =
+  let xs = make 3 0 in
+  push xs 3;
+  set xs 0 1;
+  set xs 1 1;
+  set xs 2 2;
+  push xs 5;
+  print_endline ([%derive.show: int t] xs);
+  [%expect "[|1; 1; 2; 3; 5|]"]
+
+let%expect_test _ =
+  let xs = make 0 false in
+  let add x =
+    let index = length xs in
+    push xs x;
+    index
+  in
+
+  let i = add false and j = add true and k = add true in
+
+  print_endline
+    ([%derive.show: int list * bool list]
+       ([ i; j; k ], [ get xs i; get xs j; get xs k ]));
+  [%expect "([0; 1; 2], [false; true; true])"]
diff --git a/discocaml/arraylist.mli b/discocaml/arraylist.mli
new file mode 100644
index 0000000..e42d194
--- /dev/null
+++ b/discocaml/arraylist.mli
@@ -0,0 +1,9 @@
+type 'a t
+
+val make : int -> 'a -> 'a t
+val length : 'a t -> int
+val get : 'a t -> int -> 'a
+val set : 'a t -> int -> 'a -> unit
+val push : 'a t -> 'a -> unit
+val to_array : 'a t -> 'a array
+val pp : (Format.formatter -> 'a -> unit) -> Format.formatter -> 'a t -> unit
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)
diff --git a/discocaml/discocaml.ml b/discocaml/discocaml.ml
index 2140fa5..f11a81b 100644
--- a/discocaml/discocaml.ml
+++ b/discocaml/discocaml.ml
@@ -1,3 +1,5 @@
+open Ast
+
 type command = [ `Parse | `DrawTree | `RunCBN | `RunCBV | `StepCBN | `StepCBV ]
 
 let command_of_yojson = function
@@ -31,21 +33,21 @@ let handle_request { expr; command } : response =
   try
     let buf = Lexing.from_string expr in
     let expr = Parse.expression buf in
-    let expr = Ast.expr_of_parsetree expr in
+    let expr = expr_of_parsetree expr in
 
-    let expr_response (expr : Ast.expr) : response =
+    let expr_response (expr : expr ast) : response =
       let buf = Buffer.create 16 in
       (let fmt = Format.formatter_of_buffer buf in
-       Pprintast.expression fmt (Ast.parsetree_of_expr expr);
+       Pprintast.expression fmt (parsetree_of_expr expr);
        Format.pp_print_flush fmt ());
       let has_redex = Option.is_some (Eval.find_redex_cbn expr) in
       `Expr { expr = Buffer.contents buf; has_redex }
     in
 
-    let step_with (find_redex : Ast.expr -> Eval.path option) (expr : Ast.expr)
-        : Ast.expr =
+    let step_with (find_redex : expr ast -> expr index option) (expr : expr ast)
+        : expr ast =
       match find_redex expr with
-      | Some path -> Eval.reduce expr path
+      | Some i -> Eval.reduce expr i
       | None -> failwith "no redex"
     in
 
diff --git a/discocaml/dune b/discocaml/dune
index 50b12c2..83bf5ed 100644
--- a/discocaml/dune
+++ b/discocaml/dune
@@ -2,7 +2,7 @@
  (name discocaml)
  (inline_tests)
  (libraries compiler-libs.common yojson)
- (modules ast discocaml eval)
+ (modules arraylist ast discocaml eval)
  (package discocaml)
  (preprocess
   (pps ppx_deriving.show ppx_deriving_yojson ppx_expect)))
diff --git a/discocaml/eval.ml b/discocaml/eval.ml
index 63dd6f7..816d269 100644
--- a/discocaml/eval.ml
+++ b/discocaml/eval.ml
@@ -1,8 +1,5 @@
 open Ast
 
-type path = int list
-(** A path from the root of a term to a subterm. *)
-
-let find_redex_cbv : expr -> path option = function _ -> None
-let find_redex_cbn : expr -> path option = function _ -> None
-let reduce : expr -> path -> expr = fun e _ -> e
+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