1 files changed, 135 insertions, 0 deletions
diff --git a/map.sml b/map.sml
new file mode 100644
index 0000000..27a60f5
--- /dev/null
+++ b/map.sml
@@ -0,0 +1,135 @@
+(* map is a red-black tree implementing an immutable key-value store. *)
+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
+end
+
+functor Map (type k
+             val cmp : k * k -> order) :> MAP where type key = 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
+
+  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))
+        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
+        end
+
+  fun sub1 (SOME (Node (Black, x, a, b))) = Node (Red, x, a, b)
+    | sub1 _ = raise Fail "invariance violation"
+
+  (* Thanks to https://github.com/sweirich/dth/blob/master/examples/red-black/RedBlack.lhs *)
+  fun delete t k =
+        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)
+
+          and delLeft nodeKV (a as SOME (Node (Black, _, _, _))) b = balLeft nodeKV (del a) b
+            | delLeft nodeKV a b = Node (Red, nodeKV, del a, b)
+
+          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"
+
+          and delRight nodeKV a (b as SOME (Node (Black, _, _, _))) = balRight nodeKV a (del b)
+            | delRight nodeKV a b = Node (Red, nodeKV, a, del b)
+
+          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"
+
+          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
+end