Leave DOT output on disk

[tiger.ml.git] / exercises / ch01 / tree.ml

diff --git a/exercises/ch01/tree.ml b/exercises/ch01/tree.ml
index 6b0ab43..195a70c 100644 (file)
--- a/exercises/ch01/tree.ml
+++ b/exercises/ch01/tree.ml
@@ -1,4 +1,5 @@
 module Array = ArrayLabels
+module List = ListLabels
 
 module type TREE = sig
   type ('k, 'v) t
@@ -10,6 +11,8 @@ module type TREE = sig
   val get : ('k, 'v) t -> k:'k -> 'v option
 
   val member : ('k, 'v) t -> k:'k -> bool
+
+  val to_dot : ('k, 'v) t -> k_to_string:('k -> string) -> string
 end
 
 module BinaryTree : TREE = struct
@@ -39,6 +42,28 @@ module BinaryTree : TREE = struct
     | Node (k', _, l, _) when k < k' -> member l ~k
     | Node (k', _, _, r) when k > k' -> member r ~k
     | Node _ -> true
+
+  let to_edges t =
+    let rec to_edges_from k1 t =
+      match t with
+      | Leaf -> []
+      | Node (k2, _, l, r) ->
+          (k1, k2) :: ((to_edges_from k2 l) @ (to_edges_from k2 r))
+    in
+    match t with
+    | Leaf -> []
+    | Node (k, _, l, r) -> (to_edges_from k l) @ (to_edges_from k r)
+
+  let to_dot t ~k_to_string =
+    let (edges, _) =
+      List.fold_left (to_edges t)
+        ~init:("", "\n")
+        ~f:(fun (edges, sep) (k1, k2) ->
+          let k1, k2 = k_to_string k1, k_to_string k2 in
+          (Printf.sprintf "%s%s    %S -> %S;\n" edges sep k1 k2, "")
+      )
+    in
+    "digraph G {" ^ edges ^ "}"
 end
 
 let () =
@@ -50,9 +75,8 @@ let () =
   assert (Some "v1" = BinaryTree.get tree_a ~k:"k1");
   assert (Some "v2" = BinaryTree.get tree_a ~k:"k2");
   let tree_b =
-    Array.fold_left
-      (Sys.argv)
+    Array.fold_left (Sys.argv)
       ~init:BinaryTree.empty
       ~f:(fun t k -> BinaryTree.set t ~k ~v:())
   in
-  Printf.printf "%B\n" (BinaryTree.member tree_b ~k:"a")
+  print_endline (BinaryTree.to_dot tree_b ~k_to_string:(fun x -> x))