Merge pull request #59 from natefaubion/perf

Performance improvements

Author: garyb

src/Data/Map.purs

@@ -116,14 +116,24 @@ checkValid tree = length (nub (allHeights tree)) == one
 -- | Lookup a value for the specified key
 lookup :: forall k v. (Ord k) => k -> Map k v -> Maybe v
 lookup _ Leaf = Nothing
-lookup k (Two _ k1 v _) | k == k1 = Just v
-lookup k (Two left k1 _ _) | k < k1 = lookup k left
-lookup k (Two _ _ _ right) = lookup k right
-lookup k (Three _ k1 v1 _ _ _ _) | k == k1 = Just v1
-lookup k (Three _ _ _ _ k2 v2 _) | k == k2 = Just v2
-lookup k (Three left k1 _ _ _ _ _) | k < k1 = lookup k left
-lookup k (Three _ k1 _ mid k2 _ _) | k1 < k && k <= k2 = lookup k mid
-lookup k (Three _ _ _ _ _ _ right) = lookup k right
+lookup k tree =
+  let comp :: k -> k -> Ordering
+      comp = compare
+  in case tree of
+    Two left k1 v right ->
+      case comp k k1 of
+        EQ -> Just v
+        LT -> lookup k left
+        _  -> lookup k right
+    Three left k1 v1 mid k2 v2 right ->
+      case comp k k1 of
+        EQ -> Just v1
+        c1 ->
+          case c1, comp k k2 of
+            _ , EQ -> Just v2
+            LT, _  -> lookup k left
+            _ , GT -> lookup k right
+            _ , _  -> lookup k mid
 
 -- | Test if a key is a member of a map
 member :: forall k v. (Ord k) => k -> Map k v -> Boolean
@@ -138,82 +148,104 @@ data TreeContext k v
 
 fromZipper :: forall k v. (Ord k) => List (TreeContext k v) -> Map k v -> Map k v
 fromZipper Nil tree = tree
-fromZipper (Cons (TwoLeft k1 v1 right) ctx) left = fromZipper ctx (Two left k1 v1 right)
-fromZipper (Cons (TwoRight left k1 v1) ctx) right = fromZipper ctx (Two left k1 v1 right)
-fromZipper (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) left = fromZipper ctx (Three left k1 v1 mid k2 v2 right)
-fromZipper (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) mid = fromZipper ctx (Three left k1 v1 mid k2 v2 right)
-fromZipper (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) right = fromZipper ctx (Three left k1 v1 mid k2 v2 right)
+fromZipper (Cons x ctx) tree =
+  case x of
+    TwoLeft k1 v1 right -> fromZipper ctx (Two tree k1 v1 right)
+    TwoRight left k1 v1 -> fromZipper ctx (Two left k1 v1 tree)
+    ThreeLeft k1 v1 mid k2 v2 right -> fromZipper ctx (Three tree k1 v1 mid k2 v2 right)
+    ThreeMiddle left k1 v1 k2 v2 right -> fromZipper ctx (Three left k1 v1 tree k2 v2 right)
+    ThreeRight left k1 v1 mid k2 v2 -> fromZipper ctx (Three left k1 v1 mid k2 v2 tree)
 
 data KickUp k v = KickUp (Map k v) k v (Map k v)
 
 -- | Insert a key/value pair into a map
 insert :: forall k v. (Ord k) => k -> v -> Map k v -> Map k v
 insert = down Nil
   where
+  comp :: k -> k -> Ordering
+  comp = compare
+
   down :: List (TreeContext k v) -> k -> v -> Map k v -> Map k v
   down ctx k v Leaf = up ctx (KickUp Leaf k v Leaf)
-  down ctx k v (Two left k1 _ right) | k == k1 = fromZipper ctx (Two left k v right)
-  down ctx k v (Two left k1 v1 right) | k < k1 = down (Cons (TwoLeft k1 v1 right) ctx) k v left
-  down ctx k v (Two left k1 v1 right) = down (Cons (TwoRight left k1 v1) ctx) k v right
-  down ctx k v (Three left k1 _ mid k2 v2 right) | k == k1 = fromZipper ctx (Three left k v mid k2 v2 right)
-  down ctx k v (Three left k1 v1 mid k2 _ right) | k == k2 = fromZipper ctx (Three left k1 v1 mid k v right)
-  down ctx k v (Three left k1 v1 mid k2 v2 right) | k < k1 = down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k v left
-  down ctx k v (Three left k1 v1 mid k2 v2 right) | k1 < k && k <= k2 = down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k v mid
-  down ctx k v (Three left k1 v1 mid k2 v2 right) = down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k v right
+  down ctx k v (Two left k1 v1 right) =
+    case comp k k1 of
+      EQ -> fromZipper ctx (Two left k v right)
+      LT -> down (Cons (TwoLeft k1 v1 right) ctx) k v left
+      _  -> down (Cons (TwoRight left k1 v1) ctx) k v right
+  down ctx k v (Three left k1 v1 mid k2 v2 right) =
+    case comp k k1 of
+      EQ -> fromZipper ctx (Three left k v mid k2 v2 right)
+      c1 ->
+        case c1, comp k k2 of
+          _ , EQ -> fromZipper ctx (Three left k1 v1 mid k v right)
+          LT, _  -> down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k v left
+          GT, LT -> down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k v mid
+          _ , _  -> down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k v right
 
   up :: List (TreeContext k v) -> KickUp k v -> Map k v
   up Nil (KickUp left k v right) = Two left k v right
-  up (Cons (TwoLeft k1 v1 right) ctx) (KickUp left k v mid) = fromZipper ctx (Three left k v mid k1 v1 right)
-  up (Cons (TwoRight left k1 v1) ctx) (KickUp mid k v right) = fromZipper ctx (Three left k1 v1 mid k v right)
-  up (Cons (ThreeLeft k1 v1 c k2 v2 d) ctx) (KickUp a k v b) = up ctx (KickUp (Two a k v b) k1 v1 (Two c k2 v2 d))
-  up (Cons (ThreeMiddle a k1 v1 k2 v2 d) ctx) (KickUp b k v c) = up ctx (KickUp (Two a k1 v1 b) k v (Two c k2 v2 d))
-  up (Cons (ThreeRight a k1 v1 b k2 v2) ctx) (KickUp c k v d) = up ctx (KickUp (Two a k1 v1 b) k2 v2 (Two c k v d))
+  up (Cons x ctx) kup =
+    case x, kup of
+      TwoLeft k1 v1 right, KickUp left k v mid -> fromZipper ctx (Three left k v mid k1 v1 right)
+      TwoRight left k1 v1, KickUp mid k v right -> fromZipper ctx (Three left k1 v1 mid k v right)
+      ThreeLeft k1 v1 c k2 v2 d, KickUp a k v b -> up ctx (KickUp (Two a k v b) k1 v1 (Two c k2 v2 d))
+      ThreeMiddle a k1 v1 k2 v2 d, KickUp b k v c -> up ctx (KickUp (Two a k1 v1 b) k v (Two c k2 v2 d))
+      ThreeRight a k1 v1 b k2 v2, KickUp c k v d -> up ctx (KickUp (Two a k1 v1 b) k2 v2 (Two c k v d))
 
 -- | Delete a key and its corresponding value from a map
 delete :: forall k v. (Ord k) => k -> Map k v -> Map k v
 delete = down Nil
   where
+  comp :: k -> k -> Ordering
+  comp = compare
+
   down :: List (TreeContext k v) -> k -> Map k v -> Map k v
   down ctx _ Leaf = fromZipper ctx Leaf
-  down ctx k (Two Leaf k1 _ Leaf)
-    | k == k1 = up ctx Leaf
-  down ctx k (Two left k1 v1 right)
-    | k == k1   = let max = maxNode left
-                    in removeMaxNode (Cons (TwoLeft max.key max.value right) ctx) left
-    | k < k1    = down (Cons (TwoLeft k1 v1 right) ctx) k left
-    | otherwise = down (Cons (TwoRight left k1 v1) ctx) k right
-  down ctx k (Three Leaf k1 v1 Leaf k2 v2 Leaf)
-    | k == k1 = fromZipper ctx (Two Leaf k2 v2 Leaf)
-    | k == k2 = fromZipper ctx (Two Leaf k1 v1 Leaf)
-  down ctx k (Three left k1 v1 mid k2 v2 right)
-    | k == k1 = let max = maxNode left
-                  in removeMaxNode (Cons (ThreeLeft max.key max.value mid k2 v2 right) ctx) left
-    | k == k2 = let max = maxNode mid
-                  in removeMaxNode (Cons (ThreeMiddle left k1 v1 max.key max.value right) ctx) mid
-    | k < k1               = down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k left
-    | k1 < k && k < k2 = down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k mid
-    | otherwise            = down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k right
+  down ctx k (Two left k1 v1 right) =
+    case right, comp k k1 of
+      Leaf, EQ -> up ctx Leaf
+      _   , EQ -> let max = maxNode left
+                   in removeMaxNode (Cons (TwoLeft max.key max.value right) ctx) left
+      _   , LT -> down (Cons (TwoLeft k1 v1 right) ctx) k left
+      _   , _  -> down (Cons (TwoRight left k1 v1) ctx) k right
+  down ctx k (Three left k1 v1 mid k2 v2 right) =
+    let leaves =
+          case left, mid, right of
+            Leaf, Leaf, Leaf -> true
+            _   , _   , _    -> false
+    in case leaves, comp k k1, comp k k2 of
+      true, EQ, _  -> fromZipper ctx (Two Leaf k2 v2 Leaf)
+      true, _ , EQ -> fromZipper ctx (Two Leaf k1 v1 Leaf)
+      _   , EQ, _  -> let max = maxNode left
+                       in removeMaxNode (Cons (ThreeLeft max.key max.value mid k2 v2 right) ctx) left
+      _   , _ , EQ -> let max = maxNode mid
+                       in removeMaxNode (Cons (ThreeMiddle left k1 v1 max.key max.value right) ctx) mid
+      _   , LT, _  -> down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k left
+      _   , GT, LT -> down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k mid
+      _   , _ , _  -> down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k right
 
   up :: List (TreeContext k v) -> Map k v -> Map k v
   up Nil tree = tree
-  up (Cons (TwoLeft k1 v1 Leaf) ctx) Leaf = fromZipper ctx (Two Leaf k1 v1 Leaf)
-  up (Cons (TwoRight Leaf k1 v1) ctx) Leaf = fromZipper ctx (Two Leaf k1 v1 Leaf)
-  up (Cons (TwoLeft k1 v1 (Two m k2 v2 r)) ctx) l = up ctx (Three l k1 v1 m k2 v2 r)
-  up (Cons (TwoRight (Two l k1 v1 m) k2 v2) ctx) r = up ctx (Three l k1 v1 m k2 v2 r)
-  up (Cons (TwoLeft k1 v1 (Three b k2 v2 c k3 v3 d)) ctx) a = fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
-  up (Cons (TwoRight (Three a k1 v1 b k2 v2 c) k3 v3) ctx) d = fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
-  up (Cons (ThreeLeft k1 v1 Leaf k2 v2 Leaf) ctx) Leaf = fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
-  up (Cons (ThreeMiddle Leaf k1 v1 k2 v2 Leaf) ctx) Leaf = fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
-  up (Cons (ThreeRight Leaf k1 v1 Leaf k2 v2) ctx) Leaf = fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
-  up (Cons (ThreeLeft k1 v1 (Two b k2 v2 c) k3 v3 d) ctx) a = fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
-  up (Cons (ThreeMiddle (Two a k1 v1 b) k2 v2 k3 v3 d) ctx) c = fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
-  up (Cons (ThreeMiddle a k1 v1 k2 v2 (Two c k3 v3 d)) ctx) b = fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
-  up (Cons (ThreeRight a k1 v1 (Two b k2 v2 c) k3 v3) ctx) d = fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
-  up (Cons (ThreeLeft k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4 e) ctx) a = fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
-  up (Cons (ThreeMiddle (Three a k1 v1 b k2 v2 c) k3 v3 k4 v4 e) ctx) d = fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
-  up (Cons (ThreeMiddle a k1 v1 k2 v2 (Three c k3 v3 d k4 v4 e)) ctx) b = fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))
-  up (Cons (ThreeRight a k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4) ctx) e = fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))
-  up _ _ = unsafeThrow "Impossible case in 'up'"
+  up (Cons x ctx) tree =
+    case x, tree of
+      TwoLeft k1 v1 Leaf, Leaf -> fromZipper ctx (Two Leaf k1 v1 Leaf)
+      TwoRight Leaf k1 v1, Leaf -> fromZipper ctx (Two Leaf k1 v1 Leaf)
+      TwoLeft k1 v1 (Two m k2 v2 r), l -> up ctx (Three l k1 v1 m k2 v2 r)
+      TwoRight (Two l k1 v1 m) k2 v2, r -> up ctx (Three l k1 v1 m k2 v2 r)
+      TwoLeft k1 v1 (Three b k2 v2 c k3 v3 d), a -> fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
+      TwoRight (Three a k1 v1 b k2 v2 c) k3 v3, d -> fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
+      ThreeLeft k1 v1 Leaf k2 v2 Leaf, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
+      ThreeMiddle Leaf k1 v1 k2 v2 Leaf, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
+      ThreeRight Leaf k1 v1 Leaf k2 v2, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
+      ThreeLeft k1 v1 (Two b k2 v2 c) k3 v3 d, a -> fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
+      ThreeMiddle (Two a k1 v1 b) k2 v2 k3 v3 d, c -> fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
+      ThreeMiddle a k1 v1 k2 v2 (Two c k3 v3 d), b -> fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
+      ThreeRight a k1 v1 (Two b k2 v2 c) k3 v3, d -> fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
+      ThreeLeft k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4 e, a -> fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
+      ThreeMiddle (Three a k1 v1 b k2 v2 c) k3 v3 k4 v4 e, d -> fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
+      ThreeMiddle a k1 v1 k2 v2 (Three c k3 v3 d k4 v4 e), b -> fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))
+      ThreeRight a k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4, e -> fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))
+      _, _ -> unsafeThrow "Impossible case in 'up'"
 
   maxNode :: Map k v -> { key :: k, value :: v }
   maxNode (Two _ k v Leaf) = { key: k, value: v }