Introduce size-aware 'insertWithInternal' function

Author: rockbmb

Data/HashMap/Base.hs

@@ -709,6 +709,50 @@ insertWith f k0 v0 m0 = go h0 k0 v0 0 m0
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
 {-# INLINABLE insertWith #-}
 
+-- | /O(log n)/ Associate the value with the key in this map.  If
+-- this map previously contained a mapping for the key, the old value
+-- is replaced by the result of applying the given function to the new
+-- and old value.  Example:
+--
+-- > insertWithInternal f k v map
+-- >   where f new old = new + old
+insertWithInternal :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
+            -> (Int, HashMap k v)
+insertWithInternal f k0 v0 m0 = go h0 k0 v0 0 m0
+  where
+    h0 = hash k0
+    go !h !k x !_ Empty = (1, Leaf h (L k x))
+    go h k x s (Leaf hy l@(L ky y))
+        | hy == h = if ky == k
+                    then (0, Leaf h (L k (f x y)))
+                    else (1, collision h l (L k x))
+        | otherwise = (0, runST (two s h k x hy ky y))
+    go h k x s (BitmapIndexed b ary)
+        | b .&. m == 0 =
+            let ary' = A.insert ary i $! Leaf h (L k x)
+            in (1, bitmapIndexedOrFull (b .|. m) ary')
+        | otherwise =
+            let st   = A.index ary i
+                (sz, st')  = go h k x (s+bitsPerSubkey) st
+                ary' = A.update ary i $! st'
+            in (sz, BitmapIndexed b ary')
+      where m = mask h s
+            i = sparseIndex b m
+    go h k x s (Full ary) =
+        let st   = A.index ary i
+            (sz, st')  = go h k x (s+bitsPerSubkey) st
+            ary' = update16 ary i $! st'
+        in (sz, Full ary')
+      where i = index h s
+    go h k x s t@(Collision hy v)
+        | h == hy   =
+            let !start = A.length v
+                !newV = updateOrSnocWith const k x v
+                !end = A.length newV
+            in (end - start, Collision h newV)
+        | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
+{-# INLINABLE insertWithInternal #-}
+
 -- | In-place update version of insertWith
 unsafeInsertWith :: forall k v. (Eq k, Hashable k)
                  => (v -> v -> v) -> k -> v -> HashMap k v