{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE PatternSynonyms        #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeInType             #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}

module Data.Type.Util (
    runzipWith
  , rzipWithM_
  , Replicate
  , VecT(.., (:+))
  , vmap
  , withVec
  , vecToRec
  , fillRec
  , zipVecList
  , splitRec
  , p1, p2, s1, s2
  ) where

import           Data.Bifunctor
import           Data.Functor.Identity
import           Data.Kind
import           Data.Proxy
import           Data.Vinyl.Core
import           Data.Vinyl.TypeLevel
import           GHC.Generics
import           Lens.Micro

runzipWith
    :: forall f g h. ()
    => (forall x. f x -> (g x, h x))
    -> (forall xs. Rec f xs -> (Rec g xs, Rec h xs))
runzipWith :: forall {k} (f :: k -> *) (g :: k -> *) (h :: k -> *).
(forall (x :: k). f x -> (g x, h x))
-> forall (xs :: [k]). Rec f xs -> (Rec g xs, Rec h xs)
runzipWith forall (x :: k). f x -> (g x, h x)
f = forall (ys :: [k]). Rec f ys -> (Rec g ys, Rec h ys)
go
  where
    go :: forall ys. Rec f ys -> (Rec g ys, Rec h ys)
    go :: forall (ys :: [k]). Rec f ys -> (Rec g ys, Rec h ys)
go = \case
      Rec f ys
RNil    -> (forall {u} (a :: u -> *). Rec a '[]
RNil, forall {u} (a :: u -> *). Rec a '[]
RNil)
      f r
x :& Rec f rs
xs -> let (g r
y , h r
z ) = forall (x :: k). f x -> (g x, h x)
f f r
x
                     (Rec g rs
ys, Rec h rs
zs) = forall (ys :: [k]). Rec f ys -> (Rec g ys, Rec h ys)
go Rec f rs
xs
                 in  (g r
y forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec g rs
ys, h r
z forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec h rs
zs)
{-# INLINE runzipWith #-}

data VecT :: Nat -> (k -> Type) -> k -> Type where
    VNil :: VecT 'Z f a
    (:*) :: !(f a) -> VecT n f a -> VecT ('S n) f a

pattern (:+) :: a -> VecT n Identity a -> VecT ('S n) Identity a
pattern x $b:+ :: forall a (n :: Nat).
a -> VecT n Identity a -> VecT ('S n) Identity a
$m:+ :: forall {r} {a} {n :: Nat}.
VecT ('S n) Identity a
-> (a -> VecT n Identity a -> r) -> ((# #) -> r) -> r
:+ xs = Identity x :* xs

vmap
    :: forall n f g a. ()
    => (f a -> g a) -> VecT n f a -> VecT n g a
vmap :: forall {k} (n :: Nat) (f :: k -> *) (g :: k -> *) (a :: k).
(f a -> g a) -> VecT n f a -> VecT n g a
vmap f a -> g a
f = forall (m :: Nat). VecT m f a -> VecT m g a
go
  where
    go :: VecT m f a -> VecT m g a
    go :: forall (m :: Nat). VecT m f a -> VecT m g a
go = \case
      VecT m f a
VNil -> forall {k} (f :: k -> *) (a :: k). VecT 'Z f a
VNil
      f a
x :* VecT n f a
xs -> f a -> g a
f f a
x forall {k} (f :: k -> *) (a :: k) (n :: Nat).
f a -> VecT n f a -> VecT ('S n) f a
:* forall (m :: Nat). VecT m f a -> VecT m g a
go VecT n f a
xs
{-# INLINE vmap #-}

withVec
    :: [f a]
    -> (forall n. VecT n f a -> r)
    -> r
withVec :: forall {k} (f :: k -> *) (a :: k) r.
[f a] -> (forall (n :: Nat). VecT n f a -> r) -> r
withVec = \case
    []   -> \forall (n :: Nat). VecT n f a -> r
f -> forall (n :: Nat). VecT n f a -> r
f forall {k} (f :: k -> *) (a :: k). VecT 'Z f a
VNil
    f a
x:[f a]
xs -> \forall (n :: Nat). VecT n f a -> r
f -> forall {k} (f :: k -> *) (a :: k) r.
[f a] -> (forall (n :: Nat). VecT n f a -> r) -> r
withVec [f a]
xs (forall (n :: Nat). VecT n f a -> r
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. (f a
x forall {k} (f :: k -> *) (a :: k) (n :: Nat).
f a -> VecT n f a -> VecT ('S n) f a
:*))
{-# INLINE withVec #-}

type family Replicate (n :: Nat) (a :: k) = (as :: [k]) | as -> n where
    Replicate 'Z     a = '[]
    Replicate ('S n) a = a ': Replicate n a

vecToRec
    :: VecT n f a
    -> Rec f (Replicate n a)
vecToRec :: forall {k} (n :: Nat) (f :: k -> *) (a :: k).
VecT n f a -> Rec f (Replicate n a)
vecToRec = \case
    VecT n f a
VNil    -> forall {u} (a :: u -> *). Rec a '[]
RNil
    f a
x :* VecT n f a
xs -> f a
x forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& forall {k} (n :: Nat) (f :: k -> *) (a :: k).
VecT n f a -> Rec f (Replicate n a)
vecToRec VecT n f a
xs
{-# INLINE vecToRec #-}

fillRec
    :: forall f g as c. ()
    => (forall a. f a -> c -> g a)
    -> Rec f as
    -> [c]
    -> Maybe (Rec g as)
fillRec :: forall {u} (f :: u -> *) (g :: u -> *) (as :: [u]) c.
(forall (a :: u). f a -> c -> g a)
-> Rec f as -> [c] -> Maybe (Rec g as)
fillRec forall (a :: u). f a -> c -> g a
f = forall (bs :: [u]). Rec f bs -> [c] -> Maybe (Rec g bs)
go
  where
    go :: Rec f bs -> [c] -> Maybe (Rec g bs)
    go :: forall (bs :: [u]). Rec f bs -> [c] -> Maybe (Rec g bs)
go = \case
      Rec f bs
RNil -> \[c]
_ -> forall a. a -> Maybe a
Just forall {u} (a :: u -> *). Rec a '[]
RNil
      f r
x :& Rec f rs
xs -> \case
        []   -> forall a. Maybe a
Nothing
        c
y:[c]
ys -> (forall (a :: u). f a -> c -> g a
f f r
x c
y forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:&) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (bs :: [u]). Rec f bs -> [c] -> Maybe (Rec g bs)
go Rec f rs
xs [c]
ys
{-# INLINE fillRec #-}

rzipWithM_
    :: forall h f g as. Applicative h
    => (forall a. f a -> g a -> h ())
    -> Rec f as
    -> Rec g as
    -> h ()
rzipWithM_ :: forall {u} (h :: * -> *) (f :: u -> *) (g :: u -> *) (as :: [u]).
Applicative h =>
(forall (a :: u). f a -> g a -> h ())
-> Rec f as -> Rec g as -> h ()
rzipWithM_ forall (a :: u). f a -> g a -> h ()
f = forall (bs :: [u]). Rec f bs -> Rec g bs -> h ()
go
  where
    go :: forall bs. Rec f bs -> Rec g bs -> h ()
    go :: forall (bs :: [u]). Rec f bs -> Rec g bs -> h ()
go = \case
      Rec f bs
RNil -> \case
        Rec g bs
RNil -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      f r
x :& Rec f rs
xs -> \case
        g r
y :& Rec g rs
ys -> forall (a :: u). f a -> g a -> h ()
f f r
x g r
y forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall (bs :: [u]). Rec f bs -> Rec g bs -> h ()
go Rec f rs
xs Rec g rs
ys
{-# INLINE rzipWithM_ #-}

zipVecList
    :: forall a b c f g n. ()
    => (f a -> Maybe b -> g c)
    -> VecT n f a
    -> [b]
    -> VecT n g c
zipVecList :: forall {k} {k} (a :: k) b (c :: k) (f :: k -> *) (g :: k -> *)
       (n :: Nat).
(f a -> Maybe b -> g c) -> VecT n f a -> [b] -> VecT n g c
zipVecList f a -> Maybe b -> g c
f = forall (m :: Nat). VecT m f a -> [b] -> VecT m g c
go
  where
    go :: VecT m f a -> [b] -> VecT m g c
    go :: forall (m :: Nat). VecT m f a -> [b] -> VecT m g c
go = \case
      VecT m f a
VNil -> forall a b. a -> b -> a
const forall {k} (f :: k -> *) (a :: k). VecT 'Z f a
VNil
      f a
x :* VecT n f a
xs -> \case
        []   -> f a -> Maybe b -> g c
f f a
x forall a. Maybe a
Nothing  forall {k} (f :: k -> *) (a :: k) (n :: Nat).
f a -> VecT n f a -> VecT ('S n) f a
:* forall (m :: Nat). VecT m f a -> [b] -> VecT m g c
go VecT n f a
xs []
        b
y:[b]
ys -> f a -> Maybe b -> g c
f f a
x (forall a. a -> Maybe a
Just b
y) forall {k} (f :: k -> *) (a :: k) (n :: Nat).
f a -> VecT n f a -> VecT ('S n) f a
:* forall (m :: Nat). VecT m f a -> [b] -> VecT m g c
go VecT n f a
xs [b]
ys
{-# INLINE zipVecList #-}

splitRec
    :: forall f as bs. RecApplicative as
    => Rec f (as ++ bs)
    -> (Rec f as, Rec f bs)
splitRec :: forall {u} (f :: u -> *) (as :: [u]) (bs :: [u]).
RecApplicative as =>
Rec f (as ++ bs) -> (Rec f as, Rec f bs)
splitRec = forall (as' :: [u]).
Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs)
go (forall {u} (rs :: [u]) (f :: u -> *).
RecApplicative rs =>
(forall (x :: u). f x) -> Rec f rs
rpure forall {k} (t :: k). Proxy t
Proxy)
  where
    go :: Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs)
    go :: forall (as' :: [u]).
Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs)
go = \case
      Rec Proxy as'
RNil -> (forall {u} (a :: u -> *). Rec a '[]
RNil,)
      Proxy r
_ :& Rec Proxy rs
ps -> \case
        f r
x :& Rec f rs
xs -> forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (f r
x forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:&) forall a b. (a -> b) -> a -> b
$ forall (as' :: [u]).
Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs)
go Rec Proxy rs
ps Rec f rs
xs
{-# INLINE splitRec #-}

p1 :: Lens' ((f :*: g) a) (f a)
p1 :: forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Lens' ((:*:) f g a) (f a)
p1 f a -> f (f a)
f (f a
x :*: g a
y) = (forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: g a
y) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a -> f (f a)
f f a
x
{-# INLINE p1 #-}

p2 :: Lens' ((f :*: g) a) (g a)
p2 :: forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Lens' ((:*:) f g a) (g a)
p2 g a -> f (g a)
f (f a
x :*: g a
y) = (f a
x forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g a -> f (g a)
f g a
y
{-# INLINE p2 #-}

s1 :: Traversal' ((f :+: g) a) (f a)
s1 :: forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Traversal' ((:+:) f g a) (f a)
s1 f a -> f (f a)
f (L1 f a
x) = forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a -> f (f a)
f f a
x
s1 f a -> f (f a)
_ (R1 g a
y) = forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1 g a
y)
{-# INLINE s1 #-}

s2 :: Traversal' ((f :+: g) a) (g a)
s2 :: forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Traversal' ((:+:) f g a) (g a)
s2 g a -> f (g a)
_ (L1 f a
x) = forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1 f a
x)
s2 g a -> f (g a)
f (R1 g a
y) = forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g a -> f (g a)
f g a
y
{-# INLINE s2 #-}