Re-implement map.

I stole Haskell's map, and I read a paper.

1 files changed, 168 insertions, 109 deletions

diff --git a/map.sml b/map.sml
index d74cfd5..a5b316c 100644
--- a/map.sml
+++ b/map.sml
@@ -1,13 +1,23 @@
-(* map is a red-black tree implementing an immutable key-value store. *)
+(* map is a size-balanced binary tree implementing an immutable key-value store.
+
+ * Stolen with love from Haskell's Data.Map. *)
 signature MAP =
 sig
   type key
   type 'a map
 
   val empty : 'a map
-  val insert : 'a map -> key * 'a -> 'a map
-  val lookup : 'a map -> key -> 'a option
-  val delete : 'a map -> key -> 'a map
+  val null : 'a map -> bool
+  val size : 'a map -> int
+  val insert : key -> 'a -> 'a map -> 'a map
+  val lookup : key -> 'a map -> 'a option
+  val delete : key -> 'a map -> 'a map
+  val union : 'a map -> 'a map -> 'a map
+  val intersect : 'a map -> 'a map -> 'a map
+  val difference : 'a map -> 'b map -> 'a map
+  val fromList : (key * 'a) list -> 'a map
+  val toList : 'a map -> (key * 'a) list
+  val lookupMin : 'a map -> (key * 'a) option
 end
 
 functor Map (type k
@@ -15,121 +25,170 @@ functor Map (type k
 struct
   type key = k
 
-  datatype color = Red | Black
-  datatype 'a node = Node of color * (key * 'a) * 'a node option * 'a node option
-  type 'a map = 'a node option
-  
-  val empty = NONE
+  datatype 'a map = Tip
+                  | Bin of int * key * 'a * 'a map * 'a map
 
-  fun rebalance (Node (Black, g, SOME (Node (Red, p, SOME (Node (Red, n, nLeft, nRight)), pRight)), u)) : 'a node =
-        (*     b
-              /
-             r
-            /
-           r     *)
-        Node (Red, p, SOME (Node (Black, n, nLeft, nRight)),
-                      SOME (Node (Black, g, pRight, u)))
-    | rebalance (Node (Black, g, u, SOME (Node (Red, p, pLeft, SOME (Node (Red, n, nLeft, nRight)))))) =
-        (* b
-            \
-             r
-              \
-               r *)
-        Node (Red, p, SOME (Node (Black, g, u, pLeft)),
-                      SOME (Node (Black, n, nLeft, nRight)))
-    | rebalance (Node (Black, g, u, SOME (Node (Red, p, SOME (Node (Red, n, nLeft, nRight)), pRight)))) =
-        (* b
-            \
-             r
-            /
-           r   *)
-        Node (Red, n, SOME (Node (Black, g, u, nLeft)),
-                      SOME (Node (Black, p, nRight, pRight)))
-    | rebalance (Node (Black, g, SOME (Node (Red, p, pLeft, SOME (Node (Red, n, nLeft, nRight)))), u)) =
-        (*   b
-            /
-           r
-            \
-             r *)
-       Node (Red, n, SOME (Node (Black, p, pLeft, nLeft)),
-                     SOME (Node (Black, g, nRight, u)))
-    | rebalance n = n
-  
-  fun insert m (kv as (k, _)) =
-        let
-          fun insertNode NONE = Node (Red, kv, NONE, NONE)
-            | insertNode (SOME (Node (color, nodeKV as (nodeKey, _), left, right))) =
-                case cmp (k, nodeKey) of
-                  LESS => rebalance (Node (color, nodeKV, SOME (insertNode left), right))
-                | EQUAL => Node (color, kv, left, right)
-                | GREATER => rebalance (Node (color, nodeKV, left, SOME (insertNode right)))
-          val Node (_, nodeKV, left, right) = insertNode m
-        in
-          SOME (Node (Black, nodeKV, left, right))
+  val empty : 'a map = Tip
+
+  fun null Tip = true
+    | null _ = false
+
+  fun size Tip = 0
+    | size (Bin (s, _, _, _, _)) = s
+
+  fun lookup (k : key) (m : 'a map) : 'a option =
+    let fun go Tip = NONE
+          | go (Bin (_, kx, x, l, r)) =
+              case cmp (k, kx) of
+                LESS => go l
+              | GREATER => go r
+              | EQUAL => SOME x
+    in go m
+    end
+
+  fun singleton (k : key) (v : 'a) : 'a map = Bin (1, k, v, Tip, Tip)
+
+  fun size Tip = 0
+    | size (Bin (sz, _, _, _, _)) = sz
+
+  (* The bin constructor maintains the size of the tree. *)
+  fun bin k x l r = Bin (size l + size r + 1, k, x, l, r)
+
+  (* https://hackage.haskell.org/package/containers-0.4.0.0/docs/src/Data-Map.html#delta *)
+  val delta = 4
+  val ratio = 2
+
+  fun singleL k1 x1 t1 (Bin (_, k2, x2, t2, t3)) = bin k2 x2 (bin k1 x1 t1 t2) t3
+    | singleL _ _ _ Tip = raise Fail "singleL Tip"
+
+  fun singleR k1 x1 (Bin (_, k2, x2, t1, t2)) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3)
+    | singleR _ _ Tip _ = raise Fail "singleR Tip"
+
+  fun doubleL k1 x1 t1 (Bin (_, k2, x2, Bin (_, k3, x3, t2, t3), t4)) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4)
+    | doubleL _ _ _ _ = raise Fail "doubleL"
+
+  fun doubleR k1 x1 (Bin (_, k2, x2, t1, Bin (_, k3, x3, t2, t3))) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4)
+    | doubleR _ _ _ _ = raise Fail "doubleR"
+
+  fun rotateL k x l (r as Bin (_, _, _, ly, ry)) =
+        if size ly < ratio*size ry then singleL k x l r
+        else doubleL k x l r
+    | rotateL _ _ _ Tip = raise Fail "rotateL Tip"
+
+  fun rotateR k x (l as Bin (_, _, _, ly, ry)) r =
+        if size ry < ratio*size ly then singleR k x l r
+        else doubleR k x l r
+    | rotateR _ _ Tip _ = raise Fail "rotateR Tip"
+
+  fun balance k x l r =
+    let
+      val sizeL = size l
+      val sizeR = size r
+    in
+      if sizeL + sizeR <= 1 then bin k x l r
+      else if sizeR >= delta*sizeL then rotateL k x l r
+      else if sizeL >= delta*sizeR then rotateR k x l r
+      else bin k x l r
+    end
+
+  fun insertMax kx x t =
+    case t of
+      Tip => singleton kx x
+    | Bin (_, ky, y, l, r) =>
+        balance ky y l (insertMax kx x r)
+
+  fun insertMin kx x t =
+    case t of
+      Tip => singleton kx x
+    | Bin (_, ky, y, l, r) =>
+        balance ky y (insertMin kx x l) r
+
+  fun join kx x Tip r = insertMin kx x r
+    | join kx x l Tip = insertMax kx x l
+    | join kx x (l as Bin (sizeL, ky, y, ly, ry)) (r as Bin (sizeR, kz, z, lz, rz)) =
+        if delta*sizeL <= sizeR then balance kz z (join kx x l lz) rz
+        else if delta*sizeR <= sizeL then balance ky y ly (join kx x ry r)
+        else bin kx x l r
+
+  fun split (t : 'a map) (k : key) =
+    case t of
+      Tip => (Tip, false, Tip)
+    | Bin (_, km, m, l, r) =>
+        case cmp (k, km) of
+          EQUAL => (l, true, r)
+        | LESS =>
+            let val (ll, b, lr) = split l k
+            in (ll, b, join km m lr r)
+            end
+        | GREATER =>
+            let val (rl, b, rr) = split r k
+            in (join km m l rl, b, rr)
+            end
+
+  fun splitLast (Bin (_, kx, x, l, Tip)) = (l, (kx, x))
+    | splitLast (Bin (_, kx, x, l, r)) =
+        let val (t', k') = splitLast r
+        in (join kx x l t', k')
+        end
+    | splitLast Tip = raise Fail "splitLast Tip"
+
+  fun join2 Tip tr = tr
+    | join2 tl tr =
+        let val (tl', (kx, x)) = splitLast tl
+        in join kx x tl' tr
         end
 
-  fun lookup m k =
-        let fun loop NONE = NONE
-              | loop (SOME (Node (_, (nodeKey, nodeVal), left, right))) =
-                  case cmp (k, nodeKey) of
-                    LESS => loop left
-                  | EQUAL => SOME nodeVal
-                  | GREATER => loop right
-        in loop m
+  fun delete k m =
+    let val (tl, _, tr) = split m k
+    in join2 tl tr
+    end
+
+  fun union (t1 : 'a map) (t2 : 'a map) : 'a map =
+    case (t1, t2) of
+      (Tip, _) => t2
+    | (_, Tip) => t1
+    | (_, Bin (_, k2, v2, l2, r2)) =>
+        let
+          val (l1, _, r1) = split t1 k2
+          val tl = union l1 l2
+          val tr = union r1 r2
+        in join k2 v2 tl tr
         end
 
-  fun sub1 (SOME (Node (Black, x, a, b))) = Node (Red, x, a, b)
-    | sub1 _ = raise Fail "invariance violation"
+  fun intersect Tip _ = Tip
+    | intersect _ Tip = Tip
+    | intersect t1 (Bin (_, k2, v2, l2, r2)) =
+        let
+          val (l1, b, r1) = split t1 k2
+          val tl = intersect l1 l2
+          val tr = intersect r1 r2
+        in
+          if b
+            then join k2 v2 tl tr
+            else join2 tl tr
+        end
 
-  (* Thanks to https://github.com/sweirich/dth/blob/master/examples/red-black/RedBlack.lhs *)
-  fun delete t k =
+  fun difference Tip _ = Tip
+    | difference t1 Tip = t1
+    | difference t1 (Bin (_, k2, v2, l2, r2)) =
         let
-          fun del NONE = NONE
-            | del (SOME (Node (_, (nodeKV as (nodeKey, _)), left, right))) =
-                case cmp (k, nodeKey) of
-                  LESS => SOME (delLeft nodeKV left right)
-                | EQUAL => merge left right
-                | GREATER => SOME (delRight nodeKV left right)
+          val (l1, _, r1) = split t1 k2
+          val tl = difference l1 l2
+          val tr = difference r1 r2
+        in join2 tl tr
+        end
 
-          and delLeft nodeKV (a as SOME (Node (Black, _, _, _))) b = balLeft nodeKV (del a) b
-            | delLeft nodeKV a b = Node (Red, nodeKV, del a, b)
+  fun insert (k : key) (v : 'a) (m : 'a map) : 'a map = union m (singleton k v)
 
-          and balLeft x (SOME (Node (Red, y, a, b))) c = Node (Red, x, SOME (Node (Black, y, a, b)), c)
-            | balLeft x bl (SOME (Node (Black, y, a, b))) =
-                rebalance (Node (Black, x, bl, SOME (Node (Red, y, a, b))))
-            | balLeft x bl (SOME (Node (Red, z, SOME (Node (Black, y, a, b)), c))) =
-                Node (Red, y, SOME (Node (Black, x, bl, a)), SOME (rebalance (Node (Black, z, b, SOME (sub1 c)))))
-            | balLeft _ _ _ = raise Fail "unreachable"
+  fun fromList l = foldl (fn ((kx, x), acc) => insert kx x acc) empty l
 
-          and delRight nodeKV a (b as SOME (Node (Black, _, _, _))) = balRight nodeKV a (del b)
-            | delRight nodeKV a b = Node (Red, nodeKV, a, del b)
+  fun toList' Tip acc = acc
+    | toList' (Bin (_, k, v, l, r)) acc =
+        toList' l ((k, v) :: toList' r acc)
 
-          and balRight x a (SOME (Node (Red, y, b, c))) = Node (Red, x, a, SOME (Node (Black, y, b, c)))
-            | balRight y (SOME (Node (Black, x, a, b))) bl = rebalance (Node (Black, y, SOME (Node (Red, x, a, b)), bl))
-            | balRight z (SOME (Node (Red, x, a, SOME (Node (Black, y, b, c))))) bl =
-                Node (Red, y, SOME (rebalance (Node (Black, x, SOME (sub1 a), b))),
-                              SOME (Node (Black, z, c, bl)))
-            | balRight _ _ _ = raise Fail "unreachable"
+  fun toList m = toList' m []
 
-          and merge NONE x = x
-            | merge x NONE = x
-            | merge (SOME (Node (Red, x, a, b))) (SOME (Node (Red, y, c, d))) =
-                (case merge b c of
-                  SOME (Node (Red, z, b', c')) =>
-                    SOME (Node (Red, z, SOME (Node (Red, x, a, b')),
-                                        SOME (Node (Red, y, c', d))))
-                | bc => SOME (Node (Red, x, a, SOME (Node (Red, y, bc, d)))))
-            | merge (SOME (Node (Black, x, a, b))) (SOME (Node (Black, y, c, d))) =
-                (case merge b c of
-                  SOME (Node (Red, z, b', c')) =>
-                    SOME (Node (Red, z, SOME (Node (Black, x, a, b')),
-                                        SOME (Node (Black, y, c', d))))
-                | bc => SOME (balLeft x a (SOME (Node (Black, y, bc, d)))))
-            | merge a (SOME (Node (Red, x, b, c))) = SOME (Node (Red, x, merge a b, c))
-            | merge (SOME (Node (Red, x, a, b))) c = SOME (Node (Red, x, a, merge b c))
-        in
-          case del t of
-            NONE => NONE
-          | SOME (Node (_, kv, left, right)) => SOME (Node (Black, kv, left, right))
-        end
+  fun lookupMin Tip = NONE
+    | lookupMin (Bin (_, k, v, Tip, _)) = SOME (k, v)
+    | lookupMin (Bin (_, _, _, l, _)) = lookupMin l
 end