Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 43 additions & 57 deletions Data/HashMap/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,7 @@ module Data.HashMap.Internal
, deleteKeyExists
, insertModifying
, ptrEq
, hashOfLeafOrCollision
, adjust#
) where

Expand Down Expand Up @@ -1625,7 +1626,7 @@ unionWithKey f = go 0
go (nextShift s) st1 t2
in BitmapIndexed b1 ary'
where
h2 = leafHashCode t2
h2 = hashOfLeafOrCollision t2
m2 = mask h2 s
i = sparseIndex b1 m2
go s t1 (BitmapIndexed b2 ary2)
Expand All @@ -1636,24 +1637,20 @@ unionWithKey f = go 0
go (nextShift s) t1 st2
in BitmapIndexed b2 ary'
where
h1 = leafHashCode t1
h1 = hashOfLeafOrCollision t1
m1 = mask h1 s
i = sparseIndex b2 m1
go s (Full ary1) t2 =
let h2 = leafHashCode t2
let h2 = hashOfLeafOrCollision t2
i = index h2 s
ary' = updateFullArrayWith' ary1 i $ \st1 -> go (nextShift s) st1 t2
in Full ary'
go s t1 (Full ary2) =
let h1 = leafHashCode t1
let h1 = hashOfLeafOrCollision t1
i = index h1 s
ary' = updateFullArrayWith' ary2 i $ \st2 -> go (nextShift s) t1 st2
in Full ary'

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)
Expand Down Expand Up @@ -1942,7 +1939,26 @@ differenceWithKey f = go_differenceWithKey 0
go_differenceWithKey _s a@(Collision hA aryA) (Leaf hB (L kB vB))
| hA == hB = updateCollision (\vA -> f kB vA vB) hA kB aryA a
| otherwise = a
go_differenceWithKey s a@(BitmapIndexed bA aryA) b@(Leaf hB _)
go_differenceWithKey s a@(Collision hA _) (BitmapIndexed bB aryB)
| bB .&. m == 0 = a
| otherwise =
case A.index# aryB (sparseIndex bB m) of
(# stB #) -> go_differenceWithKey (nextShift s) a stB
where m = mask hA s
go_differenceWithKey s a@(Collision hA _) (Full aryB)
= case A.index# aryB (index hA s) of
(# stB #) -> go_differenceWithKey (nextShift s) a stB
go_differenceWithKey s a@(BitmapIndexed bA aryA) (BitmapIndexed bB aryB)
= differenceWithKey_Arrays s bA aryA a bB aryB
go_differenceWithKey s a@(Full aryA) (BitmapIndexed bB aryB)
= differenceWithKey_Arrays s fullBitmap aryA a bB aryB
go_differenceWithKey s a@(BitmapIndexed bA aryA) (Full aryB)
= differenceWithKey_Arrays s bA aryA a fullBitmap aryB
go_differenceWithKey s a@(Full aryA) (Full aryB)
= differenceWithKey_Arrays s fullBitmap aryA a fullBitmap aryB
go_differenceWithKey _s a@(Collision hA aryA) (Collision hB aryB)
= differenceWithKey_Collisions f hA aryA a hB aryB
go_differenceWithKey s a@(BitmapIndexed bA aryA) b
| bA .&. m == 0 = a
| otherwise = case A.index# aryA i of
(# !stA #) -> case go_differenceWithKey (nextShift s) stA b of
Expand All @@ -1956,26 +1972,10 @@ differenceWithKey f = go_differenceWithKey 0
| stA `ptrEq` stA' -> a
| otherwise -> BitmapIndexed bA (A.update aryA i stA')
where
hB = hashOfLeafOrCollision b
m = mask hB s
i = sparseIndex bA m
go_differenceWithKey s a@(BitmapIndexed bA aryA) b@(Collision hB _)
| bA .&. m == 0 = a
| otherwise =
case A.index# aryA i of
(# !st #) -> case go_differenceWithKey (nextShift s) st b of
Empty | A.length aryA == 2
, (# l #) <- A.index# aryA (otherOfOneOrZero i)
, isLeafOrCollision l
-> l
| otherwise
-> BitmapIndexed (bA .&. complement m) (A.delete aryA i)
st' | isLeafOrCollision st' && A.length aryA == 1 -> st'
| st `ptrEq` st' -> a
| otherwise -> BitmapIndexed bA (A.update aryA i st')
where
m = mask hB s
i = sparseIndex bA m
go_differenceWithKey s a@(Full aryA) b@(Leaf hB _)
go_differenceWithKey s a@(Full aryA) b
= case A.index# aryA i of
(# !stA #) -> case go_differenceWithKey (nextShift s) stA b of
Empty ->
Expand All @@ -1984,36 +1984,9 @@ differenceWithKey f = go_differenceWithKey 0
in BitmapIndexed bm aryA'
stA' | stA `ptrEq` stA' -> a
| otherwise -> Full (updateFullArray aryA i stA')
where i = index hB s
go_differenceWithKey s a@(Full aryA) b@(Collision hB _)
= case A.index# aryA i of
(# !stA #) -> case go_differenceWithKey (nextShift s) stA b of
Empty ->
let aryA' = A.delete aryA i
bm = fullBitmap .&. complement (1 `unsafeShiftL` i)
in BitmapIndexed bm aryA'
stA' | stA `ptrEq` stA' -> a
| otherwise -> Full (updateFullArray aryA i stA')
where i = index hB s
go_differenceWithKey s a@(Collision hA _) (BitmapIndexed bB aryB)
| bB .&. m == 0 = a
| otherwise =
case A.index# aryB (sparseIndex bB m) of
(# stB #) -> go_differenceWithKey (nextShift s) a stB
where m = mask hA s
go_differenceWithKey s a@(Collision hA _) (Full aryB)
= case A.index# aryB (index hA s) of
(# stB #) -> go_differenceWithKey (nextShift s) a stB
go_differenceWithKey s a@(BitmapIndexed bA aryA) (BitmapIndexed bB aryB)
= differenceWithKey_Arrays s bA aryA a bB aryB
go_differenceWithKey s a@(Full aryA) (BitmapIndexed bB aryB)
= differenceWithKey_Arrays s fullBitmap aryA a bB aryB
go_differenceWithKey s a@(BitmapIndexed bA aryA) (Full aryB)
= differenceWithKey_Arrays s bA aryA a fullBitmap aryB
go_differenceWithKey s a@(Full aryA) (Full aryB)
= differenceWithKey_Arrays s fullBitmap aryA a fullBitmap aryB
go_differenceWithKey _s a@(Collision hA aryA) (Collision hB aryB)
= differenceWithKey_Collisions f hA aryA a hB aryB
where
hB = hashOfLeafOrCollision b
i = index hB s

differenceWithKey_Arrays !s !bA !aryA tA !bB !aryB
| bA .&. bB == 0 = tA
Expand Down Expand Up @@ -2841,6 +2814,19 @@ otherOfOneOrZero :: Int -> Int
otherOfOneOrZero i = 1 - i
{-# INLINE otherOfOneOrZero #-}

------------------------------------------------------------------------
-- Tools for reducing duplication in code handling 'Leaf' and 'Collision' nodes

-- | The 'Hash' of a 'Leaf' or 'Collision' node.
--
-- This function is marked @NOINLINE@ to prevent GHC from generating separate
-- alternatives for 'Leaf' and 'Collision' nodes.
hashOfLeafOrCollision :: HashMap k v -> Hash
hashOfLeafOrCollision (Leaf h _) = h
hashOfLeafOrCollision (Collision h _) = h
hashOfLeafOrCollision _ = error "hashOfLeafOrCollision"
{-# NOINLINE hashOfLeafOrCollision #-}

------------------------------------------------------------------------
-- IsList instance
instance (Eq k, Hashable k) => Exts.IsList (HashMap k v) where
Expand Down
12 changes: 4 additions & 8 deletions Data/HashMap/Internal/Strict.hs
Original file line number Diff line number Diff line change
Expand Up @@ -505,7 +505,7 @@ unionWithKey f = go 0
go (nextShift s) st1 t2
in BitmapIndexed b1 ary'
where
h2 = leafHashCode t2
h2 = HM.hashOfLeafOrCollision t2
m2 = mask h2 s
i = sparseIndex b1 m2
go s t1 (BitmapIndexed b2 ary2)
Expand All @@ -516,24 +516,20 @@ unionWithKey f = go 0
go (nextShift s) t1 st2
in BitmapIndexed b2 ary'
where
h1 = leafHashCode t1
h1 = HM.hashOfLeafOrCollision t1
m1 = mask h1 s
i = sparseIndex b2 m1
go s (Full ary1) t2 =
let h2 = leafHashCode t2
let h2 = HM.hashOfLeafOrCollision t2
i = index h2 s
ary' = HM.updateFullArrayWith' ary1 i $ \st1 -> go (nextShift s) st1 t2
in Full ary'
go s t1 (Full ary2) =
let h1 = leafHashCode t1
let h1 = HM.hashOfLeafOrCollision t1
i = index h1 s
ary' = HM.updateFullArrayWith' ary2 i $ \st2 -> go (nextShift s) t1 st2
in Full ary'

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)
Expand Down