3 files changed, 52 insertions, 3 deletions
diff --git a/discocaml/discocaml.ml b/discocaml/discocaml.ml
index f11a81b..df5cef3 100644
--- a/discocaml/discocaml.ml
+++ b/discocaml/discocaml.ml
@@ -17,7 +17,8 @@ type request = { expr : string; command : command }
 type response_expr = { expr : string; has_redex : bool }
 [@@deriving to_yojson { exn = true }]
 
-type response = [ `Error of string | `Expr of response_expr ]
+type response =
+  [ `Error of string | `Expr of response_expr | `Graphviz of string ]
 [@@deriving to_yojson { exn = true }]
 
 let%expect_test _ =
@@ -55,7 +56,8 @@ let handle_request { expr; command } : response =
     | `Parse -> expr_response expr
     | `StepCBN -> expr_response (step_with Eval.find_redex_cbn expr)
     | `StepCBV -> expr_response (step_with Eval.find_redex_cbv expr)
-    | `DrawTree | `RunCBN | `RunCBV -> failwith "not implemented"
+    | `DrawTree -> `Graphviz (Draw_tree.draw_tree expr)
+    | `RunCBN | `RunCBV -> failwith "not implemented"
   with
   | Failure msg -> `Error msg
   | exn -> `Error ("uncaught exception: " ^ Printexc.to_string exn)
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
diff --git a/discocaml/dune b/discocaml/dune
index 83bf5ed..bd96d64 100644
--- a/discocaml/dune
+++ b/discocaml/dune
@@ -2,7 +2,7 @@
  (name discocaml)
  (inline_tests)
  (libraries compiler-libs.common yojson)
- (modules arraylist ast discocaml eval)
+ (modules arraylist ast discocaml draw_tree eval)
  (package discocaml)
  (preprocess
   (pps ppx_deriving.show ppx_deriving_yojson ppx_expect)))