module Data.Zip where

import qualified Data.NonEmpty.Class as C

import qualified Data.Traversable as Trav
import Data.Traversable (Traversable, )
import Control.Applicative (Applicative, pure, (<*>), )
import Control.DeepSeq (NFData, rnf, )


{- |
Wrap a container such that its Applicative instance is based on zip.
-}
newtype T f a = Cons {T f a -> f a
decons :: f a}

instance Functor f => Functor (T f) where
   fmap :: (a -> b) -> T f a -> T f b
fmap a -> b
f (Cons f a
xs) = f b -> T f b
forall (f :: * -> *) a. f a -> T f a
Cons (f b -> T f b) -> f b -> T f b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f f a
xs

instance (C.Zip f, C.Repeat f) => Applicative (T f) where
   pure :: a -> T f a
pure a
a = f a -> T f a
forall (f :: * -> *) a. f a -> T f a
Cons (f a -> T f a) -> f a -> T f a
forall a b. (a -> b) -> a -> b
$ a -> f a
forall (f :: * -> *) a. Repeat f => a -> f a
C.repeat a
a
   Cons f (a -> b)
f <*> :: T f (a -> b) -> T f a -> T f b
<*> Cons f a
x = f b -> T f b
forall (f :: * -> *) a. f a -> T f a
Cons (f b -> T f b) -> f b -> T f b
forall a b. (a -> b) -> a -> b
$ ((a -> b) -> a -> b) -> f (a -> b) -> f a -> f b
forall (f :: * -> *) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
C.zipWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($) f (a -> b)
f f a
x


instance (C.NFData f, NFData a) => NFData (T f a) where
   rnf :: T f a -> ()
rnf = T f a -> ()
forall (f :: * -> *) a. (NFData f, NFData a) => f a -> ()
C.rnf

instance (C.NFData f) => C.NFData (T f) where
   rnf :: T f a -> ()
rnf = f a -> ()
forall (f :: * -> *) a. (NFData f, NFData a) => f a -> ()
C.rnf (f a -> ()) -> (T f a -> f a) -> T f a -> ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T f a -> f a
forall (f :: * -> *) a. T f a -> f a
decons


{- |
Always returns a rectangular list
by clipping all dimensions to the shortest slice.
Be aware that @transpose [] == repeat []@.
-}
transposeClip ::
   (Traversable f, C.Zip g, C.Repeat g) =>
   f (g a) -> g (f a)
transposeClip :: f (g a) -> g (f a)
transposeClip =
   T g (f a) -> g (f a)
forall (f :: * -> *) a. T f a -> f a
decons (T g (f a) -> g (f a))
-> (f (g a) -> T g (f a)) -> f (g a) -> g (f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (T g a) -> T g (f a)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
Trav.sequenceA (f (T g a) -> T g (f a))
-> (f (g a) -> f (T g a)) -> f (g a) -> T g (f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (g a -> T g a) -> f (g a) -> f (T g a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap g a -> T g a
forall (f :: * -> *) a. f a -> T f a
Cons