Skip to content

Make intersections much faster #406

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 30 commits into from
Apr 15, 2022
Merged

Make intersections much faster #406

Show file tree
Hide file tree
Changes from 24 commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
21f238b
fast intersection
oberblastmeister Apr 9, 2022
16f1f7f
cleanup
oberblastmeister Apr 9, 2022
bcc13fc
add show back
oberblastmeister Apr 9, 2022
d5262bf
inline
oberblastmeister Apr 9, 2022
a16456b
debug checks
oberblastmeister Apr 9, 2022
f72011c
inline function
oberblastmeister Apr 9, 2022
678a38c
refactor to use snoc
oberblastmeister Apr 9, 2022
ec24215
Try the unboxed result thing
treeowl Apr 9, 2022
767ae6e
Remove redundant internal constraint
treeowl Apr 9, 2022
72510b4
Merge pull request #3 from treeowl/unboxedness
oberblastmeister Apr 9, 2022
fd43ba7
shrink compat
oberblastmeister Apr 9, 2022
3612645
fix import
oberblastmeister Apr 9, 2022
b484042
use clone
oberblastmeister Apr 9, 2022
9e48bc0
oof
oberblastmeister Apr 9, 2022
48119cb
don't shrink to zero
oberblastmeister Apr 9, 2022
d9d295d
Leaf special case
oberblastmeister Apr 9, 2022
88a9c2c
add strict verisons
oberblastmeister Apr 10, 2022
b3cdbd8
Update Data/HashMap/Internal.hs
oberblastmeister Apr 11, 2022
bf9a27f
Update Data/HashSet/Internal.hs
oberblastmeister Apr 11, 2022
1c20739
naming
oberblastmeister Apr 11, 2022
92e4b2a
Exts.inline
oberblastmeister Apr 11, 2022
5a439cc
add haddocks for searchSwap
oberblastmeister Apr 12, 2022
1c118c4
cleanup
oberblastmeister Apr 12, 2022
1256cf3
Update Data/HashMap/Internal/Array.hs
oberblastmeister Apr 12, 2022
b0210c8
Update Data/HashMap/Internal/Array.hs
oberblastmeister Apr 12, 2022
69f8f28
refactor
oberblastmeister Apr 13, 2022
06cc511
formatting
oberblastmeister Apr 13, 2022
d9a50d7
breakup lines
oberblastmeister Apr 13, 2022
d24cc1f
use Exts.inline
oberblastmeister Apr 14, 2022
64f3f2f
Merge branch 'master' into fast-intersection
oberblastmeister Apr 15, 2022
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
194 changes: 149 additions & 45 deletions Data/HashMap/Internal.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ module Data.HashMap.Internal
, intersection
, intersectionWith
, intersectionWithKey
, intersectionWithKey#

-- * Folds
, foldr'
Expand Down Expand Up @@ -143,16 +144,16 @@ import Control.Applicative (Const (..))
import Control.DeepSeq (NFData (..), NFData1 (..), NFData2 (..))
import Control.Monad.ST (ST, runST)
import Data.Bifoldable (Bifoldable (..))
import Data.Bits (complement, popCount, unsafeShiftL,
unsafeShiftR, (.&.), (.|.), countTrailingZeros)
import Data.Bits (complement, countTrailingZeros, popCount,
unsafeShiftL, unsafeShiftR, (.&.), (.|.))
import Data.Coerce (coerce)
import Data.Data (Constr, Data (..), DataType)
import Data.Functor.Classes (Eq1 (..), Eq2 (..), Ord1 (..), Ord2 (..),
Read1 (..), Show1 (..), Show2 (..))
import Data.Functor.Identity (Identity (..))
import Data.HashMap.Internal.List (isPermutationBy, unorderedCompare)
import Data.Hashable (Hashable)
import Data.Hashable.Lifted (Hashable1, Hashable2)
import Data.HashMap.Internal.List (isPermutationBy, unorderedCompare)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FWIW, the changed sorting of imports is probably due to haskell/stylish-haskell#385, which was recently released.

import Data.Semigroup (Semigroup (..), stimesIdempotentMonoid)
import GHC.Exts (Int (..), Int#, TYPE, (==#))
import GHC.Stack (HasCallStack)
Expand All @@ -163,9 +164,9 @@ import Text.Read hiding (step)
import qualified Data.Data as Data
import qualified Data.Foldable as Foldable
import qualified Data.Functor.Classes as FC
import qualified Data.HashMap.Internal.Array as A
import qualified Data.Hashable as H
import qualified Data.Hashable.Lifted as H
import qualified Data.HashMap.Internal.Array as A
import qualified Data.List as List
import qualified GHC.Exts as Exts
import qualified Language.Haskell.TH.Syntax as TH
Expand Down Expand Up @@ -819,17 +820,9 @@ insertNewKey !h0 !k0 x0 !m0 = go h0 k0 x0 0 m0
in Full (update32 ary i st')
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = Collision h (snocNewLeaf (L k x) v)
| h == hy = Collision h (A.snoc v (L k x))
| otherwise =
go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
where
snocNewLeaf :: Leaf k v -> A.Array (Leaf k v) -> A.Array (Leaf k v)
snocNewLeaf leaf ary = A.run $ do
let n = A.length ary
mary <- A.new_ (n + 1)
A.copy ary 0 mary 0 n
A.write mary n leaf
return mary
{-# NOINLINE insertNewKey #-}


Expand Down Expand Up @@ -1008,12 +1001,8 @@ insertModifyingArr :: Eq k => v -> (v -> (# v #)) -> k -> A.Array (Leaf k v)
insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0)
where
go !k !ary !i !n
| i >= n = A.run $ do
-- Not found, append to the end.
mary <- A.new_ (n + 1)
A.copy ary 0 mary 0 n
A.write mary n (L k x)
return mary
-- Not found, append to the end.
| i >= n = A.snoc ary $ L k x
| otherwise = case A.index ary i of
(L kx y) | k == kx -> case f y of
(# y' #) -> if ptrEq y y'
Expand Down Expand Up @@ -1639,7 +1628,7 @@ unionArrayBy f !b1 !b2 !ary1 !ary2 = A.run $ do
A.write mary i =<< A.indexM ary2 i2
go (i+1) i1 (i2+1) b'
where
m = 1 `unsafeShiftL` (countTrailingZeros b)
m = 1 `unsafeShiftL` countTrailingZeros b
testBit x = x .&. m /= 0
b' = b .&. complement m
go 0 0 0 bCombined
Expand Down Expand Up @@ -1771,37 +1760,156 @@ differenceWith f a b = foldlWithKey' go empty a
-- | /O(n*log m)/ Intersection of two maps. Return elements of the first
-- map for keys existing in the second.
intersection :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v
intersection a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Just _ -> unsafeInsert k v m
_ -> m
intersection = Exts.inline intersectionWith const
{-# INLINABLE intersection #-}

-- | /O(n*log m)/ Intersection of two maps. If a key occurs in both maps
-- the provided function is used to combine the values from the two
-- maps.
intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1
-> HashMap k v2 -> HashMap k v3
intersectionWith f a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Just w -> unsafeInsert k (f v w) m
_ -> m
intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
intersectionWith f = Exts.inline intersectionWithKey $ const f
{-# INLINABLE intersectionWith #-}

-- | /O(n*log m)/ Intersection of two maps. If a key occurs in both maps
-- the provided function is used to combine the values from the two
-- maps.
intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3)
-> HashMap k v1 -> HashMap k v2 -> HashMap k v3
intersectionWithKey f a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Just w -> unsafeInsert k (f k v w) m
_ -> m
intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3) -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
intersectionWithKey f = intersectionWithKey# $ \k v1 v2 -> (# f k v1 v2 #)
{-# INLINABLE intersectionWithKey #-}

intersectionWithKey# :: Eq k => (k -> v1 -> v2 -> (# v3 #)) -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
intersectionWithKey# f = go 0
where
-- empty vs. anything
go !_ _ Empty = Empty
go _ Empty _ = Empty
-- leaf vs. anything
go s (Leaf h1 (L k1 v1)) t2 = lookupCont (\_ -> Empty) (\v _ -> case f k1 v1 v of (# v' #) -> Leaf h1 $ L k1 v') h1 k1 s t2
go s t1 (Leaf h2 (L k2 v2)) = lookupCont (\_ -> Empty) (\v _ -> case f k2 v v2 of (# v' #) -> Leaf h2 $ L k2 v') h2 k2 s t1
-- collision vs. collision
go _ (Collision h1 ls1) (Collision h2 ls2)
| h1 == h2 = runST $ do
(len, mary) <- intersectionCollisions f ls1 ls2
case len of
0 -> pure Empty
1 -> Leaf h1 <$> A.read mary 0
_ -> Collision h1 <$> (A.unsafeFreeze =<< A.shrink mary len)
| otherwise = Empty
-- branch vs. branch
go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) = intersectionArray s b1 b2 ary1 ary2
go s (BitmapIndexed b1 ary1) (Full ary2) = intersectionArray s b1 fullNodeMask ary1 ary2
go s (Full ary1) (BitmapIndexed b2 ary2) = intersectionArray s fullNodeMask b2 ary1 ary2
go s (Full ary1) (Full ary2) = intersectionArray s fullNodeMask fullNodeMask ary1 ary2
-- collision vs. branch
go s (BitmapIndexed b1 ary1) t2@(Collision h2 _ls2)
| b1 .&. m2 == 0 = Empty
| otherwise = go (s + bitsPerSubkey) (A.index ary1 i) t2
where
m2 = mask h2 s
i = sparseIndex b1 m2
go s t1@(Collision h1 _ls1) (BitmapIndexed b2 ary2)
| b2 .&. m1 == 0 = Empty
| otherwise = go (s + bitsPerSubkey) t1 (A.index ary2 i)
where
m1 = mask h1 s
i = sparseIndex b2 m1
go s (Full ary1) t2@(Collision h2 _ls2) = go (s + bitsPerSubkey) (A.index ary1 i) t2
where
i = index h2 s
go s t1@(Collision h1 _ls1) (Full ary2) = go (s + bitsPerSubkey) t1 (A.index ary2 i)
where
i = index h1 s

intersectionArray s b1 b2 ary1 ary2
-- don't create an array of size zero in intersectionArrayBy
| b1 .&. b2 == 0 = Empty
| otherwise = runST $ do
(b, len, ary) <- intersectionArrayBy (go (s + bitsPerSubkey)) b1 b2 ary1 ary2
case len of
0 -> pure Empty
1 -> A.read ary 0
_ -> bitmapIndexedOrFull b <$> (A.unsafeFreeze =<< A.shrink ary len)
{-# INLINE intersectionWithKey# #-}

intersectionArrayBy ::
( HashMap k v1 ->
HashMap k v2 ->
HashMap k v3
) ->
Bitmap ->
Bitmap ->
A.Array (HashMap k v1) ->
A.Array (HashMap k v2) ->
ST s (Bitmap, Int, A.MArray s (HashMap k v3))
intersectionArrayBy f !b1 !b2 !ary1 !ary2 = do
mary <- A.new_ $ popCount bIntersect
-- iterate over nonzero bits of b1 .&. b2
let go !i !i1 !i2 !b !bFinal
| b == 0 = pure (i, bFinal)
| testBit $ b1 .&. b2 = do
x1 <- A.indexM ary1 i1
x2 <- A.indexM ary2 i2
case f x1 x2 of
Empty -> go i (i1 + 1) (i2 + 1) b' (bFinal .&. complement m)
_ -> do
A.write mary i $! f x1 x2
go (i + 1) (i1 + 1) (i2 + 1) b' bFinal
| testBit b1 = go i (i1 + 1) i2 b' bFinal
| otherwise = go i i1 (i2 + 1) b' bFinal
where
m = 1 `unsafeShiftL` countTrailingZeros b
testBit x = x .&. m /= 0
b' = b .&. complement m
(maryLen, bFinal) <- go 0 0 0 bCombined bIntersect
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment at the top seems incorrect: Currently the loop actually iterates over b1 .|. b2. It would be nice to change this though. In that case the i1 and i2 indices could be computed with sparseIndex.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case the i1 and i2 indices could be computed with sparseIndex.

For comparison, here is a version of unionArrayBy that uses sparseIndex to compute all the indices:

unordered-containers/Data/HashMap/Internal.hs

Lines 1626 to 1647 in a780a8d

unionArrayBy f !b1 !b2 !ary1 !ary2 = A.run $ do
let b' = b1 .|. b2
mary <- A.new_ (popCount b')
-- iterate over nonzero bits of b1 .|. b2
let go !b
| b == 0 = return ()
| otherwise = do
let ba = b1 .&. b2
c = countTrailingZeros b
m = bit c
i = sparseIndex b' m
i1 = sparseIndex b1 m
i2 = sparseIndex b2 m
t <- if | testBit ba c -> do
x1 <- A.indexM ary1 i1
x2 <- A.indexM ary2 i2
return $! f x1 x2
| testBit b1 c -> A.indexM ary1 i1
| otherwise -> A.indexM ary2 i2
A.write mary i t
go (clearBit b c)
go b'

I expect that keeping i as a loop argument will be more efficient than recomputing it on each iteration though.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using sparseIndex makes benchmarks slower

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you show me the diff?

Using sparseIndex makes benchmarks slower

By how much? I also think the benchmark data might be a bit weird.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How do I show you the diff?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally both.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sparse index:

All
  HashMap
    intersection
      Int:        OK (0.34s)
        52.1 μs ± 3.3 μs
      ByteString: OK (0.26s)
        62.6 μs ± 6.1 μs

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Without sparse index:

All
  HashMap
    intersection
      Int:        OK (0.86s)
        42.1 μs ± 1.5 μs
      ByteString: OK (0.33s)
        47.7 μs ± 2.8 μs

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alright, thanks! I think we might get different results with data where there's less overlap between the two maps. But that can be investigated at a different time.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up issue in #416.

pure (bFinal, maryLen, mary)
where
bCombined = b1 .|. b2
bIntersect = b1 .&. b2
{-# INLINE intersectionArrayBy #-}

intersectionCollisions :: Eq k => (k -> v1 -> v2 -> (# v3 #)) -> A.Array (Leaf k v1) -> A.Array (Leaf k v2) -> ST s (Int, A.MArray s (Leaf k v3))
intersectionCollisions f ary1 ary2 = do
mary2 <- A.thaw ary2 0 $ A.length ary2
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder whether we actually need to allocate two arrays for this. The alternative would be to perform the search-and-swap operations on the output array itself.

It might be a bit tricky though – maybe leave it for a follow-up PR, so this one doesn't get too huge.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the issue with this is that the type could change. For example if we have two arrays with the numbers as keys, and the arrays are both different types
1 2 3 4
3 4 2 1
Let's thaw the first array, and mutate it to
(f 3 3) 2 1 4
f 3 3 could change the type to be something difference than the 2 1 4.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, yes, good point. Unsafe coercions might work for this, but I'd prefer not trying this in this PR.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is only an issue for intersectionWithKey and such; intersection itself has no type issue.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's another thing about intersection that's special: we can reuse the leaves.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it would be better if intersection had custom code for handling collisions. Maybe this can be achieved by changing intersectionWithKey# to something similar to filterMapAux.

I'd slightly prefer if we'd leave this for a follow-up PR though.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is the branch https://github.com/oberblastmeister/unordered-containers/tree/fast-intersection-sparseIndex

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have recorded these ideas in #415.

mary <- A.new_ $ min (A.length ary1) (A.length ary2)
let go i j
| i >= A.length ary1 || j >= A.lengthM mary2 = pure j
| otherwise = do
L k1 v1 <- A.indexM ary1 i
searchSwap k1 j mary2 >>= \case
Just (L _k2 v2) -> do
let !(# v3 #) = f k1 v1 v2
A.write mary j $ L k1 v3
go (i + 1) (j + 1)
Nothing -> do
go (i + 1) j
maryLen <- go 0 0
pure (maryLen, mary)
{-# INLINE intersectionCollisions #-}

-- | Say we have
-- @
-- 1 2 3 4
-- @
-- and we search for @3@. Then we can mutate the array to
-- @
-- undefined 2 1 4
-- @
-- We don't actually need to write undefined, we just have to make sure that the next search starts 1 after the current one.
searchSwap :: Eq k => k -> Int -> A.MArray s (Leaf k v) -> ST s (Maybe (Leaf k v))
searchSwap toFind start = go start toFind start
where
go i0 k i mary
| i >= A.lengthM mary = pure Nothing
| otherwise = do
l@(L k' _v) <- A.read mary i
if k == k'
then do
A.write mary i =<< A.read mary i0
pure $ Just l
else go i0 k (i + 1) mary
{-# INLINE searchSwap #-}


------------------------------------------------------------------------
-- * Folds

Expand Down Expand Up @@ -2164,12 +2272,8 @@ updateOrSnocWithKey :: Eq k => (k -> v -> v -> (# v #)) -> k -> v -> A.Array (Le
updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
where
go !k v !ary !i !n
| i >= n = A.run $ do
-- Not found, append to the end.
mary <- A.new_ (n + 1)
A.copy ary 0 mary 0 n
A.write mary n (L k v)
return mary
-- Not found, append to the end.
| i >= n = A.snoc ary $ L k v
| L kx y <- A.index ary i
, k == kx
, (# v2 #) <- f k v y
Expand Down
26 changes: 26 additions & 0 deletions Data/HashMap/Internal/Array.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ module Data.HashMap.Internal.Array
, new_
, singleton
, singletonM
, snoc
, pair

-- * Basic interface
Expand Down Expand Up @@ -76,6 +77,7 @@ module Data.HashMap.Internal.Array
, toList
, fromList
, fromList'
, shrink
) where

import Control.Applicative (liftA2)
Expand All @@ -91,6 +93,7 @@ import GHC.Exts (Int (..), SmallArray#, SmallMutableArray#,
thawSmallArray#, unsafeCoerce#,
unsafeFreezeSmallArray#, unsafeThawSmallArray#,
writeSmallArray#)
import qualified GHC.Exts as Exts
import GHC.ST (ST (..))
import Prelude hiding (all, filter, foldMap, foldl, foldr, length,
map, read, traverse)
Expand Down Expand Up @@ -204,6 +207,20 @@ new _n@(I# n#) b =
new_ :: Int -> ST s (MArray s a)
new_ n = new n undefinedElem

-- | When 'Exts.shrinkSmallMutableArray#' is available, the returned array is the same as the array given, as it is shrunk in place.
-- Otherwise a copy is made.
shrink :: MArray s a -> Int -> ST s (MArray s a)
#if MIN_VERSION_GLASGOW_HASKELL(8, 10, 7, 0)
shrink mary _n@(I# n#) =
CHECK_GT("shrink", _n, (0 :: Int))
CHECK_LE("shrink", _n, (lengthM mary))
ST $ \s -> case Exts.shrinkSmallMutableArray# (unMArray mary) n# s of
s' -> (# s', mary #)
#else
shrink mary n = cloneM mary 0 n
#endif
{-# INLINE shrink #-}

singleton :: a -> Array a
singleton x = runST (singletonM x)
{-# INLINE singleton #-}
Expand All @@ -212,6 +229,15 @@ singletonM :: a -> ST s (Array a)
singletonM x = new 1 x >>= unsafeFreeze
{-# INLINE singletonM #-}

snoc :: Array a -> a -> Array a
snoc ary x = run $ do
mary <- new (n + 1) x
copy ary 0 mary 0 n
pure mary
where
n = length ary
{-# INLINE snoc #-}

pair :: a -> a -> Array a
pair x y = run $ do
ary <- new 2 x
Expand Down
13 changes: 3 additions & 10 deletions Data/HashMap/Internal/Strict.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ import Prelude hiding (lookup, map)
import qualified Data.HashMap.Internal as HM
import qualified Data.HashMap.Internal.Array as A
import qualified Data.List as List
import GHC.Exts (inline)

{-
Note [Imports from Data.HashMap.Internal]
Expand Down Expand Up @@ -616,23 +617,15 @@ differenceWith f a b = HM.foldlWithKey' go HM.empty a
-- maps.
intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1
-> HashMap k v2 -> HashMap k v3
intersectionWith f a b = HM.foldlWithKey' go HM.empty a
where
go m k v = case HM.lookup k b of
Just w -> let !x = f v w in HM.unsafeInsert k x m
_ -> m
intersectionWith f = inline intersectionWithKey $ const f
{-# INLINABLE intersectionWith #-}

-- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps
-- the provided function is used to combine the values from the two
-- maps.
intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3)
-> HashMap k v1 -> HashMap k v2 -> HashMap k v3
intersectionWithKey f a b = HM.foldlWithKey' go HM.empty a
where
go m k v = case HM.lookup k b of
Just w -> let !x = f k v w in HM.unsafeInsert k x m
_ -> m
intersectionWithKey f = HM.intersectionWithKey# $ \k v1 v2 -> let !v3 = f k v1 v2 in (# v3 #)
{-# INLINABLE intersectionWithKey #-}

------------------------------------------------------------------------
Expand Down
6 changes: 5 additions & 1 deletion benchmarks/Benchmarks.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,13 +318,17 @@ main = do
[ bench "Int" $ whnf (HM.union hmi) hmi2
, bench "ByteString" $ whnf (HM.union hmbs) hmbsSubset
]

, bgroup "intersection"
[ bench "Int" $ whnf (HM.intersection hmi) hmi2
, bench "ByteString" $ whnf (HM.intersection hmbs) hmbsSubset
]

-- Transformations
, bench "map" $ whnf (HM.map (\ v -> v + 1)) hmi

-- * Difference and intersection
, bench "difference" $ whnf (HM.difference hmi) hmi2
, bench "intersection" $ whnf (HM.intersection hmi) hmi2

-- Folds
, bench "foldl'" $ whnf (HM.foldl' (+) 0) hmi
Expand Down
Loading