Zip.lhs

> {-# LANGUAGE DataKinds, FlexibleContexts, FlexibleInstances,
>     TypeFamilies, MultiParamTypeClasses, UndecidableInstances,
>     RankNTypes, ScopedTypeVariables, FunctionalDependencies, TypeOperators #-}
>
> module Zip where
>
> import Data.Proxy
> import GHC.TypeLits

Let's first assume that the two nested lists have the same depth. E.g., depth 2:

< zipDeep0 ((+) :: Int -> Int -> Int) [[1,2],[3,4,5]] [[10,20],[30,40]] :: [[Int]]
< [[11,22],[33,44]]

Implementation:

> zipDeep0
>   :: forall n a b c x y z
>   .  (ZipDeep0 n a b c x y z, n ~ Levels a x, n ~ Levels b y, n ~ Levels c z)
>   => (a -> b -> c) -> (x -> y -> z)
> zipDeep0 = zipDeep0_ (Proxy :: Proxy n)

`Levels a x` computes the depth of a in the nested list type x.
Closed type families allow us to do some non-linear type-level pattern matching
(where `a` occurs twice in a clause).

> type family Levels a x :: Nat where
>   Levels a a = 0
>   Levels a [x] = 1 + Levels a x

We use that depth to select the `ZipDeep0` instance implementing the zip,
This way is neater than relying only on the six other ordinary type parameters,
as it avoids a problem with overlapping instances when some list is empty (so
we can't infer its actual type from itself), or when `a`, `b`, `c` are also
list types.

> class ZipDeep0 (n :: Nat) a b c x y z where
>   zipDeep0_ :: proxy n -> (a -> b -> c) -> x -> y -> z
>
> -- Moving the equality constraints into the context helps type inference.
> instance {-# OVERLAPPING #-} (a ~ x, b ~ y, c ~ z) => ZipDeep0 0 a b c x y z where
>   zipDeep0_ _ = id
>
> instance (ZipDeep0 (n - 1) a b c x y z, xs ~ [x], ys ~ [y], zs ~ [z])
>   => ZipDeep0 n a b c xs ys zs where
>   zipDeep0_ _ f = zipWith (zipDeep0_ (Proxy :: Proxy (n - 1)) f)

---

When the two lists are not of the same depth, let's first assume the second one
is deeper, so we must first distribute over it.
We start losing some type inference, we must know at least `Levels a x` (and
thus `a` and `x`) and either `Levels b y` or `Levels c z` before this function
can be applied.

Example:

< zipDeep1 (+) [10,20 :: Int] [[1,2],[3,4]] :: [[Int]]
< [[11,22],[13,24]]

Implementation:

> zipDeep1
>  :: forall n m a b c x y z
>  .  (n ~ Levels a x, m ~ Levels b y, m ~ Levels c z, ZipDeep1 (m - n) a b c x y z)
>  => (a -> b -> c) -> x -> y -> z
> zipDeep1 = zipDeep1_ (Proxy :: Proxy (m - n))

The difference between levels `(m - n)` tells us how many layers we must "distribute"
through before falling back to `zipDeep0`.

> class ZipDeep1 (n :: Nat) a b c x y z where
>   zipDeep1_ :: proxy n -> (a -> b -> c) -> x -> y -> z
>
> instance {-# OVERLAPPING #-} ZipDeep0 (Levels a x) a b c x y z => ZipDeep1 0 a b c x y z where
>   zipDeep1_ _ = zipDeep0_ (Proxy :: Proxy (Levels a x))
>
> instance (ZipDeep1 (n - 1) a b c x y z, ys ~ [y], zs ~ [z]) => ZipDeep1 n a b c x ys zs where
>   zipDeep1_ proxy f xs = fmap (zipDeep1_ (Proxy :: Proxy (n - 1)) f xs)

---

Finally, we can do a type-level comparison when either list may be the deeper
one. We lose all type inference though.

Example:

< zipDeep ((+) :: Int -> Int -> Int) [[1,2 :: Int],[3,4]] [10 :: Int,20] :: [[Int]]
< [[11,22],[13,24]]

Some type inference is recovered by instead specifying the expected depth of
each list with TypeApplications.

< zipDeep @2 @1 ((+) :: Int -> Int -> Int) [[1,2],[3,4]] [10,20]
< [[11,22],[13,24]]

Implementation:

> zipDeep
>   :: forall n m a b c x y z
>   .  (ZipDeep2 (CmpNat n m) n m a b c x y z, n ~ Levels a x, m ~ Levels b y)
>   => (a -> b -> c) -> x -> y -> z
> zipDeep = zipDeep2_ (Proxy :: Proxy '(CmpNat n m, n, m))

> class ZipDeep2 (cmp :: Ordering) (n :: Nat) (m :: Nat) a b c x y z where
>   zipDeep2_ :: proxy '(cmp, n, m) -> (a -> b -> c) -> x -> y -> z

> instance {-# OVERLAPPING #-} (n ~ Levels a x, m ~ Levels b y, m ~ Levels c z, ZipDeep1 (m - n) a b c x y z)
>   => ZipDeep2 'LT n m a b c x y z where
>   zipDeep2_ _ = zipDeep1
>
> instance (n ~ Levels a x, m ~ Levels b y, n ~ Levels c z, ZipDeep1 (n - m) b a c y x z)
>   => ZipDeep2 cmp n m a b c x y z where
>   zipDeep2_ _ = flip . zipDeep1 . flip