Make `unionWith` more flexible

[garyb]

As per @fehrenbach in https://github.com/purescript/purescript-ordered-collections/pull/5#discussion_r219065421:

unionWith :: forall k a b c. Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c

[fehrenbach]

unionWith cannot have that type, because values from both input maps may end up in the result (without passing through the combining function). Not sure what I was thinking. (I guess a merge :: forall k u v w. Ord k => (Maybe u -> Maybe v -> Maybe w) -> Map k u -> Map k v -> Map k w could be neat. Also complicated and probably slow.)

[garyb]

I didn't spot it either :smile:. Might be worth having as zipWith though.

[fehrenbach]

or intersect(ion)With. I like zipWith better, though.

And this would be a better type for merge, I guess. Now can I stop thinking about this, please?!

merge :: forall k u v w. Ord k => (k -> These u v -> Maybe w) -> Map k u -> Map k v -> Map k w

[garyb]

Just remembered this discussion too: https://github.com/purescript-contrib/purescript-these/issues/21

[karshan]

Slightly related: I just made pr #15 to add intersection and intersectionWith.

[np]

For the record here is advanced API used in Haskell: Data.Map.Merge.Strict.

Would such an API be considered if implemented?

[np]

Let's forget about my previous comment.

I'm focusing on the type signature from @fehrenbach. Here is what I think could be a reference :

merge :: forall k a b c. Ord k => (k -> These a b -> Maybe c) -> Map k a -> Map k b -> Map k c
merge f m1 m2 = mapMaybeWithKey f (unionWith g (This <$> m1) (That <$> m2))
  where
    g (This p) (That v) = Both p v
    g x _ = x -- impossible

What do you think? Shall we integrate something functionally equivalent but which avoid these 3 extra traversals.

[xgrommx]

How about?

mergeMapWithKey :: forall k a b c. Ord k => (k -> These a b -> c -> c) -> c -> M.Map k a -> M.Map k b -> c
mergeMapWithKey fn init = go
  where
    step :: k -> Tuple (List (Tuple k a)) c -> b -> Tuple (List (Tuple k a)) c
    step rKey (Tuple l res) rValue =
      case l of
        Nil -> Tuple l (fn rKey (That rValue) res)
        Cons (Tuple lKey lValue) rest -> case compare lKey rKey of
          LT -> step rKey (Tuple rest (fn lKey (This lValue) res)) rValue
          GT -> Tuple l (fn rKey (That rValue) res)
          EQ -> Tuple rest (fn lKey (Both lValue rValue) res)
    go :: M.Map k a -> M.Map k b -> c
    go m1 = uncurry (flip (foldl (\result (Tuple k v) -> fn k (This v) result))) <<< foldlWithIndex step (Tuple (M.toUnfoldable m1) init)

type DeltaUnit a = { old :: a, new :: a }

data Delta a
  = Delta (DeltaUnit a)
  | Same a
  | Old a
  | New a

derive instance genericDelta :: Generic (Delta a) _

instance showDelta :: Show a => Show (Delta a) where
  show d = genericShow d

instance eqDelta :: Eq a => Eq (Delta a) where
  eq d1 d2 = genericEq d1 d2

diff :: forall k v. Eq v => Ord k => M.Map k v -> M.Map k v -> M.Map k (Delta v)  
diff = mergeMapWithKey fn M.empty
  where
    fn k t = M.insert k (go t)
      where
        go = case _ of
          This x -> Old x
          That x -> New x
          Both x y
            | x == y -> Same x
            | otherwise -> Delta { old: x, new: y }

unionWithKey :: forall k c. Ord k => (k -> c -> c -> c) -> M.Map k c -> M.Map k c -> M.Map k c
unionWithKey f = mergeMapWithKey fn M.empty
  where
    fn k t = M.insert k (go t)
      where
        go = case _ of
          This x -> x
          That x -> x
          Both x y -> f k x y

intersectionWithKey :: forall k a b c. Ord k => (k -> a -> b -> c) -> M.Map k a -> M.Map k b -> M.Map k c
intersectionWithKey f = mergeMapWithKey fn M.empty
  where
    fn k t m = case t of
      This x -> m
      That x -> m
      Both x y -> M.insert k (f k x y) m

differenceWith :: forall k a b. Ord k => (k -> a -> b -> a) -> M.Map k a -> M.Map k b -> M.Map k a
differenceWith f = mergeMapWithKey fn M.empty
  where
    fn k t m = case t of
      This x -> M.insert k x m
      That x -> m
      Both x y -> M.insert k (f k x y) m

symmetricDifference :: forall k a. Ord k => M.Map k a -> M.Map k a -> M.Map k a
symmetricDifference = mergeMapWithKey fn M.empty
  where
    fn k t m = case t of
      This x -> M.insert k x m
      That x -> M.insert k x m
      Both _ _ -> m 

[JordanMartinez]

Seems like this would work to me.

[xgrommx]

@JordanMartinez but probably we should eliminate toUnfoldable and try to fuse foldlWithIndex and foldl if it possible

[JordanMartinez]

@xgrommx Let's say that could be fused, would the type signature change of any of these functions? If not, then couldn't that be done in a future PR?

[xgrommx]

@JordanMartinez what do u mean about signature?

[JordanMartinez]

@xgrommx Meaning, the type signature for mergeMapWithKey is currently below. If fusion was possible, would that implementation detail change the type signature of mergeMapWithKey's (which acts as a public API)?

mergeMapWithKey :: forall k a b c. Ord k => (k -> These a b -> c -> c) -> c -> M.Map k a -> M.Map k b -> c
mergeMapWithKey fn init = go
  where
    step :: k -> Tuple (List (Tuple k a)) c -> b -> Tuple (List (Tuple k a)) c
    step rKey (Tuple l res) rValue =
      case l of
        Nil -> Tuple l (fn rKey (That rValue) res)
        Cons (Tuple lKey lValue) rest -> case compare lKey rKey of
          LT -> step rKey (Tuple rest (fn lKey (This lValue) res)) rValue
          GT -> Tuple l (fn rKey (That rValue) res)
          EQ -> Tuple rest (fn lKey (Both lValue rValue) res)
    go :: M.Map k a -> M.Map k b -> c
    go m1 = uncurry (flip (foldl (\result (Tuple k v) -> fn k (This v) result))) <<< foldlWithIndex step (Tuple (M.toUnfoldable m1) init)