Add a representation for the splay tree

1 files changed, 84 insertions, 69 deletions

diff --git a/splay.ml b/splay.ml
index 6f740eb..ec5750c 100644
--- a/splay.ml
+++ b/splay.ml
@@ -1,21 +1,34 @@
-module Make (El : Tools.COMPARABLE_TYPE) = struct
+module type KEY = sig
+
+    type 'a t
+
+    val comp: 'a t -> 'b t -> ('a, 'b) Tools.cmp
+
+    val repr: Format.formatter -> 'a t -> unit
+
+end
+
+module Make (El : KEY) = struct
 
   type 'a elem = 'a El.t
 
-  type treeVal =
-    | Leaf : treeVal
-    | Node : treeVal * ('a elem * 'a) * treeVal -> treeVal
+  type leaf (** Fantom type for typing the tree *)
+  type node (** Fantom type for typing the tree *)
 
-  type tree = treeVal ref
+  type 'a treeVal =
+    | Leaf : leaf treeVal
+    | Node : _ treeVal * ('a elem * 'a) * _ treeVal -> node treeVal
 
-  type splay = S : treeVal * ('a elem * 'a) * treeVal -> splay
+  type t = T : 'a treeVal ref -> t [@@unboxed]
 
-  let empty = ref Leaf;;
+  let empty = T (ref Leaf)
 
-  let isEmpty tree = !tree = Leaf
+  let isEmpty (T tree) = match !tree with
+    | Leaf -> true
+    | _ -> false
 
-  let rec splay : type a. a elem -> splay -> splay = fun x t -> begin
-    let S (l, y, r) = t in
+  let rec splay : type a. a elem -> node treeVal -> node treeVal = fun x t -> begin
+    let Node (l, y, r) = t in
     begin match El.comp x (fst y)  with
       | Tools.Eq -> t
       | Tools.Lt ->
@@ -23,22 +36,20 @@ module Make (El : Tools.COMPARABLE_TYPE) = struct
           | Leaf -> t
           | Node (ll, z, rr) ->
             begin match El.comp x (fst z)  with
-              | Tools.Eq -> S (ll, z, Node (rr, y, r))
+              | Tools.Eq -> Node (ll, z, Node (rr, y, r))
               | Tools.Lt ->
                 begin match ll with
-                  | Leaf -> S (ll, z, Node (rr, y, r))
-                  | Node (t1, k, t2 ) ->
-                    let ll = S (t1, k, t2) in
-                    let S (newL, newV, newR) = splay x ll
-                    in S (newL, newV, Node (newR, z, Node (rr, y, r)))
+                  | Leaf -> Node (ll, z, Node (rr, y, r))
+                  | Node _ as ll ->
+                    let Node (newL, newV, newR) = splay x ll
+                    in Node (newL, newV, Node (newR, z, Node (rr, y, r)))
                 end
               | Tools.Gt ->
                 begin match rr with
-                  | Leaf -> S (ll, z, Node (rr, y, r))
-                  | Node (t1, k, t2 ) ->
-                    let rr = S (t1, k, t2) in
-                    let S (newL, newV, newR) = splay x rr
-                    in S (Node (ll, z, newL), newV, Node (newR, y, r))
+                  | Leaf -> Node (ll, z, Node (rr, y, r))
+                  | Node _ as rr ->
+                    let Node (newL, newV, newR) = splay x rr
+                    in Node (Node (ll, z, newL), newV, Node (newR, y, r))
                 end
             end
         end
@@ -47,81 +58,85 @@ module Make (El : Tools.COMPARABLE_TYPE) = struct
           | Leaf -> t
           | Node (ll, z, rr) ->
             begin match El.comp x (fst z)  with
-              | Tools.Eq -> S (Node (l, y, ll), z, rr)
+              | Tools.Eq -> Node (Node (l, y, ll), z, rr)
               | Tools.Lt ->
                 begin match ll with
-                  | Leaf -> S (Node (l, y, ll), z, rr)
-                  | Node (t1, k, t2 ) ->
-                    let ll = S (t1, k, t2) in
-                    let S (newL, newV, newR) = splay x ll
-                    in S (Node (l, y, newL), newV, Node (newR, z, rr))
+                  | Leaf -> Node (Node (l, y, ll), z, rr)
+                  | Node _ as ll ->
+                    let Node (newL, newV, newR) = splay x ll
+                    in Node (Node (l, y, newL), newV, Node (newR, z, rr))
                 end
               | Tools.Gt ->
                 begin match rr with
-                  | Leaf -> S (Node (l, y, ll), z, rr)
-                  | Node (t1, k, t2 ) ->
-                    let rr = S (t1, k, t2) in
-                    let S (newL, newV, newR) = splay x rr
-                    in S (Node (Node(l, y, ll), z, newL), newV, newR)
+                  | Leaf -> Node (Node (l, y, ll), z, rr)
+                  | Node _ as rr ->
+                    let Node (newL, newV, newR) = splay x rr
+                    in Node (Node (Node(l, y, ll), z, newL), newV, newR)
                 end
             end
         end
     end
   end
 
-  let member: type a. a elem -> treeVal ref -> bool = fun x t ->
-
-    match !t with
+  let member: type a. a elem -> t -> bool = fun x (T t) -> match !t with
     | Leaf -> false
-    | Node (l, c, r) ->
-        let S (l', c', r') = splay x (S (l, c, r)) in
-        t := Node (l', c', r');
+    | Node _ as root ->
+        let root' = splay x root in
+        t := root';
+        let Node (_, c', _) = root' in
         begin match El.comp (fst c') x  with
           | Tools.Eq -> true
           | _ -> false
         end
 
-  let find: type a. a elem -> treeVal ref -> a = fun x t ->
-
-    match !t with
+  let find: type a. a elem -> t -> a = fun x (T t) -> match !t with
     | Leaf -> raise Not_found
-    | Node (l, c, r) ->
-        let S (l', c', r') = splay x (S (l, c, r)) in
-        t := Node (l', c', r');
+    | Node _ as root ->
+        let root' = splay x root in
+        t := root';
+        let Node (_, c', _) = root' in
         begin match El.comp (fst c') x with
           | Tools.Eq -> snd c'
           | _ -> raise Not_found
         end
 
-  let add: type a. a elem -> a -> treeVal ref -> treeVal ref = fun key value t ->
-    begin match !t with
-    | Leaf -> ref (Node (Leaf, (key, value), Leaf))
-    | Node (l, c, r) ->
-        let S (l, y, r) = splay key (S (l, c, r)) in
+  let add: type a. a elem -> a -> t -> t = fun key value (T t) -> match !t with
+    | Leaf -> T (ref (Node (Leaf, (key, value), Leaf)))
+    | Node _ as root ->
+        let root' = splay key root in
+        let Node (l, y, r) = root' in
         begin match El.comp key (fst y)  with
-          | Tools.Eq -> ref (Node (l, y, r))
-          | Tools.Lt -> ref (Node (l, (key, value), Node (Leaf, y, r)))
-          | Tools.Gt -> ref (Node (Node (l, y, Leaf), (key, value), r))
+          | Tools.Eq -> T (ref root')
+          | Tools.Lt -> T (ref (Node (l, (key, value), Node (Leaf, y, r))))
+          | Tools.Gt -> T (ref (Node (Node (l, y, Leaf), (key, value), r)))
         end
-    end
 
-  let delete: type a. a elem -> treeVal ref -> treeVal ref = fun x t ->
+  let repr formatter (T t) = begin
+
+    let repr_edge from formatter dest = begin
+      Format.fprintf formatter "\"%a\" -> \"%a\"\n"
+        El.repr from
+        El.repr dest
+    end in
+
+    let rec repr': type a b. a El.t -> Format.formatter -> b treeVal -> unit = fun parent formatter -> function
+    | Leaf -> ()
+    | Node (l, c, r) ->
+        let key = fst c in
+        Format.fprintf formatter "%a%a%a"
+        (repr_edge parent) key
+        (repr' key) l
+        (repr' key) r in
+
     begin match !t with
-    | Leaf -> ref Leaf
+    | Leaf -> Format.fprintf formatter "digraph G {}"
     | Node (l, c, r) ->
-        let S (l, y, r) = splay x (S (l, c, r)) in
-        begin match El.comp x (fst y)  with
-            Tools.Eq ->
-            begin match (l,r) with
-              | (Leaf, _) -> ref r
-              | (_, Leaf) -> ref l
-              | (Node (t1, c, t2), _) ->
-                let S (newL, newV, newR) = splay x (S (t1, c, t2))
-                in ref (Node (newL, newV, r)) end
-          | _ -> ref (Node (l, y, r)) end
-        end
+      let key = fst c in
+      Format.fprintf formatter "digraph G {\n%a%a}"
+        (repr' key) l
+        (repr' key) r
+    end
+
+  end
 
-  let rec depth tree = match tree with
-    | Node (l, _, r) -> max (depth l) (depth r) + 1
-    | Leaf -> 0
 end