[commit: packages/containers] changelog-foldtree, cleaned_bugfix394, master, merge-doc-target, merge-fixes-5.9, merge-restrict-fix-5.8, revert-408-bugfix_394: Speed up adjust and adjustWithKey (2988826)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/containers

On branches: changelog-foldtree,cleaned_bugfix394,master,merge-doc-target,merge-fixes-5.9,merge-restrict-fix-5.8,revert-408-bugfix_394
Link       : http://git.haskell.org/packages/containers.git/commitdiff/2988826ffd37bb69fb02061ce7981c9482a8ccbb

>---------------------------------------------------------------

commit 2988826ffd37bb69fb02061ce7981c9482a8ccbb
Author: David Feuer <David.Feuer at gmail.com>
Date:   Mon May 2 13:07:19 2016 -0400

    Speed up adjust and adjustWithKey
    
    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.


>---------------------------------------------------------------

2988826ffd37bb69fb02061ce7981c9482a8ccbb
 Data/Map/Base.hs   | 10 +++++++++-
 Data/Map/Strict.hs | 11 ++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/Data/Map/Base.hs b/Data/Map/Base.hs
index 6401c0c..789b4a7 100644
--- a/Data/Map/Base.hs
+++ b/Data/Map/Base.hs
@@ -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
diff --git a/Data/Map/Strict.hs b/Data/Map/Strict.hs
index c061cab..7b82e2e 100644
--- a/Data/Map/Strict.hs
+++ b/Data/Map/Strict.hs
@@ -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