diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs index 6cad1028..7726ae32 100644 --- a/Data/HashMap/Internal.hs +++ b/Data/HashMap/Internal.hs @@ -126,6 +126,7 @@ module Data.HashMap.Internal , nextShift , sparseIndex , two + , two' , unionArrayBy , updateFullArray , updateFullArrayM @@ -164,7 +165,7 @@ import Data.Hashable (Hashable) import Data.Hashable.Lifted (Hashable1, Hashable2) import Data.HashMap.Internal.List (isPermutationBy, unorderedCompare) import Data.Semigroup (Semigroup (..), stimesIdempotentMonoid) -import GHC.Exts (Int (..), Int#, TYPE, (==#)) +import GHC.Exts (Int (..), Int#, TYPE, Word (..), (==#)) import GHC.Stack (HasCallStack) import Prelude hiding (Foldable (..), filter, lookup, map, pred) @@ -832,7 +833,7 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0 where i = index h s go h k x s t@(Collision hy v) | h == hy = Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v) - | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) + | otherwise = runST (two s h k x hy t) {-# INLINABLE insert' #-} -- | Insert optimized for the case when we know the key is not in the map. @@ -866,8 +867,7 @@ insertNewKey !h0 !k0 x0 !m0 = go h0 k0 x0 0 m0 where i = index h s go h k x s t@(Collision hy v) | h == hy = Collision h (A.snoc v (L k x)) - | otherwise = - go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) + | otherwise = runST (two s h k x hy t) {-# NOINLINE insertNewKey #-} @@ -884,33 +884,19 @@ insertKeyExists !collPos0 !h0 !k0 x0 !m0 = go collPos0 h0 k0 x0 m0 = Leaf h (L k x) go collPos shiftedHash k x (BitmapIndexed b ary) = let !st = A.index ary i - !st' = go collPos (shiftHash shiftedHash) k x st + !st' = go collPos (nextSH shiftedHash) k x st in BitmapIndexed b (A.update ary i st') - where m = mask' shiftedHash + where m = maskSH shiftedHash i = sparseIndex b m go collPos shiftedHash k x (Full ary) = let !st = A.index ary i - !st' = go collPos (shiftHash shiftedHash) k x st + !st' = go collPos (nextSH shiftedHash) k x st in Full (updateFullArray ary i st') - where i = index' shiftedHash + where i = indexSH shiftedHash go collPos _shiftedHash k x (Collision h v) | collPos >= 0 = Collision h (setAtPosition collPos k x v) | otherwise = Empty -- error "Internal error: go {collPos negative}" go _ _ _ _ Empty = Empty -- error "Internal error: go Empty" - - -- Customized version of 'index' that doesn't require a 'Shift'. - index' :: Hash -> Int - index' w = fromIntegral $ w .&. subkeyMask - {-# INLINE index' #-} - - -- Customized version of 'mask' that doesn't require a 'Shift'. - mask' :: Word -> Bitmap - mask' w = 1 `unsafeShiftL` index' w - {-# INLINE mask' #-} - - shiftHash h = h `unsafeShiftR` bitsPerSubkey - {-# INLINE shiftHash #-} - {-# NOINLINE insertKeyExists #-} -- | Replace the ith Leaf with Leaf k v. @@ -953,44 +939,52 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0) where i = index h s go h k x s t@(Collision hy v) | h == hy = return $! Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v) - | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) + | otherwise = two s h k x hy t {-# INLINABLE unsafeInsert #-} --- | 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. +-- | Create a map from a key-value pair and a 'Leaf' or 'Collision' node with a +-- different hash. -- --- Note: to avoid silly thunks, this function must be strict in the --- key. See issue #232. We don't need to force the HashMap argument --- because it's already in WHNF (having just been matched) and we --- just put it directly in an array. +-- It is the caller's responsibility to ensure that the HashMap argument is in +-- WHNF. two :: Shift -> Hash -> k -> v -> Hash -> HashMap k v -> ST s (HashMap k v) -two = go - where - go s h1 k1 v1 h2 t2 - | bp1 == bp2 = do - st <- go (nextShift s) h1 k1 v1 h2 t2 - ary <- A.singletonM st - return $ BitmapIndexed bp1 ary - | otherwise = do - mary <- A.new 2 $! Leaf h1 (L k1 v1) - A.write mary idx2 t2 - ary <- A.unsafeFreeze mary - return $ BitmapIndexed (bp1 .|. bp2) ary - where - bp1 = mask h1 s - bp2 = mask h2 s - !(I# i1) = index h1 s - !(I# i2) = index h2 s - idx2 = I# (i1 Exts.<# i2) - -- This way of computing idx2 saves us a branch compared to the previous approach: - -- - -- idx2 | index h1 s < index h2 s = 1 - -- | otherwise = 0 - -- - -- See https://github.com/haskell-unordered-containers/unordered-containers/issues/75#issuecomment-1128419337 +two s h1 k1 v1 = two' s h1 l + where !l = Leaf h1 (L k1 v1) {-# INLINE two #-} +-- | Create a map from two 'Leaf' or 'Collision' nodes whose hashes are +-- distinct. +-- +-- It is the caller's responsibility to ensure that both HashMap arguments are +-- in WHNF. +two' :: Shift -> Hash -> HashMap k v -> Hash -> HashMap k v -> ST s (HashMap k v) +two' s h1 lc1 h2 lc2 = two_go (shiftHash s h1) lc1 (shiftHash s h2) lc2 +{-# INLINE two' #-} + +-- | This function lives at the top-level so 'two' and `two'` can be inlined +-- without inlining this loop. +two_go :: ShiftedHash -> HashMap k v -> ShiftedHash -> HashMap k v -> ST s (HashMap k v) +two_go !sh1 t1 !sh2 t2 + | bp1 == bp2 = do + st <- two_go (nextSH sh1) t1 (nextSH sh2) t2 + ary <- A.singletonM st + return $ BitmapIndexed bp1 ary + | otherwise = do + mary <- A.new 2 t1 + A.write mary idx2 t2 + ary <- A.unsafeFreeze mary + return $ BitmapIndexed (bp1 .|. bp2) ary + where + !bp1@(W# bp1#) = maskSH sh1 + !bp2@(W# bp2#) = maskSH sh2 + idx2 = I# (bp1# `Exts.ltWord#` bp2#) + -- This way of computing idx2 saves us a branch compared to the previous approach: + -- + -- idx2 | index h1 s < index h2 s = 1 + -- | otherwise = 0 + -- + -- See https://github.com/haskell-unordered-containers/unordered-containers/issues/75#issuecomment-1128419337 + -- | \(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 @@ -1050,7 +1044,7 @@ insertModifying x f k0 m0 = go h0 k0 0 m0 in if A.unsafeSameArray v v' then t else Collision h v' - | otherwise = go h k s $ BitmapIndexed (mask hy s) (A.singleton t) + | otherwise = runST (two s h k x hy t) {-# INLINABLE insertModifying #-} -- | Like insertModifying for arrays; used to implement insertModifying @@ -1109,7 +1103,7 @@ unsafeInsertWithKey f k0 v0 m0 = runST (go h0 k0 v0 0 m0) where i = index h s go h k x s t@(Collision hy v) | h == hy = return $! Collision h (updateOrSnocWithKey f k x v) - | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) + | otherwise = two s h k x hy t {-# INLINABLE unsafeInsertWithKey #-} -- | \(O(\log n)\) Remove the mapping for the specified key from this map @@ -1178,11 +1172,11 @@ delete' h0 k0 m0 = go h0 k0 0 m0 deleteKeyExists :: Int -> Hash -> k -> HashMap k v -> HashMap k v deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 m0 where - go :: Int -> Word -> k -> HashMap k v -> HashMap k v + go :: Int -> ShiftedHash -> k -> HashMap k v -> HashMap k v go !_collPos !_shiftedHash !_k (Leaf _ _) = Empty go collPos shiftedHash k (BitmapIndexed b ary) = let !st = A.index ary i - !st' = go collPos (shiftHash shiftedHash) k st + !st' = go collPos (nextSH shiftedHash) k st in case st' of Empty | A.length ary == 1 -> Empty | A.length ary == 2 -> @@ -1195,18 +1189,18 @@ deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 m0 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' shiftedHash + where m = maskSH shiftedHash i = sparseIndex b m go collPos shiftedHash k (Full ary) = let !st = A.index ary i - !st' = go collPos (shiftHash shiftedHash) k st + !st' = go collPos (nextSH shiftedHash) k st in case st' of Empty -> let ary' = A.delete ary i bm = fullBitmap .&. complement (1 `unsafeShiftL` i) in BitmapIndexed bm ary' _ -> Full (A.update ary i st') - where i = index' shiftedHash + where i = indexSH shiftedHash go collPos _shiftedHash _k (Collision h v) | A.length v == 2 = if collPos == 0 @@ -1214,20 +1208,6 @@ deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 m0 else Leaf h (A.index v 0) | otherwise = Collision h (A.delete v collPos) go !_ !_ !_ Empty = Empty -- error "Internal error: deleteKeyExists empty" - - -- Customized version of 'index' that doesn't require a 'Shift'. - index' :: Hash -> Int - index' w = fromIntegral $ w .&. subkeyMask - {-# INLINE index' #-} - - -- Customized version of 'mask' that doesn't require a 'Shift'. - mask' :: Word -> Bitmap - mask' w = 1 `unsafeShiftL` index' w - {-# INLINE mask' #-} - - shiftHash h = h `unsafeShiftR` bitsPerSubkey - {-# INLINE shiftHash #-} - {-# NOINLINE deleteKeyExists #-} -- | \(O(\log n)\) Adjust the value tied to a given key in this map only @@ -1610,16 +1590,16 @@ unionWithKey f = go 0 | h1 == h2 = if k1 == k2 then Leaf h1 (L k1 (f k1 v1 v2)) else collision h1 l1 l2 - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (two' s h1 t1 h2 t2) go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2) | h1 == h2 = Collision h1 (updateOrSnocWithKey (\k a b -> (# f k a b #)) k1 v1 ls2) - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (two' s h1 t1 h2 t2) go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2)) | h1 == h2 = Collision h1 (updateOrSnocWithKey (\k a b -> (# f k b a #)) k2 v2 ls1) - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (two' s h1 t1 h2 t2) go s t1@(Collision h1 ls1) t2@(Collision h2 ls2) | h1 == h2 = Collision h1 (updateOrConcatWithKey (\k a b -> (# f k a b #)) ls1 ls2) - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (two' s h1 t1 h2 t2) -- branch vs. branch go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) = let b' = b1 .|. b2 @@ -1672,14 +1652,6 @@ unionWithKey f = go 0 leafHashCode (Leaf h _) = h leafHashCode (Collision h _) = h leafHashCode _ = error "leafHashCode" - - goDifferentHash s h1 h2 t1 t2 - | m1 == m2 = BitmapIndexed m1 (A.singleton $! goDifferentHash (nextShift s) h1 h2 t1 t2) - | m1 < m2 = BitmapIndexed (m1 .|. m2) (A.pair t1 t2) - | otherwise = BitmapIndexed (m1 .|. m2) (A.pair t2 t1) - where - m1 = mask h1 s - m2 = mask h2 s {-# INLINE unionWithKey #-} -- | Strict in the result of @f@. @@ -2510,6 +2482,34 @@ nextShift :: Shift -> Shift nextShift s = s + bitsPerSubkey {-# INLINE nextShift #-} +------------------------------------------------------------------------ +-- ShiftedHash + +-- | Sometimes it's more efficient to right-shift the hashes directly instead +-- of keeping track of an additional 'Shift' value. +type ShiftedHash = Hash + +-- | Construct a 'ShiftedHash' from a 'Shift' and a 'Hash'. +shiftHash :: Shift -> Hash -> ShiftedHash +shiftHash s h = h `unsafeShiftR` s +{-# INLINE shiftHash #-} + +-- | Update a 'ShiftedHash' for the next level of the tree. +nextSH :: ShiftedHash -> ShiftedHash +nextSH sh = sh `unsafeShiftR` bitsPerSubkey +{-# INLINE nextSH #-} + +-- | Version of 'index' for use with @'ShiftedHash'es@. +indexSH :: ShiftedHash -> Int +indexSH sh = fromIntegral $ sh .&. subkeyMask +{-# INLINE indexSH #-} + +-- | Version of 'mask' for use with @'ShiftedHash'es@. +maskSH :: ShiftedHash -> Bitmap +maskSH sh = 1 `unsafeShiftL` indexSH sh +{-# INLINE maskSH #-} + + ------------------------------------------------------------------------ -- Pointer equality diff --git a/Data/HashMap/Internal/Strict.hs b/Data/HashMap/Internal/Strict.hs index ce9a48fa..1c064ab1 100644 --- a/Data/HashMap/Internal/Strict.hs +++ b/Data/HashMap/Internal/Strict.hs @@ -216,7 +216,7 @@ insertWith f k0 v0 m0 = go h0 k0 v0 0 m0 where i = index h s go h k x s t@(Collision hy v) | h == hy = Collision h (updateOrSnocWith f k x v) - | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) + | otherwise = x `seq` runST (HM.two s h k x hy t) {-# INLINABLE insertWith #-} -- | In-place update version of insertWith @@ -257,7 +257,7 @@ unsafeInsertWithKey f k0 v0 m0 = runST (go h0 k0 v0 0 m0) where i = index h s go h k x s t@(Collision hy v) | h == hy = return $! Collision h (updateOrSnocWithKey f k x v) - | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) + | otherwise = x `seq` HM.two s h k x hy t {-# INLINABLE unsafeInsertWithKey #-} -- | \(O(\log n)\) Adjust the value tied to a given key in this map only @@ -465,16 +465,16 @@ unionWithKey f = go 0 | h1 == h2 = if k1 == k2 then leaf h1 k1 (f k1 v1 v2) else HM.collision h1 l1 l2 - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (HM.two' s h1 t1 h2 t2) go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2) | h1 == h2 = Collision h1 (updateOrSnocWithKey f k1 v1 ls2) - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (HM.two' s h1 t1 h2 t2) go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2)) | h1 == h2 = Collision h1 (updateOrSnocWithKey (flip . f) k2 v2 ls1) - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (HM.two' s h1 t1 h2 t2) go s t1@(Collision h1 ls1) t2@(Collision h2 ls2) | h1 == h2 = Collision h1 (HM.updateOrConcatWithKey (\k a b -> let !v = f k a b in (# v #)) ls1 ls2) - | otherwise = goDifferentHash s h1 h2 t1 t2 + | otherwise = runST (HM.two' s h1 t1 h2 t2) -- branch vs. branch go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) = let b' = b1 .|. b2 @@ -527,14 +527,6 @@ unionWithKey f = go 0 leafHashCode (Leaf h _) = h leafHashCode (Collision h _) = h leafHashCode _ = error "leafHashCode" - - goDifferentHash s h1 h2 t1 t2 - | m1 == m2 = BitmapIndexed m1 (A.singleton $! goDifferentHash (nextShift s) h1 h2 t1 t2) - | m1 < m2 = BitmapIndexed (m1 .|. m2) (A.pair t1 t2) - | otherwise = BitmapIndexed (m1 .|. m2) (A.pair t2 t1) - where - m1 = mask h1 s - m2 = mask h2 s {-# INLINE unionWithKey #-} ------------------------------------------------------------------------