-- | Overloaded 'zip' and 'zipWith'.
--
-- This module provides the 'Zip' class to overload zipping functions
-- for variants of the list type ('[]', 'List1', 'List2', 'ListInf').
--
-- Motivation:
-- In practice, we often zip a finite list with an infinite list,
-- e.g. @zipWith f xs (ys ++ repeat z)@.
-- We can know statically in this case that the resulting considers
-- each @xs@ by rewriting this to @zipWith f xs (ListInf.pad ys z)@
-- statically exposing the infinite nature of the second list.
-- To avoid cluttering the namespace with 4*4 variants of 'zipWith'
-- for all combination of our list types, we introduce an overloaded version.
--
-- Import this module as follows:
--
-- @
-- import Prelude hiding (zip, zipWith)
-- import Agda.Utils.Zip
-- @

module Agda.Utils.Zip where

import Prelude ((.), flip)

import Data.List          qualified as List
import Data.List.Infinite ( heteroZipWith )

-- import Agda.Utils.List    ( pattern (:) )
import Agda.Utils.List1   ( List1, pattern (:|) )
import Agda.Utils.List2   ( List2, pattern List2 )
import Agda.Utils.ListInf ( ListInf, pattern (:<) )

import Agda.Utils.List    qualified as List
import Agda.Utils.List1   qualified as List1
import Agda.Utils.List2   qualified as List2
import Agda.Utils.ListInf qualified as ListInf

class Zip f g h | f g -> h where
  zip :: f a -> g b -> h (a, b)
  zip = (a -> b -> (a, b)) -> f a -> g b -> h (a, b)
forall a b c. (a -> b -> c) -> f a -> g b -> h c
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) a b c.
Zip f g h =>
(a -> b -> c) -> f a -> g b -> h c
zipWith (,)

  zipWith :: (a -> b -> c) -> f a -> g b -> h c
  -- zipWith f as bs = map (uncurry f) (zip as bs) -- needs Functor h

  {-# MINIMAL zipWith #-}

-- List instances

-- 0/0
instance Zip [] [] [] where
  zipWith :: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith = (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith
  zip :: forall a b. [a] -> [b] -> [(a, b)]
zip     = [a] -> [b] -> [(a, b)]
forall a b. [a] -> [b] -> [(a, b)]
List.zip

-- 0/1
instance Zip [] List1 [] where
  zipWith :: forall a b c. (a -> b -> c) -> [a] -> List1 b -> [c]
zipWith a -> b -> c
_ [] List1 b
_ = []
  zipWith a -> b -> c
f (a
a : [a]
as) (b
b :| [b]
bs) = a -> b -> c
f a
a b
b c -> [c] -> [c]
forall a. a -> [a] -> [a]
: (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) a b c.
Zip f g h =>
(a -> b -> c) -> f a -> g b -> h c
zipWith a -> b -> c
f [a]
as [b]
bs

-- 1/0
instance Zip List1 [] [] where
  zipWith :: forall a b c. (a -> b -> c) -> List1 a -> [b] -> [c]
zipWith a -> b -> c
_ List1 a
_ [] = []
  zipWith a -> b -> c
f (a
a :| [a]
as) (b
b : [b]
bs) = a -> b -> c
f a
a b
b c -> [c] -> [c]
forall a. a -> [a] -> [a]
: (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) a b c.
Zip f g h =>
(a -> b -> c) -> f a -> g b -> h c
zipWith a -> b -> c
f [a]
as [b]
bs

-- 0/2
instance Zip [] List2 [] where
  zipWith :: forall a b c. (a -> b -> c) -> [a] -> List2 b -> [c]
zipWith a -> b -> c
_ []  List2 b
_ = []
  zipWith a -> b -> c
f (a
a1 : [a]
as) (List2 b
b1 b
b2 [b]
bs) = a -> b -> c
f a
a1 b
b1 c -> [c] -> [c]
forall a. a -> [a] -> [a]
: (a -> b -> c) -> [a] -> NonEmpty b -> [c]
forall a b c. (a -> b -> c) -> [a] -> List1 b -> [c]
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) a b c.
Zip f g h =>
(a -> b -> c) -> f a -> g b -> h c
zipWith a -> b -> c
f [a]
as (b
b2 b -> [b] -> NonEmpty b
forall a. a -> [a] -> NonEmpty a
:| [b]
bs)

-- 2/0
instance Zip List2 [] [] where
  zipWith :: forall a b c. (a -> b -> c) -> List2 a -> [b] -> [c]
zipWith a -> b -> c
_ List2 a
_ [] = []
  zipWith a -> b -> c
f (List2 a
a1 a
a2 [a]
as) (b
b1 : [b]
bs) = a -> b -> c
f a
a1 b
b1 c -> [c] -> [c]
forall a. a -> [a] -> [a]
: (a -> b -> c) -> NonEmpty a -> [b] -> [c]
forall a b c. (a -> b -> c) -> List1 a -> [b] -> [c]
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) a b c.
Zip f g h =>
(a -> b -> c) -> f a -> g b -> h c
zipWith a -> b -> c
f (a
a2 a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
as) [b]
bs

-- 0/∞
instance Zip [] ListInf [] where
  zipWith :: forall a b c. (a -> b -> c) -> [a] -> ListInf b -> [c]
zipWith = (ListInf b -> [a] -> [c]) -> [a] -> ListInf b -> [c]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((ListInf b -> [a] -> [c]) -> [a] -> ListInf b -> [c])
-> ((a -> b -> c) -> ListInf b -> [a] -> [c])
-> (a -> b -> c)
-> [a]
-> ListInf b
-> [c]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> a -> c) -> ListInf b -> [a] -> [c]
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> c) -> Infinite a -> t b -> t c
heteroZipWith ((b -> a -> c) -> ListInf b -> [a] -> [c])
-> ((a -> b -> c) -> b -> a -> c)
-> (a -> b -> c)
-> ListInf b
-> [a]
-> [c]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip

-- ∞/0
instance Zip ListInf [] [] where
  zipWith :: forall a b c. (a -> b -> c) -> ListInf a -> [b] -> [c]
zipWith = (a -> b -> c) -> Infinite a -> [b] -> [c]
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> c) -> Infinite a -> t b -> t c
heteroZipWith

-- List1 instances

-- 1/1
instance Zip List1 List1 List1 where
  zipWith :: forall a b c. (a -> b -> c) -> List1 a -> List1 b -> List1 c
zipWith = (a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
forall a b c. (a -> b -> c) -> List1 a -> List1 b -> List1 c
List1.zipWith
  zip :: forall a b. List1 a -> List1 b -> List1 (a, b)
zip     = NonEmpty a -> NonEmpty b -> NonEmpty (a, b)
forall a b. List1 a -> List1 b -> List1 (a, b)
List1.zip

-- 1/2
instance Zip List1 List2 List1 where
  zipWith :: forall a b c. (a -> b -> c) -> List1 a -> List2 b -> List1 c
zipWith a -> b -> c
f (a
a1 :| [a]
as) (List2 b
b1 b
b2 [b]
bs) = a -> b -> c
f a
a1 b
b1 c -> [c] -> NonEmpty c
forall a. a -> [a] -> NonEmpty a
:| (a -> b -> c) -> [a] -> NonEmpty b -> [c]
forall a b c. (a -> b -> c) -> [a] -> List1 b -> [c]
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) a b c.
Zip f g h =>
(a -> b -> c) -> f a -> g b -> h c
zipWith a -> b -> c
f [a]
as (b
b2 b -> [b] -> NonEmpty b
forall a. a -> [a] -> NonEmpty a
:| [b]
bs)

-- 2/1
instance Zip List2 List1 List1 where
  zipWith :: forall a b c. (a -> b -> c) -> List2 a -> List1 b -> List1 c
zipWith a -> b -> c
f (List2 a
a1 a
a2 [a]
as) (b
b1 :| [b]
bs) = a -> b -> c
f a
a1 b
b1 c -> [c] -> NonEmpty c
forall a. a -> [a] -> NonEmpty a
:| (a -> b -> c) -> NonEmpty a -> [b] -> [c]
forall a b c. (a -> b -> c) -> List1 a -> [b] -> [c]
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) a b c.
Zip f g h =>
(a -> b -> c) -> f a -> g b -> h c
zipWith a -> b -> c
f (a
a2 a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
as) [b]
bs

-- 1/∞
instance Zip List1 ListInf List1 where
  zipWith :: forall a b c. (a -> b -> c) -> List1 a -> ListInf b -> List1 c
zipWith = (ListInf b -> List1 a -> List1 c)
-> List1 a -> ListInf b -> List1 c
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((ListInf b -> List1 a -> List1 c)
 -> List1 a -> ListInf b -> List1 c)
-> ((a -> b -> c) -> ListInf b -> List1 a -> List1 c)
-> (a -> b -> c)
-> List1 a
-> ListInf b
-> List1 c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> a -> c) -> ListInf b -> List1 a -> List1 c
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> c) -> Infinite a -> t b -> t c
heteroZipWith ((b -> a -> c) -> ListInf b -> List1 a -> List1 c)
-> ((a -> b -> c) -> b -> a -> c)
-> (a -> b -> c)
-> ListInf b
-> List1 a
-> List1 c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip

-- ∞/1
instance Zip ListInf List1 List1 where
  zipWith :: forall a b c. (a -> b -> c) -> ListInf a -> List1 b -> List1 c
zipWith = (a -> b -> c) -> Infinite a -> NonEmpty b -> NonEmpty c
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> c) -> Infinite a -> t b -> t c
heteroZipWith

-- List2 instances

-- 2/2
instance Zip List2 List2 List2 where
  zipWith :: forall a b c. (a -> b -> c) -> List2 a -> List2 b -> List2 c
zipWith = (a -> b -> c) -> List2 a -> List2 b -> List2 c
forall a b c. (a -> b -> c) -> List2 a -> List2 b -> List2 c
List2.zipWith
  zip :: forall a b. List2 a -> List2 b -> List2 (a, b)
zip     = List2 a -> List2 b -> List2 (a, b)
forall a b. List2 a -> List2 b -> List2 (a, b)
List2.zip

-- 2/∞
instance Zip List2 ListInf List2 where
  zipWith :: forall a b c. (a -> b -> c) -> List2 a -> ListInf b -> List2 c
zipWith = (ListInf b -> List2 a -> List2 c)
-> List2 a -> ListInf b -> List2 c
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((ListInf b -> List2 a -> List2 c)
 -> List2 a -> ListInf b -> List2 c)
-> ((a -> b -> c) -> ListInf b -> List2 a -> List2 c)
-> (a -> b -> c)
-> List2 a
-> ListInf b
-> List2 c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> a -> c) -> ListInf b -> List2 a -> List2 c
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> c) -> Infinite a -> t b -> t c
heteroZipWith ((b -> a -> c) -> ListInf b -> List2 a -> List2 c)
-> ((a -> b -> c) -> b -> a -> c)
-> (a -> b -> c)
-> ListInf b
-> List2 a
-> List2 c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip

-- ∞/2
instance Zip ListInf List2 List2 where
  zipWith :: forall a b c. (a -> b -> c) -> ListInf a -> List2 b -> List2 c
zipWith = (a -> b -> c) -> Infinite a -> List2 b -> List2 c
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> c) -> Infinite a -> t b -> t c
heteroZipWith

-- ListInf instances

-- ∞/∞
instance Zip ListInf ListInf ListInf where
  zipWith :: forall a b c. (a -> b -> c) -> ListInf a -> ListInf b -> ListInf c
zipWith = (a -> b -> c) -> Infinite a -> Infinite b -> Infinite c
forall a b c. (a -> b -> c) -> ListInf a -> ListInf b -> ListInf c
ListInf.zipWith
  zip :: forall a b. ListInf a -> ListInf b -> ListInf (a, b)
zip     = Infinite a -> Infinite b -> Infinite (a, b)
forall a b. ListInf a -> ListInf b -> ListInf (a, b)
ListInf.zip