Track height once instead of at every node

1 files changed, 88 insertions, 85 deletions

diff --git a/Map.sml b/Map.sml
index d8b3a37..a60c5cc 100644
--- a/Map.sml
+++ b/Map.sml
@@ -22,31 +22,30 @@ functor Map (type k
 struct
   type key = k
 
-  datatype 'a map =
+  datatype 'a node =
     Tip
-  | Two of int * 'a map * key * 'a * 'a map
-  | Three of int * 'a map * key * 'a * 'a map * key * 'a * 'a map
+  | Two of 'a node * key * 'a * 'a node
+  | Three of 'a node * key * 'a * 'a node * key * 'a * 'a node
 
-  val empty : 'a map = Tip
+  type 'a map = int * 'a node
 
-  fun null Tip = true
-    | null _ = false
+  val empty : 'a map = (0, Tip)
 
-  fun height (Tip : 'a map) : int = 0
-    | height (Two (h, _, _, _, _)) = h
-    | height (Three (h, _, _, _, _, _, _, _)) = h
+  fun height ((h, _) : 'a map) : int = h
 
-  fun two (l : 'a map) (k : key) (v : 'a) (r : 'a map) : 'a map =
-    if height l <> height r then raise Fail "two: height mismatch" else
-    Two (height l + 1, l, k, v, r)
+  fun null (m : 'a map) : bool = height m = 0
 
-  fun three (a : 'a map) (k1 : key) (v1 : 'a) (b : 'a map) (k2 : key) (v2 : 'a) (c : 'a map) : 'a map =
-    if not (height a = height b andalso height b = height c) then raise Fail "three: height mismatch" else
-    Three (height a + 1, a, k1, v1, b, k2, v2, c)
+  fun two ((lh, l) : 'a map) (k : key) (v : 'a) ((rh, r) : 'a map) : 'a map =
+    if lh <> rh then raise Fail "two: height mismatch" else
+    (lh + 1, Two (l, k, v, r))
 
-  fun view (Tip : 'a map) : ('a map * key * 'a * 'a map) option = NONE
-    | view (Two (_, l, k, v, r)) = SOME (l, k, v, r)
-    | view (Three (h, a, k1, v1, b, k2, v2, c)) = SOME (a, k1, v1, two b k2 v2 c)
+  fun three ((ah, a) : 'a map) (k1 : key) (v1 : 'a) ((bh, b) : 'a map) (k2 : key) (v2 : 'a) ((ch, c) : 'a map) : 'a map =
+    if not (ah = bh andalso bh = ch) then raise Fail "three: height mismatch" else
+    (ah + 1, Three (a, k1, v1, b, k2, v2, c))
+
+  fun view ((_, Tip) : 'a map) : ('a map * key * 'a * 'a map) option = NONE
+    | view (h, Two (l, k, v, r)) = SOME ((h - 1, l), k, v, (h - 1, r))
+    | view (h, Three (a, k1, v1, b, k2, v2, c)) = SOME ((h - 1, a), k1, v1, two (h - 1, b) k2 v2 (h - 1, c))
 
   fun lookup (k : key) (m : 'a map) : 'a option =
     case view m of
@@ -66,26 +65,26 @@ struct
       Split (left, k, v, right)
     else if height left < height right then
       case right of
-        Tip => raise Fail "unreachable"
-      | Two (_, rl, rk, rv, rr) =>
-          (case join' left k v rl of
-            One newNode => One (two newNode rk rv rr)
-          | Split (left, k, v, right) => One (three left k v right rk rv rr))
-      | Three (_, ra, rk1, rv1, rb, rk2, rv2, rc) =>
-          case join' left k v ra of
-            One newNode => One (three newNode rk1 rv1 rb rk2 rv2 rc)
-          | Split (left, k, v, right) => Split (two left k v right, rk1, rv1, two rb rk2 rv2 rc)
+        (_, Tip) => raise Fail "unreachable"
+      | (rh, Two (rl, rk, rv, rr)) =>
+          (case join' left k v (rh - 1, rl) of
+            One newNode => One (two newNode rk rv (rh - 1, rr))
+          | Split (left, k, v, right) => One (three left k v right rk rv (rh - 1, rr)))
+      | (rh, Three (ra, rk1, rv1, rb, rk2, rv2, rc)) =>
+          case join' left k v (rh - 1, ra) of
+            One newNode => One (three newNode rk1 rv1 (rh - 1, rb) rk2 rv2 (rh - 1, rc))
+          | Split (left, k, v, right) => Split (two left k v right, rk1, rv1, two (rh - 1, rb) rk2 rv2 (rh - 1, rc))
     else
       case left of
-        Tip => raise Fail "unreachable"
-      | Two (_, ll, lk, lv, lr) =>
-          (case join' lr k v right of
-            One newNode => One (two ll lk lv newNode)
-          | Split (left, k, v, right) => One (three ll lk lv left k v right))
-      | Three (_, la, lk1, lv1, lb, lk2, lv2, lc) =>
-          case join' lc k v right of
-            One newNode => One (three la lk1 lv1 lb lk2 lv2 newNode)
-          | Split (left, k, v, right) => Split (two la lk1 lv1 lb, lk2, lv2, two left k v right)
+        (_, Tip) => raise Fail "unreachable"
+      | (lh, Two (ll, lk, lv, lr)) =>
+          (case join' (lh - 1, lr) k v right of
+            One newNode => One (two (lh - 1, ll) lk lv newNode)
+          | Split (left, k, v, right) => One (three (lh - 1, ll) lk lv left k v right))
+      | (lh, Three (la, lk1, lv1, lb, lk2, lv2, lc)) =>
+          case join' (lh - 1, lc) k v right of
+            One newNode => One (three (lh - 1, la) lk1 lv1 (lh - 1, lb) lk2 lv2 newNode)
+          | Split (left, k, v, right) => Split (two (lh - 1, la) lk1 lv1 (lh - 1, lb), lk2, lv2, two left k v right)
 
   fun join (left : 'a map) (k : key) (v : 'a) (right : 'a map) : 'a map =
     case join' left k v right of
@@ -94,7 +93,7 @@ struct
 
   fun split (m : 'a map) (k : key) : 'a map * bool * 'a map =
     case view m of
-      NONE => (Tip, false, Tip)
+      NONE => (empty, false, empty)
     | SOME (l, k', v, r) =>
         case cmp (k, k') of
           EQUAL => (l, true, r)
@@ -109,63 +108,66 @@ struct
 
   fun splitLast (m : 'a map) : 'a map * (key * 'a) =
     case view m of
-      SOME (l, k, v, Tip) => (l, (k, v))
-    | SOME (l, k, v, r) =>
-        let val (t', k') = splitLast r
-        in (join l k v t', k')
-        end
+      SOME (l, k, v, r) =>
+        if null r then (l, (k, v))
+        else
+          let val (t', k') = splitLast r
+          in (join l k v t', k') end
     | NONE => raise Fail "splitLast Tip"
 
-  fun join2 (Tip : 'a map) (tr : 'a map) : 'a map = tr
-    | join2 tl tr =
-        let val (tl', (kx, x)) = splitLast tl
-        in join tl' kx x tr
-        end
+  fun join2 (tl : 'a map) (tr : 'a map) : 'a map =
+    if null tl then tr
+    else
+      let val (tl', (kx, x)) = splitLast tl
+      in join tl' kx x tr end
 
   fun delete (k : key) (m : 'a map) : 'a map =
     let val (tl, _, tr) = split m k
     in join2 tl tr
     end
 
-  fun union (Tip : 'a map) (t2 : 'a map) : 'a map = t2
-    | union t1 t2 =
-        case view t2 of
-          NONE => t1
-        | SOME (l2, k2, v2, r2) =>
-            let
-              val (l1, _, r1) = split t1 k2
-              val tl = union l1 l2
-              val tr = union r1 r2
-            in join tl k2 v2 tr
-            end
+  fun union (t1 : 'a map) (t2 : 'a map) : 'a map =
+    if null t1 then t2
+    else
+      case view t2 of
+        NONE => t1
+      | SOME (l2, k2, v2, r2) =>
+          let
+            val (l1, _, r1) = split t1 k2
+            val tl = union l1 l2
+            val tr = union r1 r2
+          in join tl k2 v2 tr
+          end
 
-  fun intersect (Tip : 'a map) (_ : 'a map) : 'a map = Tip
-    | intersect t1 t2 =
-        case view t2 of
-          NONE => Tip
-        | SOME (l2, k2, v2, r2) =>
-            let
-              val (l1, b, r1) = split t1 k2
-              val tl = intersect l1 l2
-              val tr = intersect r1 r2
-            in
-              if b then join tl k2 v2 tr
-              else join2 tl tr
-            end
+  fun intersect (t1 : 'a map) (t2 : 'a map) : 'a map =
+    if null t1 then empty
+    else
+      case view t2 of
+        NONE => empty
+      | SOME (l2, k2, v2, r2) =>
+          let
+            val (l1, b, r1) = split t1 k2
+            val tl = intersect l1 l2
+            val tr = intersect r1 r2
+          in
+            if b then join tl k2 v2 tr
+            else join2 tl tr
+          end
 
-  fun difference (Tip : 'a map) (_ : 'b map) : 'a map = Tip
-    | difference t1 t2 =
-        case view t2 of
-          NONE => t1
-        | SOME (l2, k2, _, r2) =>
-            let
-              val (l1, _, r1) = split t1 k2
-              val tl = difference l1 l2
-              val tr = difference r1 r2
-            in join2 tl tr
-            end
+  fun difference (t1 : 'a map) (t2 : 'b map) : 'a map =
+    if null t1 then empty
+    else
+      case view t2 of
+        NONE => t1
+      | SOME (l2, k2, _, r2) =>
+          let
+            val (l1, _, r1) = split t1 k2
+            val tl = difference l1 l2
+            val tr = difference r1 r2
+          in join2 tl tr
+          end
 
-  fun singleton (k : key) (v : 'a) : 'a map = two Tip k v Tip
+  fun singleton (k : key) (v : 'a) : 'a map = two empty k v empty
 
   fun insert (k : key) (v : 'a) (m : 'a map) : 'a map = union m (singleton k v)
 
@@ -181,6 +183,7 @@ struct
   fun lookupMin (m : 'a map) : (key * 'a) option =
     case view m of
       NONE => NONE
-    | SOME (Tip, k, v, _) => SOME (k, v)
-    | SOME (l, _, _, _) => lookupMin l
+    | SOME (l, k, v, _) =>
+        if null l then SOME (k, v)
+        else lookupMin l
 end