{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ConstraintKinds            #-}
module Basement.Sized.UVect
    ( UVect
    , MUVect
    , unUVect
    , toUVect
    , empty
    , singleton
    , replicate
    , thaw
    , freeze
    , index
    , map
    , foldl'
    , foldr
    , cons
    , snoc
    , elem
    , sub
    , uncons
    , unsnoc
    , splitAt
    , all
    , any
    , find
    , reverse
    , sortBy
    , intersperse
    ) where

import           Basement.Compat.Base
import           Basement.Nat
import           Basement.NormalForm
import           Basement.Types.OffsetSize
import           Basement.Monad
import           Basement.PrimType (PrimType)
import qualified Basement.UArray as A
import qualified Basement.UArray.Mutable as A hiding (sub)
import           Data.Proxy

newtype UVect (n :: Nat) a = UVect { UVect n a -> UArray a
unUVect :: A.UArray a } deriving (UVect n a -> ()
(UVect n a -> ()) -> NormalForm (UVect n a)
forall a. (a -> ()) -> NormalForm a
forall (n :: Nat) a. UVect n a -> ()
toNormalForm :: UVect n a -> ()
$ctoNormalForm :: forall (n :: Nat) a. UVect n a -> ()
NormalForm, UVect n a -> UVect n a -> Bool
(UVect n a -> UVect n a -> Bool)
-> (UVect n a -> UVect n a -> Bool) -> Eq (UVect n a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (n :: Nat) a. PrimType a => UVect n a -> UVect n a -> Bool
/= :: UVect n a -> UVect n a -> Bool
$c/= :: forall (n :: Nat) a. PrimType a => UVect n a -> UVect n a -> Bool
== :: UVect n a -> UVect n a -> Bool
$c== :: forall (n :: Nat) a. PrimType a => UVect n a -> UVect n a -> Bool
Eq, Int -> UVect n a -> ShowS
[UVect n a] -> ShowS
UVect n a -> String
(Int -> UVect n a -> ShowS)
-> (UVect n a -> String)
-> ([UVect n a] -> ShowS)
-> Show (UVect n a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (n :: Nat) a.
(PrimType a, Show a) =>
Int -> UVect n a -> ShowS
forall (n :: Nat) a. (PrimType a, Show a) => [UVect n a] -> ShowS
forall (n :: Nat) a. (PrimType a, Show a) => UVect n a -> String
showList :: [UVect n a] -> ShowS
$cshowList :: forall (n :: Nat) a. (PrimType a, Show a) => [UVect n a] -> ShowS
show :: UVect n a -> String
$cshow :: forall (n :: Nat) a. (PrimType a, Show a) => UVect n a -> String
showsPrec :: Int -> UVect n a -> ShowS
$cshowsPrec :: forall (n :: Nat) a.
(PrimType a, Show a) =>
Int -> UVect n a -> ShowS
Show)
newtype MUVect (n :: Nat) ty st = MUVect { MUVect n ty st -> MUArray ty st
unMUVect :: A.MUArray ty st }

toUVect :: forall n ty . (PrimType ty, KnownNat n, Countable ty n) => A.UArray ty -> Maybe (UVect n ty)
toUVect :: UArray ty -> Maybe (UVect n ty)
toUVect UArray ty
b
    | CountOf ty
expected CountOf ty -> CountOf ty -> Bool
forall a. Eq a => a -> a -> Bool
== UArray ty -> CountOf ty
forall ty. UArray ty -> CountOf ty
A.length UArray ty
b = UVect n ty -> Maybe (UVect n ty)
forall a. a -> Maybe a
Just (UArray ty -> UVect n ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect UArray ty
b)
    | Bool
otherwise              = Maybe (UVect n ty)
forall a. Maybe a
Nothing
  where
    expected :: CountOf ty
expected = forall ty. (KnownNat n, Countable ty n) => CountOf ty
forall (n :: Nat) ty. (KnownNat n, Countable ty n) => CountOf ty
toCount @n

empty :: PrimType ty => UVect 0 ty
empty :: UVect 0 ty
empty = UArray ty -> UVect 0 ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect UArray ty
forall a. Monoid a => a
mempty

singleton :: PrimType ty => ty -> UVect 1 ty
singleton :: ty -> UVect 1 ty
singleton ty
a = UArray ty -> UVect 1 ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (ty -> UArray ty
forall ty. PrimType ty => ty -> UArray ty
A.singleton ty
a)

create :: forall ty (n :: Nat) . (PrimType ty, Countable ty n, KnownNat n) => (Offset ty -> ty) -> UVect n ty
create :: (Offset ty -> ty) -> UVect n ty
create Offset ty -> ty
f = UArray ty -> UVect n ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> UVect n ty) -> UArray ty -> UVect n ty
forall a b. (a -> b) -> a -> b
$ CountOf ty -> (Offset ty -> ty) -> UArray ty
forall ty.
PrimType ty =>
CountOf ty -> (Offset ty -> ty) -> UArray ty
A.create CountOf ty
sz Offset ty -> ty
f
  where
    sz :: CountOf ty
sz = Proxy n -> CountOf ty
forall (n :: Nat) ty (proxy :: Nat -> *).
(KnownNat n, NatWithinBound (CountOf ty) n) =>
proxy n -> CountOf ty
natValCountOf (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)

replicate :: forall n ty . (KnownNat n, Countable ty n, PrimType ty) => ty -> UVect n ty
replicate :: ty -> UVect n ty
replicate ty
a = UArray ty -> UVect n ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (CountOf ty -> ty -> UArray ty
forall ty. PrimType ty => CountOf ty -> ty -> UArray ty
A.replicate (forall ty. (KnownNat n, Countable ty n) => CountOf ty
forall (n :: Nat) ty. (KnownNat n, Countable ty n) => CountOf ty
toCount @n) ty
a)

thaw :: (KnownNat n, PrimMonad prim, PrimType ty) => UVect n ty -> prim (MUVect n ty (PrimState prim))
thaw :: UVect n ty -> prim (MUVect n ty (PrimState prim))
thaw UVect n ty
b = MUArray ty (PrimState prim) -> MUVect n ty (PrimState prim)
forall (n :: Nat) ty st. MUArray ty st -> MUVect n ty st
MUVect (MUArray ty (PrimState prim) -> MUVect n ty (PrimState prim))
-> prim (MUArray ty (PrimState prim))
-> prim (MUVect n ty (PrimState prim))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> UArray ty -> prim (MUArray ty (PrimState prim))
forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
UArray ty -> prim (MUArray ty (PrimState prim))
A.thaw (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

freeze ::  (PrimMonad prim, PrimType ty, Countable ty n) => MUVect n ty (PrimState prim) -> prim (UVect n ty)
freeze :: MUVect n ty (PrimState prim) -> prim (UVect n ty)
freeze MUVect n ty (PrimState prim)
b = UArray ty -> UVect n ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> UVect n ty) -> prim (UArray ty) -> prim (UVect n ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MUArray ty (PrimState prim) -> prim (UArray ty)
forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
MUArray ty (PrimState prim) -> prim (UArray ty)
A.freeze (MUVect n ty (PrimState prim) -> MUArray ty (PrimState prim)
forall (n :: Nat) ty st. MUVect n ty st -> MUArray ty st
unMUVect MUVect n ty (PrimState prim)
b)

write :: (PrimMonad prim, PrimType ty) => MUVect n ty (PrimState prim) -> Offset ty -> ty -> prim ()
write :: MUVect n ty (PrimState prim) -> Offset ty -> ty -> prim ()
write (MUVect MUArray ty (PrimState prim)
ma) Offset ty
ofs ty
v = MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
A.write MUArray ty (PrimState prim)
ma Offset ty
ofs ty
v

read :: (PrimMonad prim, PrimType ty) => MUVect n ty (PrimState prim) -> Offset ty -> prim ty
read :: MUVect n ty (PrimState prim) -> Offset ty -> prim ty
read (MUVect MUArray ty (PrimState prim)
ma) Offset ty
ofs = MUArray ty (PrimState prim) -> Offset ty -> prim ty
forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> Offset ty -> prim ty
A.read MUArray ty (PrimState prim)
ma Offset ty
ofs

indexStatic :: forall i n ty . (KnownNat i, CmpNat i n ~ 'LT, PrimType ty, Offsetable ty i) => UVect n ty -> ty
indexStatic :: UVect n ty -> ty
indexStatic UVect n ty
b = UArray ty -> Offset ty -> ty
forall ty. PrimType ty => UArray ty -> Offset ty -> ty
A.unsafeIndex (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) (forall ty. (KnownNat i, Offsetable ty i) => Offset ty
forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @i)

index :: forall i n ty . PrimType ty => UVect n ty -> Offset ty -> ty
index :: UVect n ty -> Offset ty -> ty
index UVect n ty
b Offset ty
ofs = UArray ty -> Offset ty -> ty
forall ty. PrimType ty => UArray ty -> Offset ty -> ty
A.index (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) Offset ty
ofs

map :: (PrimType a, PrimType b) => (a -> b) -> UVect n a -> UVect n b
map :: (a -> b) -> UVect n a -> UVect n b
map a -> b
f UVect n a
b = UArray b -> UVect n b
forall (n :: Nat) a. UArray a -> UVect n a
UVect ((a -> b) -> UArray a -> UArray b
forall a b.
(PrimType a, PrimType b) =>
(a -> b) -> UArray a -> UArray b
A.map a -> b
f (UVect n a -> UArray a
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n a
b))

foldl' :: PrimType ty => (a -> ty -> a) -> a -> UVect n ty -> a
foldl' :: (a -> ty -> a) -> a -> UVect n ty -> a
foldl' a -> ty -> a
f a
acc UVect n ty
b = (a -> ty -> a) -> a -> UArray ty -> a
forall ty a. PrimType ty => (a -> ty -> a) -> a -> UArray ty -> a
A.foldl' a -> ty -> a
f a
acc (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

foldr :: PrimType ty => (ty -> a -> a) -> a -> UVect n ty -> a
foldr :: (ty -> a -> a) -> a -> UVect n ty -> a
foldr ty -> a -> a
f a
acc UVect n ty
b = (ty -> a -> a) -> a -> UArray ty -> a
forall ty a. PrimType ty => (ty -> a -> a) -> a -> UArray ty -> a
A.foldr ty -> a -> a
f a
acc (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

cons :: PrimType ty => ty -> UVect n ty -> UVect (n+1) ty
cons :: ty -> UVect n ty -> UVect (n + 1) ty
cons ty
e = UArray ty -> UVect (n + 1) ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> UVect (n + 1) ty)
-> (UVect n ty -> UArray ty) -> UVect n ty -> UVect (n + 1) ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ty -> UArray ty -> UArray ty
forall ty. PrimType ty => ty -> UArray ty -> UArray ty
A.cons ty
e (UArray ty -> UArray ty)
-> (UVect n ty -> UArray ty) -> UVect n ty -> UArray ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect

snoc :: PrimType ty => UVect n ty -> ty -> UVect (n+1) ty
snoc :: UVect n ty -> ty -> UVect (n + 1) ty
snoc UVect n ty
b = UArray ty -> UVect (n + 1) ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> UVect (n + 1) ty)
-> (ty -> UArray ty) -> ty -> UVect (n + 1) ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. UArray ty -> ty -> UArray ty
forall ty. PrimType ty => UArray ty -> ty -> UArray ty
A.snoc (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

sub :: forall i j n ty
     . ( (i <=? n) ~ 'True
       , (j <=? n) ~ 'True
       , (i <=? j) ~ 'True
       , PrimType ty
       , KnownNat i
       , KnownNat j
       , Offsetable ty i
       , Offsetable ty j )
    => UVect n ty
    -> UVect (j-i) ty
sub :: UVect n ty -> UVect (j - i) ty
sub UVect n ty
block = UArray ty -> UVect (j - i) ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> Offset ty -> Offset ty -> UArray ty
forall ty.
PrimType ty =>
UArray ty -> Offset ty -> Offset ty -> UArray ty
A.sub (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
block) (forall ty. (KnownNat i, Offsetable ty i) => Offset ty
forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @i) (forall ty. (KnownNat j, Offsetable ty j) => Offset ty
forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @j))

uncons :: forall n ty . (CmpNat 0 n ~ 'LT, PrimType ty, KnownNat n, Offsetable ty n)
       => UVect n ty
       -> (ty, UVect (n-1) ty)
uncons :: UVect n ty -> (ty, UVect (n - 1) ty)
uncons UVect n ty
b = (UVect n ty -> ty
forall (i :: Nat) (n :: Nat) ty.
(KnownNat i, CmpNat i n ~ 'LT, PrimType ty, Offsetable ty i) =>
UVect n ty -> ty
indexStatic @0 UVect n ty
b, UArray ty -> UVect (n - 1) ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> Offset ty -> Offset ty -> UArray ty
forall ty.
PrimType ty =>
UArray ty -> Offset ty -> Offset ty -> UArray ty
A.sub (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) Offset ty
1 (forall ty. (KnownNat n, Offsetable ty n) => Offset ty
forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @n)))

unsnoc :: forall n ty . (CmpNat 0 n ~ 'LT, KnownNat n, PrimType ty, Offsetable ty n)
       => UVect n ty
       -> (UVect (n-1) ty, ty)
unsnoc :: UVect n ty -> (UVect (n - 1) ty, ty)
unsnoc UVect n ty
b =
    ( UArray ty -> UVect (n - 1) ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> Offset ty -> Offset ty -> UArray ty
forall ty.
PrimType ty =>
UArray ty -> Offset ty -> Offset ty -> UArray ty
A.sub (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) Offset ty
0 (forall ty. (KnownNat n, Offsetable ty n) => Offset ty
forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @n Offset ty -> Offset ty -> Offset ty
forall a. Offset a -> Offset a -> Offset a
`offsetSub` Offset ty
1))
    , UArray ty -> Offset ty -> ty
forall ty. PrimType ty => UArray ty -> Offset ty -> ty
A.unsafeIndex (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) (forall ty. (KnownNat n, Offsetable ty n) => Offset ty
forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @n Offset ty -> Offset ty -> Offset ty
forall a. Offset a -> Offset a -> Offset a
`offsetSub` Offset ty
1))

splitAt :: forall i n ty . (CmpNat i n ~ 'LT, PrimType ty, KnownNat i, Countable ty i) => UVect n ty -> (UVect i ty, UVect (n-i) ty)
splitAt :: UVect n ty -> (UVect i ty, UVect (n - i) ty)
splitAt UVect n ty
b =
    let (UArray ty
left, UArray ty
right) = CountOf ty -> UArray ty -> (UArray ty, UArray ty)
forall ty. CountOf ty -> UArray ty -> (UArray ty, UArray ty)
A.splitAt (forall ty. (KnownNat i, Countable ty i) => CountOf ty
forall (n :: Nat) ty. (KnownNat n, Countable ty n) => CountOf ty
toCount @i) (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)
     in (UArray ty -> UVect i ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect UArray ty
left, UArray ty -> UVect (n - i) ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect UArray ty
right)

elem :: PrimType ty => ty -> UVect n ty -> Bool
elem :: ty -> UVect n ty -> Bool
elem ty
e UVect n ty
b = ty -> UArray ty -> Bool
forall ty. PrimType ty => ty -> UArray ty -> Bool
A.elem ty
e (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

all :: PrimType ty => (ty -> Bool) -> UVect n ty -> Bool
all :: (ty -> Bool) -> UVect n ty -> Bool
all ty -> Bool
p UVect n ty
b = (ty -> Bool) -> UArray ty -> Bool
forall ty. PrimType ty => (ty -> Bool) -> UArray ty -> Bool
A.all ty -> Bool
p (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

any :: PrimType ty => (ty -> Bool) -> UVect n ty -> Bool
any :: (ty -> Bool) -> UVect n ty -> Bool
any ty -> Bool
p UVect n ty
b = (ty -> Bool) -> UArray ty -> Bool
forall ty. PrimType ty => (ty -> Bool) -> UArray ty -> Bool
A.any ty -> Bool
p (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

find :: PrimType ty => (ty -> Bool) -> UVect n ty -> Maybe ty
find :: (ty -> Bool) -> UVect n ty -> Maybe ty
find ty -> Bool
p UVect n ty
b = (ty -> Bool) -> UArray ty -> Maybe ty
forall ty. PrimType ty => (ty -> Bool) -> UArray ty -> Maybe ty
A.find ty -> Bool
p (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

reverse :: PrimType ty => UVect n ty -> UVect n ty
reverse :: UVect n ty -> UVect n ty
reverse = UArray ty -> UVect n ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (UArray ty -> UVect n ty)
-> (UVect n ty -> UArray ty) -> UVect n ty -> UVect n ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. UArray ty -> UArray ty
forall ty. PrimType ty => UArray ty -> UArray ty
A.reverse (UArray ty -> UArray ty)
-> (UVect n ty -> UArray ty) -> UVect n ty -> UArray ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect

sortBy :: PrimType ty => (ty -> ty -> Ordering) -> UVect n ty -> UVect n ty
sortBy :: (ty -> ty -> Ordering) -> UVect n ty -> UVect n ty
sortBy ty -> ty -> Ordering
f UVect n ty
b = UArray ty -> UVect n ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect ((ty -> ty -> Ordering) -> UArray ty -> UArray ty
forall ty.
PrimType ty =>
(ty -> ty -> Ordering) -> UArray ty -> UArray ty
A.sortBy ty -> ty -> Ordering
f (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b))

intersperse :: (CmpNat n 1 ~ 'GT, PrimType ty) => ty -> UVect n ty -> UVect (n+n-1) ty
intersperse :: ty -> UVect n ty -> UVect ((n + n) - 1) ty
intersperse ty
sep UVect n ty
b = UArray ty -> UVect ((n + n) - 1) ty
forall (n :: Nat) a. UArray a -> UVect n a
UVect (ty -> UArray ty -> UArray ty
forall ty. PrimType ty => ty -> UArray ty -> UArray ty
A.intersperse ty
sep (UVect n ty -> UArray ty
forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b))

toCount :: forall n ty . (KnownNat n, Countable ty n) => CountOf ty
toCount :: CountOf ty
toCount = Proxy n -> CountOf ty
forall (n :: Nat) ty (proxy :: Nat -> *).
(KnownNat n, NatWithinBound (CountOf ty) n) =>
proxy n -> CountOf ty
natValCountOf (Proxy n
forall k (t :: k). Proxy t
Proxy @n)

toOffset :: forall n ty . (KnownNat n, Offsetable ty n) => Offset ty
toOffset :: Offset ty
toOffset = Proxy n -> Offset ty
forall (n :: Nat) ty (proxy :: Nat -> *).
(KnownNat n, NatWithinBound (Offset ty) n) =>
proxy n -> Offset ty
natValOffset (Proxy n
forall k (t :: k). Proxy t
Proxy @n)