compiler/src/lib/tiger/tiger_map_red_black.ml

   1 module List = ListLabels
   2
   3 type color = R | B
   4 type ('k, 'v) t =
   5   | Leaf
   6   | Node of color * ('k * 'v) * ('k, 'v) t * ('k, 'v) t
   7
   8 let empty = Leaf
   9
  10 let set t ~k ~v =
  11   let balance = function
  12              | Leaf -> assert false
  13     (* LL *) | Node (B, x, Node (R, lx, Node (R, llx, lll, llr), lr                     ), r                                                              ) -> Node (R, lx , Node (B, llx, lll, llr), Node (B, x  , lr , r  ))
  14     (* LR *) | Node (B, x, Node (R, lx, ll                     , Node (R, lrx, lrl, lrr)), r                                                              ) -> Node (R, lrx, Node (B, lx , ll , lrl), Node (B, x  , lrr, r  ))
  15     (* RL *) | Node (B, x, l                                                             , Node (R, rx, Node (R, rlx, rll, rlr), rr                      )) -> Node (R, rlx, Node (B, x  , l  , rll), Node (B, rx , rlr, rr ))
  16     (* RR *) | Node (B, x, l                                                             , Node (R, rx, rl                      , Node (R, rrx, rrl, rrr))) -> Node (R, rx , Node (B, x  , l  , rl ), Node (B, rrx, rrl, rrr))
  17              | node -> node  (* Fragile pattern. Shall we reconsider? *)
  18   in
  19   let rec set t k v =
  20     match t with
  21     (* Because we recur, we cannot at this stage know if the following Node
  22      * with Leaf children will end-up as root (in which case it should be
  23      * black) or as the last node before leaves (in which case it should be
  24      * red).  We begin by assuming that it is the later and at the very end,
  25      * before returning to the caller, force the actual root node (which is
  26      * easy to identify at that point) to be black, regardless of what it
  27      * already is.
  28      *)
  29     | Leaf -> Node (R, (k, v), Leaf, Leaf)  (* Maybe root or last node, so R *)
  30     | Node (c, ((k', _) as x), l, r) when k < k' -> balance (Node (c, x, set l k v, r))
  31     | Node (c, ((k', _) as x), l, r) when k > k' -> balance (Node (c, x, l        , set r k v))
  32     | Node (c, _             , l, r) -> Node (c, (k, v), l, r)
  33   in
  34   match set t k v with
  35   | Leaf -> assert false  (* Can we GADT this away? *)
  36   | Node (_, (k, v), l, r) -> Node (B, (k, v), l, r)  (* Root is always Black *)
  37
  38 let rec get t ~k =
  39   match t with
  40   | Leaf -> None
  41   | Node (_, (k', _), l, _) when k < k' -> get l ~k
  42   | Node (_, (k', _), _, r) when k > k' -> get r ~k
  43   | Node (_, (_ , v), _, _) -> Some v
  44
  45 let rec member t ~k =
  46   match t with
  47   | Leaf -> false
  48   | Node (_, (k', _), l, _) when k < k' -> member l ~k
  49   | Node (_, (k', _), _, r) when k > k' -> member r ~k
  50   | Node (_, (_ , _), _, _) -> true
  51
  52 let to_edges t =
  53   let rec to_edges_from node1 t =
  54     match t with
  55     | Leaf -> [(node1, `Leaf)]
  56     | Node (c2, (k2, _), l, r) ->
  57         let node2 = k2, c2 in
  58         (node1, `Node node2) :: ((to_edges_from node2 l) @ (to_edges_from node2 r))
  59   in
  60   match t with
  61   | Leaf ->
  62       []
  63   | Node (c, (k, _), l, r) ->
  64       let node1 = k, c in
  65       (to_edges_from node1 l) @ (to_edges_from node1 r)
  66
  67 let color_to_string = function
  68   | B -> "black"
  69   | R -> "red"
  70
  71 let to_dot t ~k_to_string =
  72   let (dot_edges_and_nodes, _, _) =
  73     List.fold_left
  74       (to_edges t)
  75       ~init:("", "\n", 0)
  76       ~f:(fun (dot_edges_and_nodes, sep, leaves) ((k1, c1), node2_or_leaf) ->
  77         let k1 = k_to_string k1 in
  78         let k2, c2, leaves =
  79           match node2_or_leaf with
  80           | `Leaf ->
  81               let leaves = succ leaves in
  82               let label = Printf.sprintf "Leaf_%d" leaves in
  83               (label, B, leaves)
  84           | `Node (k2, c2) ->
  85               let label = k_to_string k2 in
  86               (label, c2, leaves)
  87         in
  88         let dot_edges_and_nodes =
  89           Printf.sprintf
  90             "%s%s    %S -> %S;\n    %S [color=%s,fontcolor=white,style=filled];\n    %S [color=%s,fontcolor=white,style=filled];\n"
  91             dot_edges_and_nodes
  92             sep
  93             k1 k2
  94             k1 (color_to_string c1)  (* Yes, it's redundant... *)
  95             k2 (color_to_string c2)
  96         in
  97         let sep = "" in
  98         (dot_edges_and_nodes, sep, leaves)
  99     )
 100   in
 101   "digraph G {" ^ dot_edges_and_nodes ^ "}"
 102
 103 let of_list pairs =
 104   List.fold_left pairs ~init:empty ~f:(fun t (k, v) -> set t ~k ~v)
 105
 106 let rec to_list = function
 107   | Leaf ->
 108       []
 109   | Node (_, pair, left, right) ->
 110       pair :: ((to_list left) @ (to_list right))