haskell / containers

Assorted concrete container types
https://hackage.haskell.org/package/containers
315 stars 178 forks source link

Add upsert #809

Open phadej opened 2 years ago

phadej commented 2 years ago

There are

adjust :: Ord k => (      a ->       a) -> k -> Map k a -> Map k a
update :: Ord k => (      a -> Maybe a) -> k -> Map k a -> Map k a
alter  :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a

but

upsert :: Ord k => (Maybe a ->       a) -> k -> Map k a -> Map k a

variant is missing.

While it's definable with upsert f = alter (Just . f), I'd like to have it already in the library.

Compare implementation for Map: even alter is tagged as INLINEABLE, the optimization may or may not happen.

alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
alter = go
  where
    go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
    go f !k Tip = case f Nothing of
               Nothing -> Tip
               Just x  -> singleton k x

    go f k (Bin sx kx x l r) = case compare k kx of
               LT -> balance kx x (go f k l) r
               GT -> balance kx x l (go f k r)
               EQ -> case f (Just x) of
                       Just x' -> Bin sx kx x' l r
                       Nothing -> glue l r

vs

upsert :: Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a
upsert = go
  where
    go :: Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a
    go f !k Tip = singleton k (f Nothing)
    go f k (Bin sx kx x l r) = case compare k kx of
               LT -> balance kx x (go f k l) r
               GT -> balance kx x l (go f k r)
               EQ -> Bin sx kx (f (Just x)) l r

In fact, I'd like having upsertF variant as well. It would be less code/branches then with alterF.

treeowl commented 2 years ago

It sounds reasonable.

Bodigrim commented 7 months ago

+1, I've just spent some time trying to locate Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a, only to realise that it's missing indeed.