Introduce size-aware 'deleteInternal' function

Author: rockbmb

Data/HashMap/Base.hs

@@ -2,6 +2,7 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE PatternGuards #-}
 {-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-}
 
@@ -889,6 +890,62 @@ delete k0 m0 = go h0 k0 0 m0
         | otherwise = t
 {-# INLINABLE delete #-}
 
+-- | /O(log n)/ Remove the mapping for the specified key from this map
+-- if present. Returns a tuple with the hashmap's change in size and the
+-- hashmap after the deletion.
+deleteInternal :: (Eq k, Hashable k) => k -> HashMap k v -> (Int, HashMap k v)
+deleteInternal k0 m0 = go h0 k0 0 m0
+  where
+    h0 = hash k0
+    go !_ !_ !_ Empty = (0, Empty)
+    go h k _ t@(Leaf hy (L ky _))
+        | hy == h && ky == k = (-1, Empty)
+        | otherwise          = (0, t)
+    go h k s t@(BitmapIndexed b ary)
+        | b .&. m == 0 = (0, t)
+        | otherwise =
+            let !st = A.index ary i
+                (!sz, !st') = go h k (s+bitsPerSubkey) st
+            in (sz,) (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) =
+        let !st   = A.index ary i
+            (!sz, !st') = go h k (s+bitsPerSubkey) st
+        in (sz,) (if st' `ptrEq` st
+            then t
+            else case st' of
+            Empty ->
+                let ary' = A.delete ary i
+                    bm   = fullNodeMask .&. complement (1 `unsafeShiftL` i)
+                in BitmapIndexed bm ary'
+            _ -> Full (A.update ary i st'))
+      where i = index h s
+    go h k _ t@(Collision hy v)
+        | h == hy = case indexOf k v of
+            Just i
+                | A.length v == 2 ->
+                    (-1,) (if i == 0
+                    then Leaf h (A.index v 1)
+                    else Leaf h (A.index v 0))
+                | otherwise -> (-1, Collision h (A.delete v i))
+            Nothing -> (0, t)
+        | otherwise = (0, t)
+{-# INLINABLE deleteInternal #-}
+
 -- | /O(log n)/ Adjust the value tied to a given key in this map only
 -- if it is present. Otherwise, leave the map alone.
 adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v