-
Notifications
You must be signed in to change notification settings - Fork 102
Open
Labels
Description
I'm not sure whether this leads to any unwanted laziness or reboxing, but I'd feel more comfortable if some bang patterns would simply rule out such issues.
unordered-containers/Data/HashMap/Internal.hs
Lines 1888 to 1933 in 19674b5
intersectionCollisions :: Eq k => (k -> v1 -> v2 -> (# v3 #)) -> Hash -> Hash -> A.Array (Leaf k v1) -> A.Array (Leaf k v2) -> HashMap k v3 | |
intersectionCollisions f h1 h2 ary1 ary2 | |
| h1 == h2 = runST $ do | |
mary2 <- A.thaw ary2 0 $ A.length ary2 | |
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 | |
len <- go 0 0 | |
case len of | |
0 -> pure Empty | |
1 -> Leaf h1 <$> A.read mary 0 | |
_ -> Collision h1 <$> (A.unsafeFreeze =<< A.shrink mary len) | |
| otherwise = Empty | |
{-# 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 #-} |
The searchSwap
variant being added in #435 should get similar bang patterns.