{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# 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 = Rec f xs -> (Rec g xs, Rec h xs)
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 -> (Rec g ys
Rec g '[]
forall {u} (a :: u -> *). Rec a '[]
RNil, Rec h ys
Rec h '[]
forall {u} (a :: u -> *). Rec a '[]
RNil)
      f r
x :& Rec f rs
xs ->
        let (g r
y, h r
z) = f r -> (g r, h r)
forall (x :: k). f x -> (g x, h x)
f f r
x
            (Rec g rs
ys, Rec h rs
zs) = Rec f rs -> (Rec g rs, Rec h rs)
forall (ys :: [k]). Rec f ys -> (Rec g ys, Rec h ys)
go Rec f rs
xs
         in (g r
y g r -> Rec g rs -> Rec g (r : rs)
forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec g rs
ys, h r
z h r -> Rec h rs -> Rec h (r : rs)
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 $m:+ :: forall {r} {a} {n :: Nat}.
VecT ('S n) Identity a
-> (a -> VecT n Identity a -> r) -> ((# #) -> r) -> r
$b:+ :: forall a (n :: Nat).
a -> VecT n Identity a -> VecT ('S n) Identity a
:+ 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 = VecT n f a -> VecT n g a
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 -> VecT m g a
VecT 'Z g a
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 g a -> VecT n g a -> VecT ('S n) g a
forall {k} (f :: k -> *) (a :: k) (n :: Nat).
f a -> VecT n f a -> VecT ('S n) f a
:* VecT n f a -> VecT n g 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 -> VecT 'Z f a -> r
forall (n :: Nat). VecT n f a -> r
f VecT 'Z f a
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 -> [f a] -> (forall (n :: Nat). VecT n f a -> r) -> r
forall {k} (f :: k -> *) (a :: k) r.
[f a] -> (forall (n :: Nat). VecT n f a -> r) -> r
withVec [f a]
xs (VecT ('S n) f a -> r
forall (n :: Nat). VecT n f a -> r
f (VecT ('S n) f a -> r)
-> (VecT n f a -> VecT ('S n) f a) -> VecT n f a -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (f a
x f a -> VecT n f a -> VecT ('S n) f a
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 -> Rec f '[]
Rec f (Replicate n a)
forall {u} (a :: u -> *). Rec a '[]
RNil
  f a
x :* VecT n f a
xs -> f a
x f a -> Rec f (Replicate n a) -> Rec f (a : Replicate n a)
forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& VecT n f a -> Rec f (Replicate n a)
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 = Rec f as -> [c] -> Maybe (Rec g as)
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]
_ -> Rec g bs -> Maybe (Rec g bs)
forall a. a -> Maybe a
Just Rec g bs
Rec g '[]
forall {u} (a :: u -> *). Rec a '[]
RNil
      f r
x :& Rec f rs
xs -> \case
        [] -> Maybe (Rec g bs)
forall a. Maybe a
Nothing
        c
y : [c]
ys -> (f r -> c -> g r
forall (a :: u). f a -> c -> g a
f f r
x c
y g r -> Rec g rs -> Rec g (r : rs)
forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:&) (Rec g rs -> Rec g bs) -> Maybe (Rec g rs) -> Maybe (Rec g bs)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rec f rs -> [c] -> Maybe (Rec g rs)
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 = Rec f as -> Rec g as -> h ()
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 -> () -> h ()
forall a. a -> h a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      f r
x :& Rec f rs
xs -> \case
        g r
y :& Rec g rs
ys -> f r -> g r -> h ()
forall (a :: u). f a -> g a -> h ()
f f r
x g r
g r
y h () -> h () -> h ()
forall a b. h a -> h b -> h b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Rec f rs -> Rec g rs -> h ()
forall (bs :: [u]). Rec f bs -> Rec g bs -> h ()
go Rec f rs
xs Rec g rs
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 = VecT n f a -> [b] -> VecT n g c
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 -> VecT m g c -> [b] -> VecT m g c
forall a b. a -> b -> a
const VecT m g c
VecT 'Z g c
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 Maybe b
forall a. Maybe a
Nothing g c -> VecT n g c -> VecT ('S n) g c
forall {k} (f :: k -> *) (a :: k) (n :: Nat).
f a -> VecT n f a -> VecT ('S n) f a
:* VecT n f a -> [b] -> VecT n g c
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 (b -> Maybe b
forall a. a -> Maybe a
Just b
y) g c -> VecT n g c -> VecT ('S n) g c
forall {k} (f :: k -> *) (a :: k) (n :: Nat).
f a -> VecT n f a -> VecT ('S n) f a
:* VecT n f a -> [b] -> VecT n g c
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 = Rec Proxy as -> Rec f (as ++ bs) -> (Rec f as, Rec f bs)
forall (as' :: [u]).
Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs)
go ((forall (x :: u). Proxy x) -> Rec Proxy as
forall {u} (rs :: [u]) (f :: u -> *).
RecApplicative rs =>
(forall (x :: u). f x) -> Rec f rs
forall (f :: u -> *). (forall (x :: u). f x) -> Rec f as
rpure Proxy x
forall (x :: u). Proxy x
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 -> (Rec f '[]
forall {u} (a :: u -> *). Rec a '[]
RNil,)
      Proxy r
_ :& Rec Proxy rs
ps -> \case
        f r
x :& Rec f rs
xs -> (Rec f rs -> Rec f as')
-> (Rec f rs, Rec f bs) -> (Rec f as', Rec f bs)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (f r
x f r -> Rec f rs -> Rec f (r : rs)
forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:&) ((Rec f rs, Rec f bs) -> (Rec f as', Rec f bs))
-> (Rec f rs, Rec f bs) -> (Rec f as', Rec f bs)
forall a b. (a -> b) -> a -> b
$ Rec Proxy rs -> Rec f (rs ++ bs) -> (Rec f rs, Rec f bs)
forall (as' :: [u]).
Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs)
go Rec Proxy rs
ps Rec f rs
Rec f (rs ++ bs)
xs
{-# INLINE splitRec #-}

p1 :: Lens' ((f :*: g) a) (f a)
p1 :: forall {k} (f :: k -> *) (g :: k -> *) (a :: k) (f :: * -> *).
Functor f =>
(f a -> f (f a)) -> (:*:) f g a -> f ((:*:) f g a)
p1 f a -> f (f a)
f (f a
x :*: g a
y) = (f a -> g a -> (:*:) f g a
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: g a
y) (f a -> (:*:) f g a) -> f (f a) -> f ((:*:) f g a)
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) (f :: * -> *).
Functor f =>
(g a -> f (g a)) -> (:*:) f g a -> f ((:*:) f g a)
p2 g a -> f (g a)
f (f a
x :*: g a
y) = (f a
x f a -> g a -> (:*:) f g a
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*:) (g a -> (:*:) f g a) -> f (g a) -> f ((:*:) f g a)
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) (f :: * -> *).
Applicative f =>
(f a -> f (f a)) -> (:+:) f g a -> f ((:+:) f g a)
s1 f a -> f (f a)
f (L1 f a
x) = f a -> (:+:) f g a
forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1 (f a -> (:+:) f g a) -> f (f a) -> f ((:+:) f g a)
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) = (:+:) f g a -> f ((:+:) f g a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (g a -> (:+:) f g a
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) (f :: * -> *).
Applicative f =>
(g a -> f (g a)) -> (:+:) f g a -> f ((:+:) f g a)
s2 g a -> f (g a)
_ (L1 f a
x) = (:+:) f g a -> f ((:+:) f g a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (f a -> (:+:) f g a
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) = g a -> (:+:) f g a
forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1 (g a -> (:+:) f g a) -> f (g a) -> f ((:+:) f g a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g a -> f (g a)
f g a
y
{-# INLINE s2 #-}