{-# 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 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
[] -> ($ 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 #-}