Draw edges for the binding structure.

1 files changed, 41 insertions, 0 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index a4e2272..92f368f 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -266,3 +266,44 @@ let show_expr (expr : expr ast) : string =
   Pprintast.expression fmt (parsetree_of_expr expr);
   Format.pp_print_flush fmt ();
   Buffer.contents buf
+
+type binders = expr index option array
+
+let get_binders (ast : expr ast) : binders =
+  let out = Array.make (Arraylist.length ast.subexprs) None in
+  let rec loop (env : (string * expr index) list) (i : expr index) : unit =
+    match get_subexpr ast i with
+    | App (f, x) ->
+        loop env f;
+        loop env x
+    | Cons (hd, tl) ->
+        loop env hd;
+        loop env tl
+    | If (cond, then_, else_) ->
+        loop env cond;
+        loop env then_;
+        loop env else_
+    | Lam (x, b) -> loop ((x, i) :: env) b
+    | Let (false, name, bound, body) ->
+        loop env bound;
+        loop ((name, i) :: env) body
+    | Let (true, name, bound, body) ->
+        loop ((name, i) :: env) bound;
+        loop ((name, i) :: env) body
+    | Prim (Add, (l, r)) ->
+        loop env l;
+        loop env r
+    | Prim (Sub, (l, r)) ->
+        loop env l;
+        loop env r
+    | Prim (Mul, (l, r)) ->
+        loop env l;
+        loop env r
+    | Prim (RelOp, (_, l, r)) ->
+        loop env l;
+        loop env r
+    | Var name -> out.(i.index) <- List.assoc_opt name env
+    | Bool _ | Int _ | Nil -> ()
+  in
+  loop [] ast.root;
+  out