diff --git a/Data/HashMap/Base.hs b/Data/HashMap/Base.hs index a37b2b55..b496e2ae 100644 --- a/Data/HashMap/Base.hs +++ b/Data/HashMap/Base.hs @@ -675,7 +675,7 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0 then t else Leaf h (L k x) else collision h l (L k x) - | otherwise = runST (two s h k x hy ky y) + | otherwise = runST (two s h k x hy t) go h k x s t@(BitmapIndexed b ary) | b .&. m == 0 = let !ary' = A.insert ary i $! Leaf h (L k x) @@ -711,9 +711,9 @@ insertNewKey :: Hash -> k -> v -> HashMap k v -> HashMap k v insertNewKey !h0 !k0 x0 !m0 = go h0 k0 x0 0 m0 where go !h !k x !_ Empty = Leaf h (L k x) - go h k x s (Leaf hy l@(L ky y)) + go h k x s t@(Leaf hy l) | hy == h = collision h l (L k x) - | otherwise = runST (two s h k x hy ky y) + | otherwise = runST (two s h k x hy t) go h k x s (BitmapIndexed b ary) | b .&. m == 0 = let !ary' = A.insert ary i $! Leaf h (L k x) @@ -800,7 +800,7 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0) then return t else return $! Leaf h (L k x) else return $! collision h l (L k x) - | otherwise = two s h k x hy ky y + | otherwise = two s h k x hy t go h k x s t@(BitmapIndexed b ary) | b .&. m == 0 = do ary' <- A.insertM ary i $! Leaf h (L k x) @@ -823,18 +823,20 @@ 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 #-} --- | 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) +-- | Create a map from two key-value pairs which hashes don't collide. To +-- enhance sharing, the second key-value pair is represented by the hash of its +-- key and a singleton HashMap pairing its key with its value. +two :: Shift -> Hash -> k -> v -> Hash -> HashMap k v -> ST s (HashMap k v) two = go where - go s h1 k1 v1 h2 k2 v2 + go s h1 k1 v1 h2 t2 | bp1 == bp2 = do - st <- go (s+bitsPerSubkey) h1 k1 v1 h2 k2 v2 + st <- go (s+bitsPerSubkey) h1 k1 v1 h2 t2 ary <- A.singletonM st - return $! BitmapIndexed bp1 ary + return $ BitmapIndexed bp1 ary | otherwise = do - mary <- A.new 2 $ Leaf h1 (L k1 v1) - A.write mary idx2 $ Leaf h2 (L k2 v2) + mary <- A.new 2 $! Leaf h1 (L k1 v1) + A.write mary idx2 $! t2 ary <- A.unsafeFreeze mary return $! BitmapIndexed (bp1 .|. bp2) ary where @@ -875,7 +877,7 @@ insertModifying x f k0 m0 = go h0 k0 0 m0 (# v' #) | ptrEq y v' -> t | otherwise -> Leaf h (L k (v')) else collision h l (L k x) - | otherwise = runST (two s h k x hy ky y) + | otherwise = runST (two s h k x hy t) go h k s t@(BitmapIndexed b ary) | b .&. m == 0 = let ary' = A.insert ary i $! Leaf h (L k x) @@ -935,11 +937,11 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0) h0 = hash k0 go :: Hash -> k -> v -> Shift -> HashMap k v -> ST s (HashMap k v) go !h !k x !_ Empty = return $! Leaf h (L k x) - go h k x s (Leaf hy l@(L ky y)) + go h k x s t@(Leaf hy l@(L ky y)) | hy == h = if ky == k then return $! Leaf h (L k (f x y)) else return $! collision h l (L k x) - | otherwise = two s h k x hy ky y + | otherwise = two s h k x hy t go h k x s t@(BitmapIndexed b ary) | b .&. m == 0 = do ary' <- A.insertM ary i $! Leaf h (L k x) diff --git a/Data/HashMap/Strict/Base.hs b/Data/HashMap/Strict/Base.hs index 890d18ff..e5b8240e 100644 --- a/Data/HashMap/Strict/Base.hs +++ b/Data/HashMap/Strict/Base.hs @@ -152,11 +152,11 @@ insertWith f k0 v0 m0 = go h0 k0 v0 0 m0 where h0 = hash k0 go !h !k x !_ Empty = leaf h k x - go h k x s (Leaf hy l@(L ky y)) + go h k x s t@(Leaf hy l@(L ky y)) | hy == h = if ky == k then leaf h k (f x y) else x `seq` (collision h l (L k x)) - | otherwise = x `seq` runST (two s h k x hy ky y) + | otherwise = x `seq` runST (two s h k x hy t) go h k x s (BitmapIndexed b ary) | b .&. m == 0 = let ary' = A.insert ary i $! leaf h k x @@ -186,13 +186,13 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0) where h0 = hash k0 go !h !k x !_ Empty = return $! leaf h k x - go h k x s (Leaf hy l@(L ky y)) + go h k x s t@(Leaf hy l@(L ky y)) | hy == h = if ky == k then return $! leaf h k (f x y) else do let l' = x `seq` (L k x) return $! collision h l l' - | otherwise = x `seq` two s h k x hy ky y + | otherwise = x `seq` two s h k x hy t go h k x s t@(BitmapIndexed b ary) | b .&. m == 0 = do ary' <- A.insertM ary i $! leaf h k x