Zipping free monad transformers

Question

The streaming package offers a zipsWith function

zipsWith
  :: (Monad m, Functor h)
  => (forall x y. f x -> g y -> h (x, y))
  -> Stream f m r -> Stream g m r -> Stream h m r

and a slightly more streamlined version,

zipsWith'
  :: Monad m
  => (forall x y p. (x -> y -> p) -> f x -> g y -> h p)
  -> Stream f m r -> Stream g m r -> Stream h m r

These can be adapted very easily to FreeT from the free package. But that package offers another version of the free monad transformer:

newtype FT f m a = FT
  { runFT
      :: forall r.
         (a -> m r)
      -> (forall x. (x -> m r) -> f x -> m r)
      -> m r }

There is also a third (rather simple) formulation:

newtype FF f m a = FF
  { runFF
      :: forall n. Monad n
      => (forall x. f x -> n x)  -- A natural transformation
      -> (forall x. m x -> n x)  -- A monad morphism
      -> n a }

It is possible to convert back and forth between FreeT and either FT or FF, which offers an indirect way to implement zipsWith and its relatives for FF and FT. But that seems quite unsatisfying. I seek a more direct solution.

The problem seems related to the challenge of zipping lists using folds. This has been addressed in a paper, Coroutining Folds with Hyperfunctions, by Launchbury et al, as well as a blog post by Donnacha Kidney. Neither of these are terribly simple, and I have no idea how they might be adapted to the FT or FF contexts.

---

As I've looked into this problem, I've realized that streaming should really offer some more powerful versions. The simplest would be something like

zipsWith''
  :: Monad m
  => (forall x y p. (x -> y -> p) -> f x -> g y -> h p)
  -> Stream f m r -> Stream g m s -> Stream h m (Either r s)

but a more powerful option would include the remainder:

zipsWithRemains
  :: Monad m
  => (forall x y p. (x -> y -> p) -> f x -> g y -> h p)
  -> Stream f m r
  -> Stream g m s
  -> Stream h m (Either (r, Stream g m s)
                        (f (Stream f m r), s))

I would guess that zipsWith'' would be no harder than zipsWith', but that zipsWithRemains might be a bigger challenge in the context of FT or FF, since the remainder will presumably have to be reconstituted somehow.

Note

Since there was some confusion previously, let me mention that I am not looking for help writing zipsWithRemains for Stream or FreeT; I am only looking for help with the functions on FT and FF.

Answer 1

I implemented zipsWith', zipsWith'' and zipsWithRemains for FT. My implementation closely mirrors the implementation of zipWith from this blog post.

First, notice that, given zipsWith', implementing zipsWith'' is trivial:

zipsWith''
  :: (Functor f, Functor g, Monad m)
  => (forall x y p. (x -> y -> p) -> f x -> g y -> h p)
  -> FT f m r
  -> FT g m s
  -> FT h m (Either r s)
zipsWith'' phi a b = zipsWith' phi (Left <$> a) (Right <$> b)

So let's implement zipsWith'.

Begin with an expanded and annotated version of zipWith using folds:

newtype RecFold a r = RecFold { runRecFold :: BFold a r }
type AFold a r = RecFold a r -> r
type BFold a r = a -> AFold a r -> r

zipWith
  :: forall f g a b c.
  (Foldable f, Foldable g)
  => (a -> b -> c)
  -> f a
  -> g b
  -> [c]
zipWith c a b = loop af bf where
  af :: AFold a [c]
  af = foldr ac ai a
  ai :: AFold a [c]
  ai _ = []
  ac :: a -> AFold a [c] -> AFold a [c]
  ac ae ar bl = runRecFold bl ae ar
  bf :: BFold a [c]
  bf = foldr bc bi b
  bi :: BFold a [c]
  bi _ _ = []
  bc :: b -> BFold a [c] -> BFold a [c]
  bc be br ae ar = c ae be : loop ar br
  loop :: AFold a [c] -> BFold a [c] -> [c]
  loop al bl = al (RecFold bl)

And turn it into zipsWith':

newtype RecFold f m r = RecFold { runRecFold :: BFold f m r }
type AFold f m r = m (RecFold f m r -> r)
type BFold f m r = m (f (AFold f m r) -> r)

zipsWith'
  :: forall f g h m r.
  (Monad m, Functor f, Functor g)
  => (forall x y p. (x -> y -> p) -> f x -> g y -> h p)
  -> FT f m r
  -> FT g m r
  -> FT h m r
zipsWith' phi a b = loop af bf where
  af :: AFold f m (FT h m r)
  af = runFT a ai ac
  ai :: r -> AFold f m (FT h m r)
  ai r = return $ const $ return r
  ac :: (x -> AFold f m (FT h m r)) -> f x -> AFold f m (FT h m r)
  ac am ae = return $ effect . fmap ($ (fmap am ae)) . runRecFold
  bf :: BFold f m (FT h m r)
  bf = runFT b bi bc
  bi :: r -> BFold f m (FT h m r)
  bi r = return $ const $ return r
  bc :: (x -> BFold f m (FT h m r)) -> g x -> BFold f m (FT h m r)
  bc bm be = return $ wrap . flip (phi loop) (fmap bm be)
  loop :: AFold f m (FT h m r) -> BFold f m (FT h m r) -> FT h m r
  loop av bv = effect $ fmap ($ (RecFold bv)) av

Here, two auxiliary functions are used: effect and wrap.

effect :: Monad m => m (FT f m r) -> FT f m r
effect m = FT $ \hr hy -> m >>= \r -> runFT r hr hy

wrap :: f (FT f m r) -> FT f m r
wrap s = FT $ \hr hy -> hy (\v -> runFT v hr hy) s

Note that the result could be any monad for which these functions are implemented.

To implement zipsWithRemains, start by implementing zipWithRemains for ordinary Foldables:

data ListWithTail a b = Nil b | Cons a (ListWithTail a b)
type Result a b c = ListWithTail c (Either [b] (a, [a]))
newtype RecFold a b c = RecFold { runRecFold :: BFold a b c }
type AFold a b c = (RecFold a b c -> Result a b c, [a])
type BFold a b c = (a -> AFold a b c -> Result a b c, [b])

zipWithRemains
  :: forall f g a b c.
  (Foldable f, Foldable g)
  => (a -> b -> c)
  -> f a
  -> g b
  -> Result a b c
zipWithRemains c a b = loop af bf where
  af :: AFold a b c
  af = foldr ac ai a
  ai :: AFold a b c
  ai = (\bl -> Nil $ Left $ snd (runRecFold bl), [])
  ac :: a -> AFold a b c -> AFold a b c
  ac ae ar = (\bl -> fst (runRecFold bl) ae ar, ae : snd ar)
  bf :: BFold a b c
  bf = foldr bc bi b
  bi :: BFold a b c
  bi = (\ae ar -> Nil $ Right (ae, snd ar), [])
  bc :: b -> BFold a b c -> BFold a b c
  bc be br = (\ae ar -> Cons (c ae be) (loop ar br), be : snd br)
  loop :: AFold a b c -> BFold a b c -> Result a b c
  loop al bl = fst al (RecFold bl)

Here, the result of a fold is not a function but a 2-tuple containing a function and a value. The latter is used to handle the "remains" case.

This can also be adapted to FT:

type Result f g h m r s = FT h m (Either (r, FT g m s) (f (FT f m r), s))
newtype RecFold f g h m r s = RecFold { runRecFold :: BFold f g h m r s }
type AFold f g h m r s = m (RecFold f g h m r s -> Result f g h m r s, FT f m r)
type BFold f g h m r s = m (f (AFold f g h m r s) -> Result f g h m r s, FT g m s)

zipsWithRemains
  :: forall f g h m r s.
  (Monad m, Functor f, Functor g)
  => (forall x y p. (x -> y -> p) -> f x -> g y -> h p)
  -> FT f m r
  -> FT g m s
  -> Result f g h m r s
zipsWithRemains phi a b = loop af bf where
  af :: AFold f g h m r s
  af = runFT a ai ac
  ai :: r -> AFold f g h m r s
  ai r = return (return . Left . (r,) . effect . fmap snd . runRecFold, return r)
  ac :: (x -> AFold f g h m r s) -> f x -> AFold f g h m r s
  ac am ae = return (effect . fmap (($ (fmap am ae)) . fst) . runRecFold, wrap $ fmap (effect . fmap snd . am) ae)
  bf :: BFold f g h m r s
  bf = runFT b bi bc
  bi :: s -> BFold f g h m r s
  bi r = return (return . Right . (,r) . fmap (effect . fmap snd), return r)
  bc :: (x -> BFold f g h m r s) -> g x -> BFold f g h m r s
  bc bm be = return (wrap . flip (phi loop) (fmap bm be), wrap $ fmap (effect . fmap snd . bm) be)
  loop :: AFold f g h m r s -> BFold f g h m r s -> Result f g h m r s
  loop av bv = effect $ fmap (($ (RecFold bv)) . fst) av

I wish Haskell had local types!

This probably answers the question for FT. Regarding FF: this type is designed such that to do anything with it, you first have to convert it to some other monad. So, the question is, which one? It is possible to convert it to Stream or FreeT, and use the functions for those types. It is also possible to convert it to FT and use the above implementations on it. Is there a monad better suited for implementing zipsWith? Maybe.

Answer 2

Applying a bit of Coyoneda to abacabadabacaba's answer and doing some juggling yields an implementation that avoids Functor f and Functor g constraints. If those functors have expensive fmaps, this may improve performance. I doubt it's actually better in typical situations where f and g are things like (,) a. I also still don't properly understand what any of this does.

type AFold f m r = m (RecFold f m r -> r)
newtype Fish f m r = Fish {unFish :: forall x. (x -> AFold f m r) -> f x -> r}
type BFold f m r = m (Fish f m r)
newtype RecFold f m r = RecFold { runRecFold :: BFold f m r }

zipsWith'
  :: forall f g h m r.
  Monad m
  => (forall x y p. (x -> y -> p) -> f x -> g y -> h p)
  -> FT f m r
  -> FT g m r
  -> FT h m r
zipsWith' phi a b = loop af bf where
  af :: AFold f m (FT h m r)
  af = runFT a ai ac

  ai :: r -> AFold f m (FT h m r)
  ai r = return $ const $ return r

  ac :: (x -> AFold f m (FT h m r)) -> f x -> AFold f m (FT h m r)
  ac am ae = return $ (lift >=> \(Fish z) -> z am ae) . runRecFold

  bf :: BFold f m (FT h m r)
  bf = runFT b bi bc

  bi :: r -> BFold f m (FT h m r)
  bi r = return $ Fish $ \_ _ -> return r

  bc :: (x -> BFold f m (FT h m r)) -> g x -> BFold f m (FT h m r)
  bc bm be = return $ Fish $ \xa z -> wrap $ phi (\q -> loop (xa q) . bm) z be

  loop :: AFold f m (FT h m r) -> BFold f m (FT h m r) -> FT h m r
  loop av bv = lift av >>= ($ (RecFold bv))