exercises/ch01/tree.ml

   1 module Array = ArrayLabels
   2 module List = ListLabels
   3
   4 module type TREE = sig
   5   type ('k, 'v) t
   6
   7   val empty : ('k, 'v) t
   8
   9   val set : ('k, 'v) t -> k:'k -> v:'v -> ('k, 'v) t
  10
  11   val get : ('k, 'v) t -> k:'k -> 'v option
  12
  13   val member : ('k, 'v) t -> k:'k -> bool
  14
  15   val to_dot : ('k, 'v) t -> k_to_string:('k -> string) -> string
  16 end
  17
  18 module BinaryTree : TREE = struct
  19   type ('k, 'v) t =
  20     | Node of 'k * 'v * ('k, 'v) t * ('k, 'v) t
  21     | Leaf
  22
  23   let empty = Leaf
  24
  25   let rec set t ~k ~v =
  26     match t with
  27     | Leaf -> Node (k, v, Leaf, Leaf)
  28     | Node (k', v', l, r) when k < k' -> Node (k', v', set l ~k ~v, r)
  29     | Node (k', v', l, r) when k > k' -> Node (k', v', l, set r ~k ~v)
  30     | Node (k, _, l, r) -> Node (k, v, l, r)
  31
  32   let rec get t ~k =
  33     match t with
  34     | Leaf -> None
  35     | Node (k', _, l, _) when k < k' -> get l ~k
  36     | Node (k', _, _, r) when k > k' -> get r ~k
  37     | Node (_, v, _, _) -> Some v
  38
  39   let rec member t ~k =
  40     match t with
  41     | Leaf -> false
  42     | Node (k', _, l, _) when k < k' -> member l ~k
  43     | Node (k', _, _, r) when k > k' -> member r ~k
  44     | Node _ -> true
  45
  46   let to_edges t =
  47     let rec to_edges_from k1 t =
  48       match t with
  49       | Leaf -> []
  50       | Node (k2, _, l, r) ->
  51           (k1, k2) :: ((to_edges_from k2 l) @ (to_edges_from k2 r))
  52     in
  53     match t with
  54     | Leaf -> []
  55     | Node (k, _, l, r) -> (to_edges_from k l) @ (to_edges_from k r)
  56
  57   let to_dot t ~k_to_string =
  58     let (edges, _) =
  59       List.fold_left (to_edges t)
  60         ~init:("", "\n")
  61         ~f:(fun (edges, sep) (k1, k2) ->
  62           let k1, k2 = k_to_string k1, k_to_string k2 in
  63           (Printf.sprintf "%s%s    %S -> %S;\n" edges sep k1 k2, "")
  64       )
  65     in
  66     "digraph G {" ^ edges ^ "}"
  67 end
  68
  69 let () =
  70   let tree_a = BinaryTree.empty in
  71   let tree_a = BinaryTree.set tree_a ~k:"k1" ~v:"v1" in
  72   let tree_a = BinaryTree.set tree_a ~k:"k2" ~v:"v2" in
  73   assert (BinaryTree.member tree_a ~k:"k1");
  74   assert (BinaryTree.member tree_a ~k:"k2");
  75   assert (Some "v1" = BinaryTree.get tree_a ~k:"k1");
  76   assert (Some "v2" = BinaryTree.get tree_a ~k:"k2");
  77   let tree_b =
  78     Array.fold_left (Sys.argv)
  79       ~init:BinaryTree.empty
  80       ~f:(fun t k -> BinaryTree.set t ~k ~v:())
  81   in
  82   print_endline (BinaryTree.to_dot tree_b ~k_to_string:(fun x -> x))