Can we make the Map.Merge API more expressive?

[meooow25]

I wanted to demonstrate partitionKeys recently (https://github.com/haskell/containers/pull/975#issuecomment-2417976839) and realized that the public Map.Merge API is not expressive enough for it.

What I need:

wm1 :: WhenMissing Pair k a a
wm1 = WhenMissing (\t -> Pair empty t) (\_ x -> Pair Nothing (Just x))

Best I can do with the public API:

wm1 :: WhenMissing Pair k a a
wm1 = traverseMaybeMissing (\_ x -> Pair Nothing (Just x))

which is terribly inefficient! (O(1) vs O(n))

Is there a safe way to allow such use cases?

[treeowl]

Interesting. My immediate intuition is that there might be a nice way to do this with a Biapplicative analogue of the merge API. I'll try to think later about whether there's a way to do it with just Applicative, but I'm not super optimistic.

[meooow25]

Going by the types we need something like

lift2
  :: (forall a b. m1 a -> m2 a -> n a)
  -> WhenMissing m1 k a b
  -> WhenMissing m2 k a b
  -> WhenMissing n k a b
lift2 f (WhenMissing f1 g1) (WhenMissing f2 g2) =
  WhenMissing
    (\t -> f (f1 t) (f2 t))
    (\k x -> f (g1 k x) (g2 k x))

where f probably needs some laws attached. Then

wm1 :: WhenMissing Pair k a a
wm1 = lift2 (\(Identity x1) (Identity x2) -> Pair x1 x2) M.dropMissing M.preserveMissing

This seems like https://hackage.haskell.org/package/mmorph territory.

Or perhaps equivalently

import qualified Data.Functor.Product as Prod

hoistMissing :: (forall a. f a -> g a) -> WhenMissing f k a b -> WhenMissing g k a b
hoistMissing f (WhenMissing f1 g1) = WhenMissing (\t -> f (f1 t)) (\k x -> f (g1 k x))

pairMissing
  :: WhenMissing m1 k a b
  -> WhenMissing m2 k a b
  -> WhenMissing (Prod.Product m1 m2) k a b
pairMissing (WhenMissing f1 g1) (WhenMissing f2 g2) =
  WhenMissing
    (\t -> Prod.Pair (f1 t) (f2 t))
    (\k x -> Prod.Pair (g1 k x) (g2 k x))

wm1 :: WhenMissing Pair k a a
wm1 =
  hoistMissing
    (\(Prod.Pair (Identity x1) (Identity x2)) -> Pair x1 x2)
    (pairMissing M.dropMissing M.preserveMissing)

[meooow25]

Alternately we could expose the constructor with a clear warning:

WARNING: A value WhenMissing f g must satisfy the law f = traverseMaybeWithKey g.

Yet another option is to be able to safely expose the constructor, as in #937, but that has it's own issues.