Introduce size-aware 'insertInternal' function

Author: rockbmb

Data/HashMap/Base.hs

@@ -157,6 +157,10 @@ data HashMap k v
 
 type role HashMap nominal representational
 
+-- | WIP. This will become the user-facing 'HashMap' after this PR is
+-- finalized.
+data HashMapW = HashMapW {-# UNPACK #-} !Int !HashMap
+
 instance (NFData k, NFData v) => NFData (HashMap k v) where
     rnf Empty                 = ()
     rnf (BitmapIndexed _ ary) = rnf ary
@@ -523,6 +527,46 @@ 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 #-}
 
+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
+    h0 = hash k0
+    go !h !k x !_ Empty = (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 (0, t)
+                         else (0, Leaf h (L k x))
+                    else (1, collision h l (L k x))
+        | otherwise = (1, runST (two s h k x hy ky y))
+    go h k x s t@(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
+            in if st' `ptrEq` st
+               then (sz, t)
+               else (sz, BitmapIndexed b (A.update ary i st'))
+      where m = mask h s
+            i = sparseIndex b m
+    go h k x s t@(Full ary) =
+        let !st  = A.index ary i
+            (!sz, !st') = go h k x (s+bitsPerSubkey) st
+        in if st' `ptrEq` st
+            then (sz, t)
+            else (sz, Full (update16 ary i st'))
+      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 insertInternal #-}
+
 -- | In-place update version of insert
 unsafeInsert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v
 unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
@@ -1342,7 +1386,7 @@ fullNodeMask :: Bitmap
 fullNodeMask = complement (complement 0 `unsafeShiftL` maxChildren)
 {-# INLINE fullNodeMask #-}
 
--- | Check if two the two arguments are the same value.  N.B. This
+-- | Check if the two arguments are the same value.  N.B. This
 -- function might give false negatives (due to GC moving objects.)
 ptrEq :: a -> a -> Bool
 ptrEq x y = isTrue# (reallyUnsafePtrEquality# x y ==# 1#)