From 001a532037d6c0eceadc535fa04b80e0ab159f3c Mon Sep 17 00:00:00 2001
From: Brian Shu <littlebubu.shu@gmail.com>
Date: Thu, 26 May 2022 12:37:39 -0400
Subject: [PATCH 1/6] alter now runs in one pass

---
 Data/HashMap/Internal.hs   | 101 +++++++++++++++++++++++++++++++------
 unordered-containers.cabal |   3 +-
 2 files changed, 87 insertions(+), 17 deletions(-)

diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
index 8f6997a7..76a41c3c 100644
--- a/Data/HashMap/Internal.hs
+++ b/Data/HashMap/Internal.hs
@@ -828,7 +828,7 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
       where i = index h s
     go h k x s t@(Collision hy v)
         | h == hy   = Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v)
-        | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
+        | otherwise = runST $ two s h k x hy t
 {-# INLINABLE insert' #-}
 
 -- Insert optimized for the case when we know the key is not in the map.
@@ -1255,11 +1255,10 @@ adjust# f k0 m0 = go h0 k0 0 m0
 -- | \(O(\log n)\)  The expression @('update' f k map)@ updates the value @x@ at @k@
 -- (if it is in the map). If @(f x)@ is 'Nothing', the element is deleted.
 -- If it is @('Just' y)@, the key @k@ is bound to the new value @y@.
-update :: (Eq k, Hashable k) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a
-update f = alter (>>= f)
+update :: (Eq k, Hashable k, Show k, Show a) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a
+update f = Exts.inline alter (>>= f)
 {-# INLINABLE update #-}
 
-
 -- | \(O(\log n)\)  The expression @('alter' f k map)@ alters the value @x@ at @k@, or
 -- absence thereof.
 --
@@ -1268,19 +1267,89 @@ update f = alter (>>= f)
 -- @
 -- 'lookup' k ('alter' f k m) = f ('lookup' k m)
 -- @
-alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
--- TODO(m-renaud): Consider using specialized insert and delete for alter.
-alter f k m =
-  case f (lookup k m) of
-    Nothing -> delete k m
-    Just v  -> insert k v m
-{-# INLINABLE alter #-}
+alter :: (Eq k, Hashable k, Show k, Show v) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
+alter f k = alter' f (hash k) k
+{-# INLINEABLE alter #-}
+
+alter' :: (Eq k, Show v) => (Maybe v -> Maybe v) -> Hash -> k -> HashMap k v -> HashMap k v
+alter' f h0 k0 m0 = go h0 k0 0 m0
+  where
+    go !h !k !_ Empty = case f Nothing of
+      Nothing -> Empty
+      Just v -> Leaf h $ L k v
+    go h k s t@(Leaf hy l@(L ky v))
+      | hy == h =
+          if ky == k
+            then case f $ Just v of
+              Nothing -> Empty
+              Just v'
+                | v `ptrEq` v' -> t
+                | otherwise -> Leaf h $ L k v'
+            else do
+              case f Nothing of
+                Nothing -> t
+                Just v' -> collision h l $ L k v'
+      | otherwise = case f Nothing of
+          Nothing -> t
+          Just v' -> runST $ two s h k v' hy t
+    go h k s t@(BitmapIndexed b ary)
+      | b .&. m == 0 = case f Nothing of
+          Nothing -> t
+          Just v' -> bitmapIndexedOrFull (b .|. m) $! A.insert ary i $! Leaf h $! L k v'
+      | otherwise = do
+          let !st = A.index ary i
+              !st' = go h k (nextShift s) st
+          if st' `ptrEq` st
+            then t
+            else case st' of
+              Empty
+                | A.length ary == 1 -> Empty
+                | A.length ary == 2 ->
+                    case (i, A.index ary 0, A.index ary 1) of
+                      (0, _, l) | isLeafOrCollision l -> l
+                      (1, l, _) | isLeafOrCollision l -> l
+                      _ -> bIndexed
+                | otherwise -> bIndexed
+                where
+                  bIndexed = BitmapIndexed (b .&. complement m) (A.delete ary i)
+              l | isLeafOrCollision l && A.length ary == 1 -> l
+              _ -> BitmapIndexed b (A.update ary i st')
+      where
+        m = mask h s
+        i = sparseIndex b m
+    go h k s t@(Full ary) = do
+      let !st = A.index ary i
+          !st' = go h k (nextShift s) st
+      if st' `ptrEq` st
+        then t
+        else case st' of
+          Empty ->
+            let ary' = A.delete ary i
+                bm = fullBitmap .&. complement (1 `unsafeShiftL` i)
+             in BitmapIndexed bm ary'
+          _ -> Full (A.update ary i st')
+      where
+        i = index h s
+    go h k s t@(Collision hy ls)
+      | h == hy = case indexOf k ls of
+          Just i -> do
+            let (# L _ v #) = A.index# ls i
+            case f $ Just v of
+              Nothing
+                | A.length ls == 2 ->
+                    if i == 0
+                      then Leaf h (A.index ls 1)
+                      else Leaf h (A.index ls 0)
+                | otherwise -> Collision hy (A.delete ls i)
+              Just v' -> Collision hy $ A.update ls i $ L k v'
+          Nothing -> case f Nothing of
+            Nothing -> t
+            Just v' -> Collision hy $ A.snoc ls $ L k v'
+      | otherwise = case f Nothing of
+          Nothing -> t
+          Just v' -> runST $ two s h k v' hy t
+{-# INLINE alter' #-}
 
--- | \(O(\log n)\)  The expression @('alterF' f k map)@ alters the value @x@ at
--- @k@, or absence thereof.
---
---  'alterF' can be used to insert, delete, or update a value in a map.
---
 -- Note: 'alterF' is a flipped version of the 'at' combinator from
 -- <https://hackage.haskell.org/package/lens/docs/Control-Lens-At.html#v:at Control.Lens.At>.
 --
diff --git a/unordered-containers.cabal b/unordered-containers.cabal
index 0271e963..96996548 100644
--- a/unordered-containers.cabal
+++ b/unordered-containers.cabal
@@ -68,7 +68,8 @@ library
     MagicHash,
     BangPatterns
 
-  ghc-options: -Wall -O2 -fwarn-tabs -ferror-spans
+  -- ghc-options: -Wall -O2 -fwarn-tabs -ferror-spans
+  ghc-options: -Wall -fwarn-tabs -ferror-spans
 
   -- For dumping the generated code:
   -- ghc-options: -ddump-simpl -ddump-stg-final -ddump-cmm -ddump-asm -ddump-to-file

From a92b48a7d2ef33f514c025db53734acd8686581d Mon Sep 17 00:00:00 2001
From: Brian Shu <littlebubu.shu@gmail.com>
Date: Thu, 26 May 2022 13:29:36 -0400
Subject: [PATCH 2/6] remove redundant constraints

---
 Data/HashMap/Internal.hs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
index 2aeb87a3..5323a2d4 100644
--- a/Data/HashMap/Internal.hs
+++ b/Data/HashMap/Internal.hs
@@ -1285,7 +1285,7 @@ adjust# f k0 m0 = go h0 k0 0 m0
 -- | \(O(\log n)\)  The expression @('update' f k map)@ updates the value @x@ at @k@
 -- (if it is in the map). If @(f x)@ is 'Nothing', the element is deleted.
 -- If it is @('Just' y)@, the key @k@ is bound to the new value @y@.
-update :: (Eq k, Hashable k, Show k, Show a) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a
+update :: (Eq k, Hashable k) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a
 update f = Exts.inline alter (>>= f)
 {-# INLINABLE update #-}
 
@@ -1297,11 +1297,11 @@ update f = Exts.inline alter (>>= f)
 -- @
 -- 'lookup' k ('alter' f k m) = f ('lookup' k m)
 -- @
-alter :: (Eq k, Hashable k, Show k, Show v) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
+alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
 alter f k = alter' f (hash k) k
 {-# INLINEABLE alter #-}
 
-alter' :: (Eq k, Show v) => (Maybe v -> Maybe v) -> Hash -> k -> HashMap k v -> HashMap k v
+alter' :: Eq k => (Maybe v -> Maybe v) -> Hash -> k -> HashMap k v -> HashMap k v
 alter' f h0 k0 m0 = go h0 k0 0 m0
   where
     go !h !k !_ Empty = case f Nothing of

From 440004316b134e86e6f27fe4367ed1c8971efec3 Mon Sep 17 00:00:00 2001
From: oberblastmeister <61095988+oberblastmeister@users.noreply.github.com>
Date: Fri, 27 May 2022 13:50:54 -0400
Subject: [PATCH 3/6] Update Data/HashMap/Internal.hs

Co-authored-by: Simon Jakobi <simon.jakobi@gmail.com>
---
 Data/HashMap/Internal.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
index 5323a2d4..97ec6a84 100644
--- a/Data/HashMap/Internal.hs
+++ b/Data/HashMap/Internal.hs
@@ -1299,7 +1299,7 @@ update f = Exts.inline alter (>>= f)
 -- @
 alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
 alter f k = alter' f (hash k) k
-{-# INLINEABLE alter #-}
+{-# INLINABLE alter #-}
 
 alter' :: Eq k => (Maybe v -> Maybe v) -> Hash -> k -> HashMap k v -> HashMap k v
 alter' f h0 k0 m0 = go h0 k0 0 m0

From fe15fe28bb9d0e8f525de2052b20cd8331867259 Mon Sep 17 00:00:00 2001
From: Brian Shu <littlebubu.shu@gmail.com>
Date: Sat, 28 May 2022 10:05:04 -0400
Subject: [PATCH 4/6] add to strict HashMap

---
 Data/HashMap/Internal.hs        |  2 +-
 Data/HashMap/Internal/Strict.hs | 18 ++++--------------
 unordered-containers.cabal      |  4 ++--
 3 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
index 97ec6a84..aeb0570e 100644
--- a/Data/HashMap/Internal.hs
+++ b/Data/HashMap/Internal.hs
@@ -1298,7 +1298,7 @@ update f = Exts.inline alter (>>= f)
 -- 'lookup' k ('alter' f k m) = f ('lookup' k m)
 -- @
 alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
-alter f k = alter' f (hash k) k
+alter f k = Exts.inline alter' f (hash k) k
 {-# INLINABLE alter #-}
 
 alter' :: Eq k => (Maybe v -> Maybe v) -> Hash -> k -> HashMap k v -> HashMap k v
diff --git a/Data/HashMap/Internal/Strict.hs b/Data/HashMap/Internal/Strict.hs
index ed87b1ec..a2b3e3ae 100644
--- a/Data/HashMap/Internal/Strict.hs
+++ b/Data/HashMap/Internal/Strict.hs
@@ -293,7 +293,7 @@ adjust f k0 m0 = go h0 k0 0 m0
 -- (if it is in the map). If @(f x)@ is 'Nothing', the element is deleted.
 -- If it is @('Just' y)@, the key @k@ is bound to the new value @y@.
 update :: (Eq k, Hashable k) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a
-update f = alter (>>= f)
+update f = Exts.inline alter (>>= f)
 {-# INLINABLE update #-}
 
 -- | \(O(\log n)\)  The expression @('alter' f k map)@ alters the value @x@ at @k@, or
@@ -305,19 +305,9 @@ update f = alter (>>= f)
 -- 'lookup' k ('alter' f k m) = f ('lookup' k m)
 -- @
 alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
-alter f k m =
-    let !h = hash k
-        !lookupRes = HM.lookupRecordCollision h k m
-    in case f (HM.lookupResToMaybe lookupRes) of
-        Nothing -> case lookupRes of
-            Absent            -> m
-            Present _ collPos -> HM.deleteKeyExists collPos h k m
-        Just !v' -> case lookupRes of
-            Absent             -> HM.insertNewKey h k v' m
-            Present v collPos ->
-                if v `ptrEq` v'
-                    then m
-                    else HM.insertKeyExists collPos h k v' m
+alter f = Exts.inline HM.alter $ \m -> case f m of
+  Nothing -> Nothing
+  Just !x -> Just x
 {-# INLINABLE alter #-}
 
 -- | \(O(\log n)\)  The expression (@'alterF' f k map@) alters the value @x@ at
diff --git a/unordered-containers.cabal b/unordered-containers.cabal
index 1614d5a1..e1a292fa 100644
--- a/unordered-containers.cabal
+++ b/unordered-containers.cabal
@@ -68,8 +68,8 @@ library
     MagicHash,
     BangPatterns
 
-  -- ghc-options: -Wall -O2 -fwarn-tabs -ferror-spans
-  ghc-options: -Wall -fwarn-tabs -ferror-spans
+  ghc-options: -Wall -O2 -fwarn-tabs -ferror-spans
+  -- ghc-options: -Wall -fwarn-tabs -ferror-spans
 
   -- For dumping the generated code:
   -- ghc-options: -ddump-simpl -ddump-stg-final -ddump-cmm -ddump-asm -ddump-to-file

From a87fecafdceeb6e602d459a2a2fc2a18286c3fa5 Mon Sep 17 00:00:00 2001
From: Brian Shu <littlebubu.shu@gmail.com>
Date: Sat, 28 May 2022 10:07:10 -0400
Subject: [PATCH 5/6] bang pattern

---
 Data/HashMap/Internal.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
index aeb0570e..3e90d226 100644
--- a/Data/HashMap/Internal.hs
+++ b/Data/HashMap/Internal.hs
@@ -1363,7 +1363,7 @@ alter' f h0 k0 m0 = go h0 k0 0 m0
     go h k s t@(Collision hy ls)
       | h == hy = case indexOf k ls of
           Just i -> do
-            let (# L _ v #) = A.index# ls i
+            let !(L _ v) = A.index ls i
             case f $ Just v of
               Nothing
                 | A.length ls == 2 ->

From a74d3af096e9f2edf085976cc58742e060e0734e Mon Sep 17 00:00:00 2001
From: Brian Shu <littlebubu.shu@gmail.com>
Date: Mon, 6 Jun 2022 10:56:19 -0400
Subject: [PATCH 6/6] remove use of two for now for insert'

---
 Data/HashMap/Internal.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
index 3e90d226..a5a7a045 100644
--- a/Data/HashMap/Internal.hs
+++ b/Data/HashMap/Internal.hs
@@ -834,7 +834,7 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
       where i = index h s
     go h k x s t@(Collision hy v)
         | h == hy   = Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v)
-        | otherwise = runST $ two s h k x hy t
+        | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
 {-# INLINABLE insert' #-}
 
 -- Insert optimized for the case when we know the key is not in the map.