| 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 ^ "}" |