From c2638e5d17b93e79851da8c43a377b7d51e723bc Mon Sep 17 00:00:00 2001
From: Nathan Ringo <nathan@remexre.com>
Date: Tue, 23 Jan 2024 23:22:37 -0600
Subject: Draw edges for the binding structure.

---
 discocaml/ast.ml       | 41 +++++++++++++++++++++++++++++++++++++++++
 discocaml/draw_tree.ml | 11 +++++++++++
 2 files changed, 52 insertions(+)

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
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
index 16719df..f0badd7 100644
--- a/discocaml/draw_tree.ml
+++ b/discocaml/draw_tree.ml
@@ -88,13 +88,24 @@ let draw_tree (ast : expr ast) : string =
   let fmt = Format.formatter_of_buffer buf in
   Format.fprintf fmt "digraph {\n";
   Format.fprintf fmt "  node [shape=\"box\", style=\"rounded\"];\n";
+
   add_expr_edges ast fmt nodes ast.root;
   Format.fprintf fmt "\n";
+
   IntSet.iter
     (fun index ->
       let i = { index } in
       add_node fmt i (get_subexpr ast i))
     !nodes;
+
+  get_binders ast
+  |> Array.iteri (fun i -> function
+       | Some j ->
+           Format.fprintf fmt
+             "  expr%d -> expr%d_var [color=\"#cccccc\", constraint=false];\n" i
+             j.index
+       | None -> ());
+
   Format.fprintf fmt "}\n";
   Format.pp_print_flush fmt ();
   Buffer.contents buf
-- 
cgit v1.2.3