1 files changed, 47 insertions, 0 deletions
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
new file mode 100644
index 0000000..30c466f
--- /dev/null
+++ b/discocaml/draw_tree.ml
@@ -0,0 +1,47 @@
+open Ast
+module IntSet = Set.Make (Int)
+
+let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
+  let label =
+    match expr with
+    | App _ -> "$"
+    | Int n -> string_of_int n
+    | Lam (x, _) -> "λ " ^ x ^ " →"
+    | Prim (`Add, _) -> "+"
+    | Var n -> n
+  in
+  Format.fprintf fmt "  expr%d [label=%S];\n" i.index label
+
+let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
+    (nodes : IntSet.t ref) : expr index -> unit =
+  let rec loop (i : expr index) : unit =
+    nodes := IntSet.add i.index !nodes;
+    let edge_to (j : expr index) : unit =
+      loop j;
+      Format.fprintf fmt "  expr%d -> expr%d;\n" i.index j.index
+    in
+    match get_subexpr ast i with
+    | App (f, x) ->
+        edge_to f;
+        edge_to x
+    | Int _ -> ()
+    | Lam (_, b) -> edge_to b
+    | Prim (_, xs) -> Array.iter edge_to xs
+    | Var _ -> ()
+  in
+  loop
+
+let draw_tree (ast : expr ast) : string =
+  let buf = Buffer.create 16 and nodes = ref IntSet.empty in
+  let fmt = Format.formatter_of_buffer buf in
+  Format.fprintf fmt "digraph {\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;
+  Format.fprintf fmt "}\n";
+  Format.pp_print_flush fmt ();
+  Buffer.contents buf