Util.hs

{-# 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 #-}