|
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.
