content: {-# 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 f = go
    where
      go :: forall ys. Rec f ys -> (Rec g ys, Rec h ys)
      go = \case
        RNil -> (RNil, RNil)
        x :& xs ->
          let (y, z) = f x
              (ys, zs) = go xs
           in (y :& ys, z :& 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 :+ xs = Identity x :* xs
  
  vmap ::
    forall n f g a.
    () =>
    (f a -> g a) -> VecT n f a -> VecT n g a
  vmap f = go
    where
      go :: VecT m f a -> VecT m g a
      go = \case
        VNil -> VNil
        x :* xs -> f x :* go xs
  {-# INLINE vmap #-}
  
  withVec ::
    [f a] ->
    (forall n. VecT n f a -> r) ->
    r
  withVec = \case
    [] -> \f -> f VNil
    x : xs -> \f -> withVec xs (f . (x :*))
  {-# 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 = \case
    VNil -> RNil
    x :* xs -> x :& vecToRec 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 f = go
    where
      go :: Rec f bs -> [c] -> Maybe (Rec g bs)
      go = \case
        RNil -> \_ -> Just RNil
        x :& xs -> \case
          [] -> Nothing
          y : ys -> (f x y :&) <$> go xs 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_ f = go
    where
      go :: forall bs. Rec f bs -> Rec g bs -> h ()
      go = \case
        RNil -> \case
          RNil -> pure ()
        x :& xs -> \case
          y :& ys -> f x y *> go xs 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 f = go
    where
      go :: VecT m f a -> [b] -> VecT m g c
      go = \case
        VNil -> const VNil
        x :* xs -> \case
          [] -> f x Nothing :* go xs []
          y : ys -> f x (Just y) :* go xs ys
  {-# INLINE zipVecList #-}
  
  splitRec ::
    forall f as bs.
    RecApplicative as =>
    Rec f (as ++ bs) ->
    (Rec f as, Rec f bs)
  splitRec = go (rpure Proxy)
    where
      go :: Rec Proxy as' -> Rec f (as' ++ bs) -> (Rec f as', Rec f bs)
      go = \case
        RNil -> (RNil,)
        _ :& ps -> \case
          x :& xs -> first (x :&) $ go ps xs
  {-# INLINE splitRec #-}
  
  p1 :: Lens' ((f :*: g) a) (f a)
  p1 f (x :*: y) = (:*: y) <$> f x
  {-# INLINE p1 #-}
  
  p2 :: Lens' ((f :*: g) a) (g a)
  p2 f (x :*: y) = (x :*:) <$> f y
  {-# INLINE p2 #-}
  
  s1 :: Traversal' ((f :+: g) a) (f a)
  s1 f (L1 x) = L1 <$> f x
  s1 _ (R1 y) = pure (R1 y)
  {-# INLINE s1 #-}
  
  s2 :: Traversal' ((f :+: g) a) (g a)
  s2 _ (L1 x) = pure (L1 x)
  s2 f (R1 y) = R1 <$> f y
  {-# INLINE s2 #-}
  