-
Notifications
You must be signed in to change notification settings - Fork 102
Open
Labels
Description
unordered-containers/Data/HashMap/Internal.hs
Lines 2181 to 2210 in b6bde46
updateOrConcatWithKey :: Eq k => (k -> v -> v -> (# v #)) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) | |
updateOrConcatWithKey f ary1 ary2 = A.run $ do | |
-- TODO: instead of mapping and then folding, should we traverse? | |
-- We'll have to be careful to avoid allocating pairs or similar. | |
-- first: look up the position of each element of ary2 in ary1 | |
let indices = A.map' (\(L k _) -> indexOf k ary1) ary2 | |
-- that tells us how large the overlap is: | |
-- count number of Nothing constructors | |
let nOnly2 = A.foldl' (\n -> maybe (n+1) (const n)) 0 indices | |
let n1 = A.length ary1 | |
let n2 = A.length ary2 | |
-- copy over all elements from ary1 | |
mary <- A.new_ (n1 + nOnly2) | |
A.copy ary1 0 mary 0 n1 | |
-- append or update all elements from ary2 | |
let go !iEnd !i2 | |
| i2 >= n2 = return () | |
| otherwise = case A.index indices i2 of | |
Just i1 -> do -- key occurs in both arrays, store combination in position i1 | |
L k v1 <- A.indexM ary1 i1 | |
L _ v2 <- A.indexM ary2 i2 | |
case f k v1 v2 of (# v3 #) -> A.write mary i1 (L k v3) | |
go iEnd (i2+1) | |
Nothing -> do -- key is only in ary2, append to end | |
A.write mary iEnd =<< A.indexM ary2 i2 | |
go (iEnd+1) (i2+1) | |
go n1 0 | |
return mary | |
{-# INLINABLE updateOrConcatWithKey #-} |
I think that we could avoid allocating the intermediary indices
array. Instead over-allocate the output array and shrink (#362) it once we've determined the final size.
When matching the keys we could also use a similar optimization as suggested in #291.