Speed up adjust and adjustWithKey
authorDavid Feuer <David.Feuer@gmail.com>
Mon, 2 May 2016 17:07:19 +0000 (13:07 -0400)
committerDavid Feuer <David.Feuer@gmail.com>
Mon, 2 May 2016 17:07:19 +0000 (13:07 -0400)
Previously, `adjustWithKey` was implemented using `updateWithKey`.
`updateWithKey` needs to rebalance as it builds the result tree.
`adjustWithKey` never changes the shape of the tree, so
rebalancing on the way up is a waste of time.

Data/Map/Base.hs
Data/Map/Strict.hs

index 6401c0c..789b4a7 100644 (file)
@@ -810,7 +810,15 @@ adjust f = adjustWithKey (\_ x -> f x)
 -- > adjustWithKey f 7 empty                         == empty
 
 adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
-adjustWithKey f = updateWithKey (\k' x' -> Just (f k' x'))
+adjustWithKey = go
+  where
+    go :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
+    go _ !_ Tip = Tip
+    go f k (Bin sx kx x l r) =
+        case compare k kx of
+           LT -> Bin sx kx x (go f k l) r
+           GT -> Bin sx kx x l (go f k r)
+           EQ -> Bin sx kx (f kx x) l r
 #if __GLASGOW_HASKELL__
 {-# INLINABLE adjustWithKey #-}
 #else
index c061cab..7b82e2e 100644 (file)
@@ -488,7 +488,16 @@ adjust f = adjustWithKey (\_ x -> f x)
 -- > adjustWithKey f 7 empty                         == empty
 
 adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
-adjustWithKey f = updateWithKey (\k' x' -> Just (f k' x'))
+adjustWithKey = go
+  where
+    go :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
+    go _ !_ Tip = Tip
+    go f k (Bin sx kx x l r) =
+        case compare k kx of
+           LT -> Bin sx kx x (go f k l) r
+           GT -> Bin sx kx x l (go f k r)
+           EQ -> Bin sx kx x' l r
+             where !x' = f kx x
 #if __GLASGOW_HASKELL__
 {-# INLINABLE adjustWithKey #-}
 #else