{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ConstraintKinds            #-}
module Basement.Sized.Vect
    ( Vect
    , MVect
    , unVect
    , toVect
    , 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.BoxedArray as A
--import qualified Basement.BoxedArray.Mutable as A hiding (sub)
import           Data.Proxy

newtype Vect (n :: Nat) a = Vect { Vect n a -> Array a
unVect :: A.Array a } deriving (Vect n a -> ()
(Vect n a -> ()) -> NormalForm (Vect n a)
forall a. (a -> ()) -> NormalForm a
forall (n :: Nat) a. NormalForm a => Vect n a -> ()
toNormalForm :: Vect n a -> ()
$ctoNormalForm :: forall (n :: Nat) a. NormalForm a => Vect n a -> ()
NormalForm, Vect n a -> Vect n a -> Bool
(Vect n a -> Vect n a -> Bool)
-> (Vect n a -> Vect n a -> Bool) -> Eq (Vect n a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (n :: Nat) a. Eq a => Vect n a -> Vect n a -> Bool
/= :: Vect n a -> Vect n a -> Bool
$c/= :: forall (n :: Nat) a. Eq a => Vect n a -> Vect n a -> Bool
== :: Vect n a -> Vect n a -> Bool
$c== :: forall (n :: Nat) a. Eq a => Vect n a -> Vect n a -> Bool
Eq, Int -> Vect n a -> ShowS
[Vect n a] -> ShowS
Vect n a -> String
(Int -> Vect n a -> ShowS)
-> (Vect n a -> String) -> ([Vect n a] -> ShowS) -> Show (Vect n a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (n :: Nat) a. Show a => Int -> Vect n a -> ShowS
forall (n :: Nat) a. Show a => [Vect n a] -> ShowS
forall (n :: Nat) a. Show a => Vect n a -> String
showList :: [Vect n a] -> ShowS
$cshowList :: forall (n :: Nat) a. Show a => [Vect n a] -> ShowS
show :: Vect n a -> String
$cshow :: forall (n :: Nat) a. Show a => Vect n a -> String
showsPrec :: Int -> Vect n a -> ShowS
$cshowsPrec :: forall (n :: Nat) a. Show a => Int -> Vect n a -> ShowS
Show)
newtype MVect (n :: Nat) ty st = MVect { MVect n ty st -> MArray ty st
unMVect :: A.MArray ty st }

instance Functor (Vect n) where
    fmap :: (a -> b) -> Vect n a -> Vect n b
fmap = (a -> b) -> Vect n a -> Vect n b
forall a b (n :: Nat). (a -> b) -> Vect n a -> Vect n b
map

toVect :: forall n ty . (KnownNat n, Countable ty n) => A.Array ty -> Maybe (Vect n ty)
toVect :: Array ty -> Maybe (Vect n ty)
toVect Array ty
b
    | CountOf ty
expected CountOf ty -> CountOf ty -> Bool
forall a. Eq a => a -> a -> Bool
== Array ty -> CountOf ty
forall a. Array a -> CountOf a
A.length Array ty
b = Vect n ty -> Maybe (Vect n ty)
forall a. a -> Maybe a
Just (Array ty -> Vect n ty
forall (n :: Nat) a. Array a -> Vect n a
Vect Array ty
b)
    | Bool
otherwise = Maybe (Vect 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 :: Vect 0 ty
empty :: Vect 0 ty
empty = Array ty -> Vect 0 ty
forall (n :: Nat) a. Array a -> Vect n a
Vect Array ty
forall a. Array a
A.empty

singleton :: ty -> Vect 1 ty
singleton :: ty -> Vect 1 ty
singleton ty
a = Array ty -> Vect 1 ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (ty -> Array ty
forall ty. ty -> Array ty
A.singleton ty
a)

create :: forall a (n :: Nat) . (Countable a n, KnownNat n) => (Offset a -> a) -> Vect n a
create :: (Offset a -> a) -> Vect n a
create Offset a -> a
f = Array a -> Vect n a
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array a -> Vect n a) -> Array a -> Vect n a
forall a b. (a -> b) -> a -> b
$ CountOf a -> (Offset a -> a) -> Array a
forall ty. CountOf ty -> (Offset ty -> ty) -> Array ty
A.create CountOf a
sz Offset a -> a
f
  where
    sz :: CountOf a
sz = Proxy n -> CountOf a
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) => ty -> Vect n ty
replicate :: ty -> Vect n ty
replicate ty
a = Array ty -> Vect n ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (CountOf ty -> ty -> Array ty
forall ty. CountOf ty -> ty -> Array 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) => Vect n ty -> prim (MVect n ty (PrimState prim))
thaw :: Vect n ty -> prim (MVect n ty (PrimState prim))
thaw Vect n ty
b = MArray ty (PrimState prim) -> MVect n ty (PrimState prim)
forall (n :: Nat) ty st. MArray ty st -> MVect n ty st
MVect (MArray ty (PrimState prim) -> MVect n ty (PrimState prim))
-> prim (MArray ty (PrimState prim))
-> prim (MVect n ty (PrimState prim))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array ty -> prim (MArray ty (PrimState prim))
forall (prim :: * -> *) ty.
PrimMonad prim =>
Array ty -> prim (MArray ty (PrimState prim))
A.thaw (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

freeze ::  (PrimMonad prim, Countable ty n) => MVect n ty (PrimState prim) -> prim (Vect n ty)
freeze :: MVect n ty (PrimState prim) -> prim (Vect n ty)
freeze MVect n ty (PrimState prim)
b = Array ty -> Vect n ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array ty -> Vect n ty) -> prim (Array ty) -> prim (Vect n ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray ty (PrimState prim) -> prim (Array ty)
forall (prim :: * -> *) ty.
PrimMonad prim =>
MArray ty (PrimState prim) -> prim (Array ty)
A.freeze (MVect n ty (PrimState prim) -> MArray ty (PrimState prim)
forall (n :: Nat) ty st. MVect n ty st -> MArray ty st
unMVect MVect n ty (PrimState prim)
b)

write :: PrimMonad prim => MVect n ty (PrimState prim) -> Offset ty -> ty -> prim ()
write :: MVect n ty (PrimState prim) -> Offset ty -> ty -> prim ()
write (MVect MArray ty (PrimState prim)
ma) Offset ty
ofs ty
v = MArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
forall (prim :: * -> *) ty.
PrimMonad prim =>
MArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
A.write MArray ty (PrimState prim)
ma Offset ty
ofs ty
v

read :: PrimMonad prim => MVect n ty (PrimState prim) -> Offset ty -> prim ty
read :: MVect n ty (PrimState prim) -> Offset ty -> prim ty
read (MVect MArray ty (PrimState prim)
ma) Offset ty
ofs = MArray ty (PrimState prim) -> Offset ty -> prim ty
forall (prim :: * -> *) ty.
PrimMonad prim =>
MArray ty (PrimState prim) -> Offset ty -> prim ty
A.read MArray ty (PrimState prim)
ma Offset ty
ofs

indexStatic :: forall i n ty . (KnownNat i, CmpNat i n ~ 'LT, Offsetable ty i) => Vect n ty -> ty
indexStatic :: Vect n ty -> ty
indexStatic Vect n ty
b = Array ty -> Offset ty -> ty
forall ty. Array ty -> Offset ty -> ty
A.unsafeIndex (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect 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 :: Vect n ty -> Offset ty -> ty
index :: Vect n ty -> Offset ty -> ty
index Vect n ty
b Offset ty
ofs = Array ty -> Offset ty -> ty
forall ty. Array ty -> Offset ty -> ty
A.index (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b) Offset ty
ofs

map :: (a -> b) -> Vect n a -> Vect n b
map :: (a -> b) -> Vect n a -> Vect n b
map a -> b
f Vect n a
b = Array b -> Vect n b
forall (n :: Nat) a. Array a -> Vect n a
Vect ((a -> b) -> Array a -> Array b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (Vect n a -> Array a
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n a
b))

foldl' :: (a -> ty -> a) -> a -> Vect n ty -> a
foldl' :: (a -> ty -> a) -> a -> Vect n ty -> a
foldl' a -> ty -> a
f a
acc Vect n ty
b = (a -> ty -> a) -> a -> Array ty -> a
forall a ty. (a -> ty -> a) -> a -> Array ty -> a
A.foldl' a -> ty -> a
f a
acc (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

foldr :: (ty -> a -> a) -> a -> Vect n ty -> a
foldr :: (ty -> a -> a) -> a -> Vect n ty -> a
foldr ty -> a -> a
f a
acc Vect n ty
b = (ty -> a -> a) -> a -> Array ty -> a
forall ty a. (ty -> a -> a) -> a -> Array ty -> a
A.foldr ty -> a -> a
f a
acc (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

cons :: ty -> Vect n ty -> Vect (n+1) ty
cons :: ty -> Vect n ty -> Vect (n + 1) ty
cons ty
e = Array ty -> Vect (n + 1) ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array ty -> Vect (n + 1) ty)
-> (Vect n ty -> Array ty) -> Vect n ty -> Vect (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 -> Array ty -> Array ty
forall ty. ty -> Array ty -> Array ty
A.cons ty
e (Array ty -> Array ty)
-> (Vect n ty -> Array ty) -> Vect n ty -> Array ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect

snoc :: Vect n ty -> ty -> Vect (n+1) ty
snoc :: Vect n ty -> ty -> Vect (n + 1) ty
snoc Vect n ty
b = Array ty -> Vect (n + 1) ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array ty -> Vect (n + 1) ty)
-> (ty -> Array ty) -> ty -> Vect (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
. Array ty -> ty -> Array ty
forall ty. Array ty -> ty -> Array ty
A.snoc (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

sub :: forall i j n ty
     . ( (i <=? n) ~ 'True
       , (j <=? n) ~ 'True
       , (i <=? j) ~ 'True
       , KnownNat i
       , KnownNat j
       , Offsetable ty i
       , Offsetable ty j )
    => Vect n ty
    -> Vect (j-i) ty
sub :: Vect n ty -> Vect (j - i) ty
sub Vect n ty
block = Array ty -> Vect (j - i) ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array ty -> Offset ty -> Offset ty -> Array ty
forall ty. Array ty -> Offset ty -> Offset ty -> Array ty
A.sub (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect 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, KnownNat n, Offsetable ty n)
       => Vect n ty
       -> (ty, Vect (n-1) ty)
uncons :: Vect n ty -> (ty, Vect (n - 1) ty)
uncons Vect n ty
b = (Vect n ty -> ty
forall (i :: Nat) (n :: Nat) ty.
(KnownNat i, CmpNat i n ~ 'LT, Offsetable ty i) =>
Vect n ty -> ty
indexStatic @0 Vect n ty
b, Array ty -> Vect (n - 1) ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array ty -> Offset ty -> Offset ty -> Array ty
forall ty. Array ty -> Offset ty -> Offset ty -> Array ty
A.sub (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect 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, Offsetable ty n)
       => Vect n ty
       -> (Vect (n-1) ty, ty)
unsnoc :: Vect n ty -> (Vect (n - 1) ty, ty)
unsnoc Vect n ty
b =
    ( Array ty -> Vect (n - 1) ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array ty -> Offset ty -> Offset ty -> Array ty
forall ty. Array ty -> Offset ty -> Offset ty -> Array ty
A.sub (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect 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))
    , Array ty -> Offset ty -> ty
forall ty. Array ty -> Offset ty -> ty
A.unsafeIndex (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect 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, KnownNat i, Countable ty i) => Vect n ty -> (Vect i ty, Vect (n-i) ty)
splitAt :: Vect n ty -> (Vect i ty, Vect (n - i) ty)
splitAt Vect n ty
b =
    let (Array ty
left, Array ty
right) = CountOf ty -> Array ty -> (Array ty, Array ty)
forall ty. CountOf ty -> Array ty -> (Array ty, Array 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) (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)
     in (Array ty -> Vect i ty
forall (n :: Nat) a. Array a -> Vect n a
Vect Array ty
left, Array ty -> Vect (n - i) ty
forall (n :: Nat) a. Array a -> Vect n a
Vect Array ty
right)

elem :: Eq ty => ty -> Vect n ty -> Bool
elem :: ty -> Vect n ty -> Bool
elem ty
e Vect n ty
b = ty -> Array ty -> Bool
forall ty. Eq ty => ty -> Array ty -> Bool
A.elem ty
e (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

all :: (ty -> Bool) -> Vect n ty -> Bool
all :: (ty -> Bool) -> Vect n ty -> Bool
all ty -> Bool
p Vect n ty
b = (ty -> Bool) -> Array ty -> Bool
forall ty. (ty -> Bool) -> Array ty -> Bool
A.all ty -> Bool
p (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

any :: (ty -> Bool) -> Vect n ty -> Bool
any :: (ty -> Bool) -> Vect n ty -> Bool
any ty -> Bool
p Vect n ty
b = (ty -> Bool) -> Array ty -> Bool
forall ty. (ty -> Bool) -> Array ty -> Bool
A.any ty -> Bool
p (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

find :: (ty -> Bool) -> Vect n ty -> Maybe ty
find :: (ty -> Bool) -> Vect n ty -> Maybe ty
find ty -> Bool
p Vect n ty
b = (ty -> Bool) -> Array ty -> Maybe ty
forall ty. (ty -> Bool) -> Array ty -> Maybe ty
A.find ty -> Bool
p (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b)

reverse :: Vect n ty -> Vect n ty
reverse :: Vect n ty -> Vect n ty
reverse = Array ty -> Vect n ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (Array ty -> Vect n ty)
-> (Vect n ty -> Array ty) -> Vect n ty -> Vect n ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Array ty -> Array ty
forall ty. Array ty -> Array ty
A.reverse (Array ty -> Array ty)
-> (Vect n ty -> Array ty) -> Vect n ty -> Array ty
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect

sortBy :: (ty -> ty -> Ordering) -> Vect n ty -> Vect n ty
sortBy :: (ty -> ty -> Ordering) -> Vect n ty -> Vect n ty
sortBy ty -> ty -> Ordering
f Vect n ty
b = Array ty -> Vect n ty
forall (n :: Nat) a. Array a -> Vect n a
Vect ((ty -> ty -> Ordering) -> Array ty -> Array ty
forall ty. (ty -> ty -> Ordering) -> Array ty -> Array ty
A.sortBy ty -> ty -> Ordering
f (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect n ty
b))

intersperse :: (CmpNat n 1 ~ 'GT) => ty -> Vect n ty -> Vect (n+n-1) ty
intersperse :: ty -> Vect n ty -> Vect ((n + n) - 1) ty
intersperse ty
sep Vect n ty
b = Array ty -> Vect ((n + n) - 1) ty
forall (n :: Nat) a. Array a -> Vect n a
Vect (ty -> Array ty -> Array ty
forall ty. ty -> Array ty -> Array ty
A.intersperse ty
sep (Vect n ty -> Array ty
forall (n :: Nat) a. Vect n a -> Array a
unVect Vect 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)