Proposal: add generic helpers for MonadZip

[David Feuer]

While we can't *zip* sums generically, we can *unzip* them. Here's a rewrite:

genericmzipWith :: (Generic1 m, GMonadZip (Rep1 m)) => (a -> b -> c)
-> m a -> m b -> m c
genericmzipWith f m1 m2 = to1 $ gmzipWith f (from1 m1) (from1 m2)

genericmunzip :: (Generic1 m, GMonadUnzip (Rep1 m)) => m (a, b) -> (m a, m b)
genericmunzip m = case gmunzip (from1 m) of
  (p, q) -> (to1 p, to1 q)

class GMonadZip f where
  gmzipWith :: (a -> b -> c) -> f a -> f b -> f c

class GMonadUnzip f where
  gmunzip :: f (a, b) -> (f a, f b)

instance MonadZip f => GMonadZip (Rec1 f) where
    gmzipWith f (Rec1 fa) (Rec1 fb) = Rec1 (mzipWith f fa fb)

instance MonadZip f => GMonadUnzip (Rec1 f) where
    gmunzip (Rec1 fab) = (Rec1 fa, Rec1 fb)
      where
        -- We want to be lazy here, because this might actually
        -- be recursive. We do a lot of NOINLINE to get selector
        -- thunks to avoid space leaks.
        {-# NOINLINE mufab #-}
        {-# NOINLINE fa #-}
        {-# NOINLINE fb #-}
        mufab = munzip fab
        (fa, fb) = mufab

nstance (MonadZip f, GMonadZip g) => GMonadZip (f :.: g) where
  gmzipWith f (Comp1 fga1) (Comp1 fga2) = Comp1 $ mzipWith (gmzipWith
f) fga1 fga2

instance (Functor g, MonadZip f, GMonadUnzip g) => GMonadUnzip (f :.: g) where
  gmunzip (Comp1 fgc) = case munzip . fmap gmunzip $ fgc of
    (p, q) -> (Comp1 p, Comp1 q)

instance GMonadZip U1 where
    gmzipWith _ _ _ = U1
instance GMonadUnzip U1 where
    gmunzip _ = (U1, U1)

instance GMonadZip Par1 where
    gmzipWith = coerce
instance GMonadUnzip Par1 where
    gmunzip = coerce

deriving instance GMonadZip f => GMonadZip (M1 i c f)
deriving instance GMonadUnzip f => GMonadUnzip (M1 i c f)

instance (GMonadZip f, GMonadZip g) => GMonadZip (f :*: g) where
    gmzipWith f (x1 :*: y1) (x2 :*: y2) = gmzipWith f x1 x2 :*:
gmzipWith f y1 y2

instance (GMonadUnzip f, GMonadUnzip g) => GMonadUnzip (f :*: g) where
    -- Why don't we need to be lazy in this munzip? If we're working with
    -- Rep1, then laziness will be added by the Rec1 instance. If we're working
    -- with Rep, then we can't have any K1s because K1 isn't an instance of
    -- Monad, let alone MonadZip!
    gmunzip (fab :*: gab)
      | (fa, fb) <- gmunzip fab
      , (ga, gb) <- gmunzip gab
      = (fa :*: ga, fb :*: gb)

instance (GMonadUnzip f, GMonadUnzip g) => GMonadUnzip (f :+: g) where
  gmunzip (L1 x) = case gmunzip x of
    (l, r) -> (L1 l, L1 r)
  gmunzip (R1 x) = case gmunzip x of
    (l, r) -> (R1 l, R1 r)

On Sat, Aug 21, 2021 at 9:44 PM David Feuer <david.feuer at gmail.com> wrote:
>
> We can offer generic helpers for writing MonadZip instances:
>
> genericmzipWith :: (Generic1 m, GMonadZip (Rep1 m)) => (a -> b -> c)
> -> m a -> m b -> m c
> genericmzipWith f m n = to1 $ gmzipWith f (from1 m) (from1 n)
>
> genericmunzip :: (Generic1 m, GMonadZip (Rep1 m)) => m (a, b) -> (m a, m b)
> genericmunzip m = case gmunzip (from1 m) of (r1, r2) -> (to1 r1, to1 r2)
>
> These can be used with appropriate product types (there's no way to
> zip sums generically). I propose adding these to Control.Monad.Zip.
> The GMonadZip class and its instances are below.
>
> class GMonadZip f where
>   gmzipWith :: (a -> b -> c) -> f a -> f b -> f c
>   gmunzip :: f (a, b) -> (f a, f b)
>
> instance MonadZip f => GMonadZip (Rec1 f) where
>     gmzipWith f (Rec1 fa) (Rec1 fb) = Rec1 (mzipWith f fa fb)
>     gmunzip (Rec1 fab) = (Rec1 fa, Rec1 fb)
>       where
>         -- We want to be lazy here, because this might actually
>         -- be recursive. We do a lot of NOINLINE to get selector
>         -- thunks to avoid space leaks.
>         {-# NOINLINE mufab #-}
>         {-# NOINLINE fa #-}
>         {-# NOINLINE fb #-}
>         mufab = munzip fab
>         (fa, fb) = mufab
>
> instance (Functor g, MonadZip f, GMonadZip g) => GMonadZip (f :.: g) where
>   gmzipWith f (Comp1 fga1) (Comp1 fga2) = Comp1 $ mzipWith (gmzipWith
> f) fga1 fga2
>   gmunzip (Comp1 fgc) = case munzip . fmap gmunzip $ fgc of
>     (p, q) -> (Comp1 p, Comp1 q)
>
> -- | @since 4.9.0.0
> instance GMonadZip U1 where
>     gmzipWith _ _ _ = U1
>     gmunzip _ = (U1, U1)
>
> -- | @since 4.9.0.0
> instance GMonadZip Par1 where
>     gmzipWith = coerce
>     gmunzip = coerce
>
> -- | @since 4.9.0.0
> deriving instance GMonadZip f => GMonadZip (M1 i c f)
>
> -- | @since 4.9.0.0
> instance (GMonadZip f, GMonadZip g) => GMonadZip (f :*: g) where
>     gmzipWith f (x1 :*: y1) (x2 :*: y2) = gmzipWith f x1 x2 :*:
> gmzipWith f y1 y2
>     -- Why don't we need to be lazy in this munzip? If we're working with
>     -- Rep1, then laziness will be added by the Rec1 instance. If we're working
>     -- with Rep, then we can't have any K1s because K1 isn't an instance of
>     -- Monad, let alone MonadZip!
>     gmunzip (fab :*: gab)
>       | (fa, fb) <- gmunzip fab
>       , (ga, gb) <- gmunzip gab
>       = (fa :*: ga, fb :*: gb)