{-# 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 { forall (n :: Nat) a. UVect n a -> UArray a
unUVect :: A.UArray a } deriving (UVect n a -> ()
forall (n :: Nat) a. UVect n a -> ()
forall a. (a -> ()) -> NormalForm a
toNormalForm :: UVect n a -> ()
$ctoNormalForm :: forall (n :: Nat) a. UVect n a -> ()
NormalForm, UVect n a -> UVect n a -> Bool
forall (n :: Nat) a. PrimType a => UVect n a -> UVect n a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: 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
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
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
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 { forall (n :: Nat) ty st. 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 :: forall (n :: Nat) ty.
(PrimType ty, KnownNat n, Countable ty n) =>
UArray ty -> Maybe (UVect n ty)
toUVect UArray ty
b
    | CountOf ty
expected forall a. Eq a => a -> a -> Bool
== forall ty. UArray ty -> CountOf ty
A.length UArray ty
b = forall a. a -> Maybe a
Just (forall (n :: Nat) a. UArray a -> UVect n a
UVect UArray ty
b)
    | Bool
otherwise              = forall a. Maybe a
Nothing
  where
    expected :: CountOf ty
expected = forall (n :: Nat) ty. (KnownNat n, Countable ty n) => CountOf ty
toCount @n

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

singleton :: PrimType ty => ty -> UVect 1 ty
singleton :: forall ty. PrimType ty => ty -> UVect 1 ty
singleton ty
a = forall (n :: Nat) a. UArray a -> UVect n a
UVect (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 :: forall ty (n :: Nat).
(PrimType ty, Countable ty n, KnownNat n) =>
(Offset ty -> ty) -> UVect n ty
create Offset ty -> ty
f = forall (n :: Nat) a. UArray a -> UVect n a
UVect forall a b. (a -> b) -> a -> b
$ 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 = forall (n :: Nat) ty (proxy :: Nat -> *).
(KnownNat n, NatWithinBound (CountOf ty) n) =>
proxy n -> CountOf ty
natValCountOf (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 :: forall (n :: Nat) ty.
(KnownNat n, Countable ty n, PrimType ty) =>
ty -> UVect n ty
replicate ty
a = forall (n :: Nat) a. UArray a -> UVect n a
UVect (forall ty. PrimType ty => CountOf ty -> ty -> UArray ty
A.replicate (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 :: forall (n :: Nat) (prim :: * -> *) ty.
(KnownNat n, PrimMonad prim, PrimType ty) =>
UVect n ty -> prim (MUVect n ty (PrimState prim))
thaw UVect n ty
b = forall (n :: Nat) ty st. MUArray ty st -> MUVect n ty st
MUVect forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
UArray ty -> prim (MUArray ty (PrimState prim))
A.thaw (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 :: forall (prim :: * -> *) ty (n :: Nat).
(PrimMonad prim, PrimType ty, Countable ty n) =>
MUVect n ty (PrimState prim) -> prim (UVect n ty)
freeze MUVect n ty (PrimState prim)
b = forall (n :: Nat) a. UArray a -> UVect n a
UVect forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
MUArray ty (PrimState prim) -> prim (UArray ty)
A.freeze (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 :: forall (prim :: * -> *) ty (n :: Nat).
(PrimMonad prim, PrimType ty) =>
MUVect n ty (PrimState prim) -> Offset ty -> ty -> prim ()
write (MUVect MUArray ty (PrimState prim)
ma) Offset ty
ofs ty
v = 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 :: forall (prim :: * -> *) ty (n :: Nat).
(PrimMonad prim, PrimType ty) =>
MUVect n ty (PrimState prim) -> Offset ty -> prim ty
read (MUVect MUArray ty (PrimState prim)
ma) Offset ty
ofs = 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 :: forall (i :: Nat) (n :: Nat) ty.
(KnownNat i, CmpNat i n ~ 'LT, PrimType ty, Offsetable ty i) =>
UVect n ty -> ty
indexStatic UVect n ty
b = forall ty. PrimType ty => UArray ty -> Offset ty -> ty
A.unsafeIndex (forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) (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 :: forall i (n :: Nat) ty.
PrimType ty =>
UVect n ty -> Offset ty -> ty
index UVect n ty
b Offset ty
ofs = forall ty. PrimType ty => UArray ty -> Offset ty -> ty
A.index (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 :: forall a b (n :: Nat).
(PrimType a, PrimType b) =>
(a -> b) -> UVect n a -> UVect n b
map a -> b
f UVect n a
b = forall (n :: Nat) a. UArray a -> UVect n a
UVect (forall a b.
(PrimType a, PrimType b) =>
(a -> b) -> UArray a -> UArray b
A.map a -> b
f (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' :: forall ty a (n :: Nat).
PrimType ty =>
(a -> ty -> a) -> a -> UVect n ty -> a
foldl' a -> ty -> a
f a
acc UVect n ty
b = forall ty a. PrimType ty => (a -> ty -> a) -> a -> UArray ty -> a
A.foldl' a -> ty -> a
f a
acc (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 :: forall ty a (n :: Nat).
PrimType ty =>
(ty -> a -> a) -> a -> UVect n ty -> a
foldr ty -> a -> a
f a
acc UVect n ty
b = forall ty a. PrimType ty => (ty -> a -> a) -> a -> UArray ty -> a
A.foldr ty -> a -> a
f a
acc (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 :: forall ty (n :: Nat).
PrimType ty =>
ty -> UVect n ty -> UVect (n + 1) ty
cons ty
e = forall (n :: Nat) a. UArray a -> UVect n a
UVect forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall ty. PrimType ty => ty -> UArray ty -> UArray ty
A.cons ty
e forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (n :: Nat) a. UVect n a -> UArray a
unUVect

snoc :: PrimType ty => UVect n ty -> ty -> UVect (n+1) ty
snoc :: forall ty (n :: Nat).
PrimType ty =>
UVect n ty -> ty -> UVect (n + 1) ty
snoc UVect n ty
b = forall (n :: Nat) a. UArray a -> UVect n a
UVect forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall ty. PrimType ty => UArray ty -> ty -> UArray ty
A.snoc (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 :: forall (i :: Nat) (j :: Nat) (n :: Nat) 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
block = forall (n :: Nat) a. UArray a -> UVect n a
UVect (forall ty.
PrimType ty =>
UArray ty -> Offset ty -> Offset ty -> UArray ty
A.sub (forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
block) (forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @i) (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 :: forall (n :: Nat) ty.
(CmpNat 0 n ~ 'LT, PrimType ty, KnownNat n, Offsetable ty n) =>
UVect n ty -> (ty, UVect (n - 1) ty)
uncons UVect n ty
b = (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, forall (n :: Nat) a. UArray a -> UVect n a
UVect (forall ty.
PrimType ty =>
UArray ty -> Offset ty -> Offset ty -> UArray ty
A.sub (forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) Offset ty
1 (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 :: forall (n :: Nat) ty.
(CmpNat 0 n ~ 'LT, KnownNat n, PrimType ty, Offsetable ty n) =>
UVect n ty -> (UVect (n - 1) ty, ty)
unsnoc UVect n ty
b =
    ( forall (n :: Nat) a. UArray a -> UVect n a
UVect (forall ty.
PrimType ty =>
UArray ty -> Offset ty -> Offset ty -> UArray ty
A.sub (forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) Offset ty
0 (forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @n forall a. Offset a -> Offset a -> Offset a
`offsetSub` Offset ty
1))
    , forall ty. PrimType ty => UArray ty -> Offset ty -> ty
A.unsafeIndex (forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b) (forall (n :: Nat) ty. (KnownNat n, Offsetable ty n) => Offset ty
toOffset @n 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 :: forall (i :: Nat) (n :: Nat) 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
b =
    let (UArray ty
left, UArray ty
right) = forall ty. CountOf ty -> UArray ty -> (UArray ty, UArray ty)
A.splitAt (forall (n :: Nat) ty. (KnownNat n, Countable ty n) => CountOf ty
toCount @i) (forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)
     in (forall (n :: Nat) a. UArray a -> UVect n a
UVect UArray ty
left, forall (n :: Nat) a. UArray a -> UVect n a
UVect UArray ty
right)

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

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

any :: PrimType ty => (ty -> Bool) -> UVect n ty -> Bool
any :: forall ty (n :: Nat).
PrimType ty =>
(ty -> Bool) -> UVect n ty -> Bool
any ty -> Bool
p UVect n ty
b = forall ty. PrimType ty => (ty -> Bool) -> UArray ty -> Bool
A.any ty -> Bool
p (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 :: forall ty (n :: Nat).
PrimType ty =>
(ty -> Bool) -> UVect n ty -> Maybe ty
find ty -> Bool
p UVect n ty
b = forall ty. PrimType ty => (ty -> Bool) -> UArray ty -> Maybe ty
A.find ty -> Bool
p (forall (n :: Nat) a. UVect n a -> UArray a
unUVect UVect n ty
b)

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

sortBy :: PrimType ty => (ty -> ty -> Ordering) -> UVect n ty -> UVect n ty
sortBy :: forall ty (n :: Nat).
PrimType ty =>
(ty -> ty -> Ordering) -> UVect n ty -> UVect n ty
sortBy ty -> ty -> Ordering
f UVect n ty
b = forall (n :: Nat) a. UArray a -> UVect n a
UVect (forall ty.
PrimType ty =>
(ty -> ty -> Ordering) -> UArray ty -> UArray ty
A.sortBy ty -> ty -> Ordering
f (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 :: forall (n :: Nat) ty.
(CmpNat n 1 ~ 'GT, PrimType ty) =>
ty -> UVect n ty -> UVect ((n + n) - 1) ty
intersperse ty
sep UVect n ty
b = forall (n :: Nat) a. UArray a -> UVect n a
UVect (forall ty. PrimType ty => ty -> UArray ty -> UArray ty
A.intersperse ty
sep (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 :: forall (n :: Nat) ty. (KnownNat n, Countable ty n) => CountOf ty
toCount = forall (n :: Nat) ty (proxy :: Nat -> *).
(KnownNat n, NatWithinBound (CountOf ty) n) =>
proxy n -> CountOf ty
natValCountOf (forall {k} (t :: k). Proxy t
Proxy @n)

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