module Data.Microgroove (
Rec(Rec#,RNil,(:&))
,thaw, thaw#, freeze, freeze#
, new, new'
,fromVectorN, fromVector, replicate
,modify
,rmap, rmapM
,crmap, crmapM
,rappend
,rzip
,crzip
,splitCons
,foldMapF
,toVector
,cfoldMapF
,ctoVector
,subRecord#
,Index(Index#,RZ,RS)
,mkIndex, checkIndex, checkDynIndex
,index, (!)
,module X
) where
import Prelude hiding (replicate)
import Data.Microgroove.Lib
import Data.Microgroove.Lib.TypeLevel as X
import Data.Microgroove.Index
import Data.Microgroove.Mutable (MRec(..))
import qualified Data.Microgroove.Mutable as M
import qualified Data.Vector as V
import Data.Vector as X (Vector)
import Control.Monad.ST
import Control.Monad.Primitive (PrimMonad(..))
import GHC.TypeLits
import GHC.Float
import GHC.Exts (IsList(..))
import Data.Monoid (Endo(..))
newtype Rec (f :: u -> *) (us :: [u]) = Rec# (V.Vector Any)
instance Show (Rec f '[]) where
show RNil = "[]"
show _ = error "Impossible! RNil inexhaustive in show @(Rec f '[])"
instance (Show (f x), Show (Rec f xs)) => Show (Rec f (x ': xs)) where
show (a :& xs) = show a ++ " :& " ++ show xs
show _ = error "Impossible! RCons inexhaustive in show @(Rec f (x ': xs))"
instance Eq (Rec f '[]) where RNil == RNil = True
instance Ord (Rec f '[]) where compare RNil RNil = EQ
data Rec' (f :: u -> *) (us :: [u]) where
RNil' :: Rec' f '[]
RCons' :: f u -> Rec f us -> Rec' f (u ': us)
upRec :: Rec f us -> Rec' f us
upRec (Rec# v) | V.null v = cast# RNil'
| otherwise = cast# $ RCons' (cast# $ V.head v) (Rec# $ V.tail v)
pattern RNil :: () => (us ~ '[]) => Rec f us
pattern RNil <- (upRec -> RNil') where
RNil = Rec# V.empty
pattern (:&) :: () => (us' ~ (u ': us)) => f u -> Rec f us -> Rec f us'
pattern (:&) x xs <- (upRec -> RCons' x xs) where
x :& (Rec# xs) = Rec# (V.cons (cast# x) xs)
infixr 5 :&
thaw :: PrimMonad m => Rec f us -> m (MRec (PrimState m) f us)
thaw (Rec# v) = MRec# <$> V.thaw v
thaw# :: PrimMonad m => Rec f us -> m (MRec (PrimState m) f us)
thaw# (Rec# v) = MRec# <$> V.unsafeThaw v
freeze :: PrimMonad m => MRec (PrimState m) f us -> m (Rec f us)
freeze (MRec# v) = Rec# <$> V.freeze v
freeze# :: PrimMonad m => MRec (PrimState m) f us -> m (Rec f us)
freeze# (MRec# v) = Rec# <$> V.unsafeFreeze v
splitCons :: Rec f (x ': xs) -> (f x,Rec f xs)
splitCons (Rec# v) = (cast# $ V.head v, Rec# $ V.tail v)
(!) :: Rec f us -> Index us u -> f u
Rec# v ! Index# i = cast# $ v V.! i
rappend :: Rec f as -> Rec f bs -> Rec f (as ++ bs)
rappend (Rec# as) (Rec# bs) = Rec# $ as V.++ bs
rmap :: (forall x. f x -> g x) -> Rec f xs -> Rec g xs
rmap f r = runST $ freeze# =<< M.rmap f =<< thaw r
rmapM :: forall m f g xs. PrimMonad m => (forall x. f x -> m (g x)) -> Rec f xs -> m (Rec g xs)
rmapM f r = freeze# =<< M.rmapM f =<< thaw r
crmapM :: forall c m g f xs. (PrimMonad m, AllF c f xs)
=> (forall x. c (f x) => f x -> m (g x)) -> Rec f xs -> m (Rec g xs)
crmapM f r = freeze# =<< M.crmapM @c f =<< thaw r
crmap :: forall c g f xs. AllF c f xs
=> (forall x. c (f x) => f x -> g x) -> Rec f xs -> Rec g xs
crmap f r = runST $ freeze# =<< M.crmap @c f =<< thaw r
rzip :: forall h (f :: k -> *) g (xs :: [k]). (forall x. f x -> g x -> h x)
-> Rec f xs -> Rec g xs -> Rec h xs
rzip f fs gs = runST $ do
fs' <- thaw# fs
gs' <- thaw gs
freeze# =<< M.rzip f fs' gs'
crzip :: forall (c :: * -> Constraint) h f g (xs :: [Type])
. (AllF c f xs, AllF c g xs) => (forall x. (c (f x), c (g x)) => f x -> g x -> h x)
-> Rec f xs -> Rec g xs -> Rec h xs
crzip = crzip' @Type @c
crzip' :: forall (c :: * -> Constraint) h f g (xs :: [k])
. (AllF c f xs, AllF c g xs) => (forall x. (c (f x), c (g x)) => f x -> g x -> h x)
-> Rec f xs -> Rec g xs -> Rec h xs
crzip' f fs gs = runST $ do
fs' <- thaw fs
gs' <- thaw gs
freeze# =<< M.crzip @k @c f fs' gs'
instance Monoid (Rec f '[]) where
mempty = RNil
mappend _ _ = RNil
instance (KnownNat (Length (x ': xs)), AllF Monoid f ((x ': xs) :: [k])) => Monoid (Rec f (x ': xs)) where
mempty = new' @k @Monoid @(x ': xs) mempty
mappend = crzip' @k @Monoid mappend
instance Num (Rec f '[]) where
fromInteger _ = RNil
_ + _ = RNil
_ _ = RNil
_ * _ = RNil
negate x = x
abs x = x
signum x = x
instance (KnownNat (Length (x ': xs)), AllF Num f ((x ': xs) :: [k])) => Num (Rec f (x ': xs)) where
fromInteger n = new' @k @Num (fromInteger n)
(+) = crzip' @k @Num (+)
() = crzip' @k @Num ()
(*) = crzip' @k @Num (*)
negate = crmap @Num negate
abs = crmap @Num abs
signum = crmap @Num signum
instance Fractional (Rec f '[]) where
fromRational _ = RNil
_ / _ = RNil
recip x = x
instance (KnownNat (Length (x ': xs)), AllF Num f (x ': xs), AllF Fractional f ((x ': xs) :: [k])) => Fractional (Rec f (x ': xs)) where
fromRational n = new' @k @Fractional (fromRational n)
(/) = crzip' @k @Fractional (/)
recip = crmap @Fractional recip
instance Floating (Rec f '[]) where
pi = RNil
exp x = x
log x = x
sqrt x = x
sin x = x
cos x = x
tan x = x
asin x = x
acos x = x
atan x = x
sinh x = x
cosh x = x
tanh x = x
asinh x = x
acosh x = x
atanh x = x
log1p x = x
expm1 x = x
log1pexp x = x
log1mexp x = x
instance (KnownNat (Length (x ': xs)), Fractional (Rec f ((x ': xs) :: [k])), AllF Floating f (x ': xs)) => Floating (Rec f (x ': xs)) where
pi = new' @k @Floating pi
exp = crmap @Floating exp
log = crmap @Floating log
sqrt = crmap @Floating sqrt
sin = crmap @Floating sin
cos = crmap @Floating cos
tan = crmap @Floating tan
asin = crmap @Floating asin
acos = crmap @Floating acos
atan = crmap @Floating atan
sinh = crmap @Floating sinh
cosh = crmap @Floating cosh
tanh = crmap @Floating tanh
asinh = crmap @Floating asinh
acosh = crmap @Floating acosh
atanh = crmap @Floating atanh
log1p = crmap @Floating log1p
expm1 = crmap @Floating expm1
log1pexp = crmap @Floating log1pexp
log1mexp = crmap @Floating log1mexp
instance IsList (Rec f '[]) where
type Item (Rec f '[]) = Any
fromList = \case
[] -> RNil
_ -> error "fromList: nonempty list literal"
fromListN 0 [] = RNil
fromListN _ _ = error "fromListN: nonempty list literal"
toList _ = []
instance KnownNat (Length (x ': xs)) => IsList (Rec f (x ': xs)) where
type Item (Rec f (x ': xs)) = Any
fromListN n xs =
let n' = intVal @(Length (x ': xs))
in if n == n'
then Rec# $ fromListN n xs
else error $ "fromListN: expected length " ++ show n' ++ " but actually " ++ show n
fromList (fromList -> xs) =
let n' = intVal @(Length (x ': xs))
n = V.length xs
in if n == n'
then Rec# xs
else error $ "fromListN: expected length " ++ show n' ++ " but actually " ++ show n
toList = (`appEndo` []) . foldMapF (\x -> Endo (\xs -> cast# @Any x : xs))
toVector :: (forall x. f x -> r) -> Rec f xs -> Vector r
toVector f r = runST $ V.unsafeFreeze =<< M.toMVector f =<< thaw r
foldMapF :: forall r f xs. Monoid r => (forall x. f x -> r) -> Rec f xs -> r
foldMapF f = go where
go :: forall as. Rec f as -> r
go = \case
RNil -> mempty
x :& xs -> f x `mappend` go xs
cfoldMapF :: forall (c :: * -> Constraint) r f xs. (AllF c f xs, Monoid r)
=> (forall x. c (f x) => f x -> r) -> Rec f xs -> r
cfoldMapF f = go where
go :: forall as. AllF c f as => Rec f as -> r
go = \case
RNil -> mempty
x :& xs -> f x `mappend` go xs
ctoVector :: forall c r f xs. AllF c f xs
=> (forall x. c (f x) => f x -> r) -> Rec f xs -> Vector r
ctoVector f r = runST $ V.unsafeFreeze =<< M.ctoMVector @c f =<< thaw r
replicate :: forall n f x. KnownNat n => f x -> Rec f (Replicate n x)
replicate = Rec# . mapCast# @Any . V.replicate (fromInteger $ natVal (Proxy @n))
new :: forall (c :: * -> Constraint) (xs :: [*]) f. (AllF c f xs, KnownNat (Length xs))
=> (forall x. c (f x) => f x) -> Rec f xs
new = new' @Type @c
new' :: forall (c :: * -> Constraint) (xs :: [k]) f. (AllF c f xs, KnownNat (Length xs))
=> (forall x. c (f x) => f x) -> Rec f xs
new' x = runST $ do
xs <- M.new# @f @xs
xs' <- M.crmap @c (\_ -> x) xs
freeze# xs'
fromVectorN :: forall n f x. KnownNat n => Vector (f x) -> Maybe (Rec f (Replicate n x))
fromVectorN v =
let
n = intVal @n
in
if n <= V.length v then Just . Rec# . mapCast# @Any $ V.take n v else Nothing
fromVector :: forall f (x :: u). Vector (f x) -> Some (Rec f)
fromVector v = case someNatVal (fromIntegral $ V.length v) of
Nothing -> error "fromVector: impossible! Negative vector length"
Just (SomeNat (Proxy :: Proxy n))
-> Some $ Rec# @u @f @(Replicate n x) $ mapCast# @Any v
modify :: forall n f xs y. KnownNat n
=> (f (xs !! n) -> f y) -> Rec f xs -> Rec f (SetAt n xs y)
modify f r = runST $ do
mr <- thaw r
mr' <- M.modify @n f mr
freeze# mr'
index :: forall n f xs. KnownNat n => Rec f xs -> f (xs !! n)
index (Rec# v) = cast# $ v V.! intVal @n
checkDynIndex :: forall (xs :: [u]) f. Rec f xs -> Int -> MaybeSome (Index xs)
checkDynIndex RNil _ = None
checkDynIndex ((_::f x) :& _) 0 = JustSome (RZ @u @x)
checkDynIndex (_ :& xs) n = case checkDynIndex xs (n1) of
None -> None
JustSome i -> JustSome $ RS i
subRecord# :: forall ns f xs. (KnownNat (Length ns), KnownNats ns)
=> Rec f xs -> Rec f (SubList# ns xs)
subRecord# (Rec# v) = Rec# $ v `V.backpermute` fromListN n ns
where ns = intList @ns
n = intVal @(Length ns)