Introduce size-aware 'unsafeInsertInternal' function

Author: rockbmb

Data/HashMap/Base.hs

@@ -527,6 +527,10 @@ insert k0 v0 m0 = go h0 k0 v0 0 m0
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
 {-# INLINABLE insert #-}
 
+-- | /O(log n)/ Associate the specified value with the specified
+-- key in this map.  If this map previously contained a mapping for
+-- the key, the old value is replaced. Returns a tuple containing the
+-- hashmap's change in size, and the hashmap after the insertion.
 insertInternal :: (Eq k, Hashable k) => k -> v -> HashMap k v -> (Int, HashMap k v)
 insertInternal k0 v0 m0 = go h0 k0 v0 0 m0
   where
@@ -602,6 +606,48 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
 {-# INLINABLE unsafeInsert #-}
 
+-- | In-place update version of insert. Returns a tuple with the
+-- HashMap's change in size and the hashmap itself.
+unsafeInsertInternal :: (Eq k, Hashable k) => k -> v -> HashMap k v -> (Int, HashMap k v)
+unsafeInsertInternal k0 v0 m0 = runST (go h0 k0 v0 0 m0)
+  where
+    h0 = hash k0
+    go !h !k x !_ Empty = return $! (1, Leaf h (L k x))
+    go h k x s t@(Leaf hy l@(L ky y))
+        | hy == h = if ky == k
+                    then if x `ptrEq` y
+                         then return (0, t)
+                         else return $! (0, Leaf h (L k x))
+                    else return $! (1, collision h l (L k x))
+        | otherwise = do
+            twoHm <- two s h k x hy ky y
+            return $! (1, twoHm)
+    go h k x s t@(BitmapIndexed b ary)
+        | b .&. m == 0 = do
+            ary' <- A.insertM ary i $! Leaf h (L k x)
+            return $! (1, bitmapIndexedOrFull (b .|. m) ary')
+        | otherwise = do
+            st <- A.indexM ary i
+            (sz, st') <- go h k x (s+bitsPerSubkey) st
+            A.unsafeUpdateM ary i st'
+            return (sz, t)
+      where m = mask h s
+            i = sparseIndex b m
+    go h k x s t@(Full ary) = do
+        st <- A.indexM ary i
+        (sz, st') <- go h k x (s+bitsPerSubkey) st
+        A.unsafeUpdateM ary i st'
+        return (sz, t)
+      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 return $! (end - start, Collision h newV)
+        | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
+{-# INLINABLE unsafeInsertInternal #-}
+
 -- | Create a map from two key-value pairs which hashes don't collide.
 two :: Shift -> Hash -> k -> v -> Hash -> k -> v -> ST s (HashMap k v)
 two = go