Skip to content

updateOrConcatWithKey could be more efficient #403

Open
@sjakobi

Description

@sjakobi

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions