Skip to content

Commit 324497b

Browse files
committed
Refactor two
* Use the "shifted hash" approach. * Create a variant `two'` that can be used to address #468.
1 parent aa19655 commit 324497b

File tree

1 file changed

+17
-7
lines changed

1 file changed

+17
-7
lines changed

Data/HashMap/Internal.hs

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -964,29 +964,39 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
964964
-- because it's already in WHNF (having just been matched) and we
965965
-- just put it directly in an array.
966966
two :: Shift -> Hash -> k -> v -> Hash -> HashMap k v -> ST s (HashMap k v)
967-
two = go
967+
two s h1 k1 v1 = two' s h1 (Leaf h1 (L k1 v1))
968+
{-# INLINE two #-}
969+
970+
two' :: Shift -> Hash -> HashMap k v -> Hash -> HashMap k v -> ST s (HashMap k v)
971+
two' s h1 lc1 h2 lc2 = go (shiftHash h1 s) lc1 (shiftHash h2 s) lc2
968972
where
969-
go s h1 k1 v1 h2 t2
973+
go !sh1 t1 !sh2 t2
970974
| bp1 == bp2 = do
971-
st <- go (nextShift s) h1 k1 v1 h2 t2
975+
st <- go (shiftHash sh1 bitsPerSubkey) t1 (shiftHash sh2 bitsPerSubkey) t2
972976
ary <- A.singletonM st
973977
return $ BitmapIndexed bp1 ary
974978
| otherwise = do
975-
mary <- A.new 2 $! Leaf h1 (L k1 v1)
979+
mary <- A.new 2 t1
976980
A.write mary idx2 t2
977981
ary <- A.unsafeFreeze mary
978982
return $ BitmapIndexed (bp1 .|. bp2) ary
979983
where
980-
!bp1@(W# bp1#) = mask h1 s
981-
!bp2@(W# bp2#) = mask h2 s
984+
!bp1@(W# bp1#) = mask' sh1
985+
!bp2@(W# bp2#) = mask' sh2
982986
idx2 = I# (bp1# `Exts.ltWord#` bp2#)
983987
-- This way of computing idx2 saves us a branch compared to the previous approach:
984988
--
985989
-- idx2 | index h1 s < index h2 s = 1
986990
-- | otherwise = 0
987991
--
988992
-- See https://github.com/haskell-unordered-containers/unordered-containers/issues/75#issuecomment-1128419337
989-
{-# INLINE two #-}
993+
994+
shiftHash :: Hash -> Shift -> Word -- type ShiftedHash
995+
shiftHash h n = h `unsafeShiftR` n
996+
997+
mask' :: Word -> Bitmap
998+
mask' w = 1 `unsafeShiftL` fromIntegral (w .&. subkeyMask)
999+
{-# INLINE two' #-} -- Really?!
9901000

9911001
-- | \(O(\log n)\) Associate the value with the key in this map. If
9921002
-- this map previously contained a mapping for the key, the old value

0 commit comments

Comments
 (0)