module Dahdit.LiftedPrimArray
  ( LiftedPrimArray (..)
  , MutableLiftedPrimArray (..)
  , emptyLiftedPrimArray
  , indexLiftedPrimArray
  , writeLiftedPrimArray
  , freezeLiftedPrimArray
  , thawLiftedPrimArray
  , unsafeFreezeLiftedPrimArray
  , unsafeThawLiftedPrimArray
  , liftedPrimArrayFromListN
  , liftedPrimArrayFromList
  , generateLiftedPrimArray
  , sizeofLiftedPrimArray
  , lengthLiftedPrimArray
  , cloneLiftedPrimArray
  , replicateLiftedPrimArray
  )
where

import Control.Monad.Primitive (PrimMonad (..))
import Dahdit.LiftedPrim
  ( LiftedPrim (..)
  , indexArrayLiftedInElems
  , writeArrayLiftedInElems
  )
import Dahdit.Proxy (proxyFor, proxyForF)
import Dahdit.Sizes (ByteCount (..), ByteSized (..), ElemCount (..), StaticByteSized (..))
import Data.Coerce (coerce)
import Data.Default (Default (..))
import Data.Foldable (for_)
import Data.Primitive.ByteArray
  ( ByteArray
  , MutableByteArray
  , cloneByteArray
  , emptyByteArray
  , freezeByteArray
  , newByteArray
  , runByteArray
  , sizeofByteArray
  , thawByteArray
  , unsafeFreezeByteArray
  , unsafeThawByteArray
  )
import Data.Proxy (Proxy (..))
import Data.STRef (modifySTRef', newSTRef, readSTRef)

newtype LiftedPrimArray a = LiftedPrimArray {forall a. LiftedPrimArray a -> ByteArray
unLiftedPrimArray :: ByteArray}
  deriving stock (Int -> LiftedPrimArray a -> ShowS
forall a. Int -> LiftedPrimArray a -> ShowS
forall a. [LiftedPrimArray a] -> ShowS
forall a. LiftedPrimArray a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LiftedPrimArray a] -> ShowS
$cshowList :: forall a. [LiftedPrimArray a] -> ShowS
show :: LiftedPrimArray a -> String
$cshow :: forall a. LiftedPrimArray a -> String
showsPrec :: Int -> LiftedPrimArray a -> ShowS
$cshowsPrec :: forall a. Int -> LiftedPrimArray a -> ShowS
Show)
  deriving newtype (LiftedPrimArray a -> LiftedPrimArray a -> Bool
forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: LiftedPrimArray a -> LiftedPrimArray a -> Bool
$c/= :: forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
== :: LiftedPrimArray a -> LiftedPrimArray a -> Bool
$c== :: forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
Eq, LiftedPrimArray a -> LiftedPrimArray a -> Bool
LiftedPrimArray a -> LiftedPrimArray a -> Ordering
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
forall a. Eq (LiftedPrimArray a)
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
forall a. LiftedPrimArray a -> LiftedPrimArray a -> Ordering
forall a.
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
min :: LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
$cmin :: forall a.
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
max :: LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
$cmax :: forall a.
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
>= :: LiftedPrimArray a -> LiftedPrimArray a -> Bool
$c>= :: forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
> :: LiftedPrimArray a -> LiftedPrimArray a -> Bool
$c> :: forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
<= :: LiftedPrimArray a -> LiftedPrimArray a -> Bool
$c<= :: forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
< :: LiftedPrimArray a -> LiftedPrimArray a -> Bool
$c< :: forall a. LiftedPrimArray a -> LiftedPrimArray a -> Bool
compare :: LiftedPrimArray a -> LiftedPrimArray a -> Ordering
$ccompare :: forall a. LiftedPrimArray a -> LiftedPrimArray a -> Ordering
Ord, NonEmpty (LiftedPrimArray a) -> LiftedPrimArray a
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
forall b. Integral b => b -> LiftedPrimArray a -> LiftedPrimArray a
forall a. NonEmpty (LiftedPrimArray a) -> LiftedPrimArray a
forall a.
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
forall a b.
Integral b =>
b -> LiftedPrimArray a -> LiftedPrimArray a
stimes :: forall b. Integral b => b -> LiftedPrimArray a -> LiftedPrimArray a
$cstimes :: forall a b.
Integral b =>
b -> LiftedPrimArray a -> LiftedPrimArray a
sconcat :: NonEmpty (LiftedPrimArray a) -> LiftedPrimArray a
$csconcat :: forall a. NonEmpty (LiftedPrimArray a) -> LiftedPrimArray a
<> :: LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
$c<> :: forall a.
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
Semigroup, LiftedPrimArray a
[LiftedPrimArray a] -> LiftedPrimArray a
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
forall a. Semigroup (LiftedPrimArray a)
forall a. LiftedPrimArray a
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
forall a. [LiftedPrimArray a] -> LiftedPrimArray a
forall a.
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
mconcat :: [LiftedPrimArray a] -> LiftedPrimArray a
$cmconcat :: forall a. [LiftedPrimArray a] -> LiftedPrimArray a
mappend :: LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
$cmappend :: forall a.
LiftedPrimArray a -> LiftedPrimArray a -> LiftedPrimArray a
mempty :: LiftedPrimArray a
$cmempty :: forall a. LiftedPrimArray a
Monoid)

instance Default (LiftedPrimArray a) where
  def :: LiftedPrimArray a
def = forall a. LiftedPrimArray a
emptyLiftedPrimArray

instance ByteSized (LiftedPrimArray a) where
  byteSize :: LiftedPrimArray a -> ByteCount
byteSize = forall a. LiftedPrimArray a -> ByteCount
sizeofLiftedPrimArray

newtype MutableLiftedPrimArray m a = MutableLiftedPrimArray {forall m a. MutableLiftedPrimArray m a -> MutableByteArray m
unMutableLiftedPrimArray :: MutableByteArray m}
  deriving newtype (MutableLiftedPrimArray m a -> MutableLiftedPrimArray m a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall m a.
MutableLiftedPrimArray m a -> MutableLiftedPrimArray m a -> Bool
/= :: MutableLiftedPrimArray m a -> MutableLiftedPrimArray m a -> Bool
$c/= :: forall m a.
MutableLiftedPrimArray m a -> MutableLiftedPrimArray m a -> Bool
== :: MutableLiftedPrimArray m a -> MutableLiftedPrimArray m a -> Bool
$c== :: forall m a.
MutableLiftedPrimArray m a -> MutableLiftedPrimArray m a -> Bool
Eq)

emptyLiftedPrimArray :: LiftedPrimArray a
emptyLiftedPrimArray :: forall a. LiftedPrimArray a
emptyLiftedPrimArray = forall a. ByteArray -> LiftedPrimArray a
LiftedPrimArray ByteArray
emptyByteArray

indexLiftedPrimArray :: LiftedPrim a => LiftedPrimArray a -> ElemCount -> a
indexLiftedPrimArray :: forall a. LiftedPrim a => LiftedPrimArray a -> ElemCount -> a
indexLiftedPrimArray (LiftedPrimArray ByteArray
arr) = forall a. LiftedPrim a => Proxy a -> ByteArray -> ElemCount -> a
indexArrayLiftedInElems forall {k} (t :: k). Proxy t
Proxy ByteArray
arr

writeLiftedPrimArray :: (LiftedPrim a, PrimMonad m) => MutableLiftedPrimArray (PrimState m) a -> ElemCount -> a -> m ()
writeLiftedPrimArray :: forall a (m :: * -> *).
(LiftedPrim a, PrimMonad m) =>
MutableLiftedPrimArray (PrimState m) a -> ElemCount -> a -> m ()
writeLiftedPrimArray (MutableLiftedPrimArray MutableByteArray (PrimState m)
arr) = forall (m :: * -> *) a.
(PrimMonad m, LiftedPrim a) =>
MutableByteArray (PrimState m) -> ElemCount -> a -> m ()
writeArrayLiftedInElems MutableByteArray (PrimState m)
arr

freezeLiftedPrimArray :: PrimMonad m => MutableLiftedPrimArray (PrimState m) a -> ElemCount -> ElemCount -> m (LiftedPrimArray a)
freezeLiftedPrimArray :: forall (m :: * -> *) a.
PrimMonad m =>
MutableLiftedPrimArray (PrimState m) a
-> ElemCount -> ElemCount -> m (LiftedPrimArray a)
freezeLiftedPrimArray (MutableLiftedPrimArray MutableByteArray (PrimState m)
arr) ElemCount
off ElemCount
len = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. ByteArray -> LiftedPrimArray a
LiftedPrimArray (forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> Int -> m ByteArray
freezeByteArray MutableByteArray (PrimState m)
arr (coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
off) (coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
len))

unsafeFreezeLiftedPrimArray :: PrimMonad m => MutableLiftedPrimArray (PrimState m) a -> m (LiftedPrimArray a)
unsafeFreezeLiftedPrimArray :: forall (m :: * -> *) a.
PrimMonad m =>
MutableLiftedPrimArray (PrimState m) a -> m (LiftedPrimArray a)
unsafeFreezeLiftedPrimArray (MutableLiftedPrimArray MutableByteArray (PrimState m)
arr) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. ByteArray -> LiftedPrimArray a
LiftedPrimArray (forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray (PrimState m)
arr)

thawLiftedPrimArray :: PrimMonad m => LiftedPrimArray a -> ElemCount -> ElemCount -> m (MutableLiftedPrimArray (PrimState m) a)
thawLiftedPrimArray :: forall (m :: * -> *) a.
PrimMonad m =>
LiftedPrimArray a
-> ElemCount
-> ElemCount
-> m (MutableLiftedPrimArray (PrimState m) a)
thawLiftedPrimArray (LiftedPrimArray ByteArray
arr) ElemCount
off ElemCount
len = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall m a. MutableByteArray m -> MutableLiftedPrimArray m a
MutableLiftedPrimArray (forall (m :: * -> *).
PrimMonad m =>
ByteArray -> Int -> Int -> m (MutableByteArray (PrimState m))
thawByteArray ByteArray
arr (coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
off) (coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
len))

unsafeThawLiftedPrimArray :: PrimMonad m => LiftedPrimArray a -> m (MutableLiftedPrimArray (PrimState m) a)
unsafeThawLiftedPrimArray :: forall (m :: * -> *) a.
PrimMonad m =>
LiftedPrimArray a -> m (MutableLiftedPrimArray (PrimState m) a)
unsafeThawLiftedPrimArray (LiftedPrimArray ByteArray
arr) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall m a. MutableByteArray m -> MutableLiftedPrimArray m a
MutableLiftedPrimArray (forall (m :: * -> *).
PrimMonad m =>
ByteArray -> m (MutableByteArray (PrimState m))
unsafeThawByteArray ByteArray
arr)

liftedPrimArrayFromListN :: LiftedPrim a => ElemCount -> [a] -> LiftedPrimArray a
liftedPrimArrayFromListN :: forall a. LiftedPrim a => ElemCount -> [a] -> LiftedPrimArray a
liftedPrimArrayFromListN ElemCount
n [a]
xs = forall a. ByteArray -> LiftedPrimArray a
LiftedPrimArray forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutableByteArray s)) -> ByteArray
runByteArray forall a b. (a -> b) -> a -> b
$ do
  let elemSize :: ByteCount
elemSize = forall a. StaticByteSized a => Proxy a -> ByteCount
staticByteSize (forall (f :: * -> *) a. f a -> Proxy a
proxyForF [a]
xs)
      len :: Int
len = coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
n forall a. Num a => a -> a -> a
* coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
elemSize
  MutableByteArray s
arr <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
len
  STRef s ByteCount
offRef <- forall a s. a -> ST s (STRef s a)
newSTRef ByteCount
0
  forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [a]
xs forall a b. (a -> b) -> a -> b
$ \a
x -> do
    ByteCount
off <- forall s a. STRef s a -> ST s a
readSTRef STRef s ByteCount
offRef
    forall a (m :: * -> *).
(LiftedPrim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> ByteCount -> a -> m ()
writeArrayLiftedInBytes MutableByteArray s
arr ByteCount
off a
x
    forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s ByteCount
offRef (ByteCount
elemSize forall a. Num a => a -> a -> a
+)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure MutableByteArray s
arr

liftedPrimArrayFromList :: LiftedPrim a => [a] -> LiftedPrimArray a
liftedPrimArrayFromList :: forall a. LiftedPrim a => [a] -> LiftedPrimArray a
liftedPrimArrayFromList [a]
xs = forall a. LiftedPrim a => ElemCount -> [a] -> LiftedPrimArray a
liftedPrimArrayFromListN (coerce :: forall a b. Coercible a b => a -> b
coerce (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)) [a]
xs

generateLiftedPrimArray :: LiftedPrim a => ElemCount -> (ElemCount -> a) -> LiftedPrimArray a
generateLiftedPrimArray :: forall a.
LiftedPrim a =>
ElemCount -> (ElemCount -> a) -> LiftedPrimArray a
generateLiftedPrimArray ElemCount
n ElemCount -> a
f = forall a. LiftedPrim a => ElemCount -> [a] -> LiftedPrimArray a
liftedPrimArrayFromListN ElemCount
n (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ElemCount -> a
f [ElemCount
0 .. ElemCount
n forall a. Num a => a -> a -> a
- ElemCount
1])

sizeofLiftedPrimArray :: LiftedPrimArray a -> ByteCount
sizeofLiftedPrimArray :: forall a. LiftedPrimArray a -> ByteCount
sizeofLiftedPrimArray (LiftedPrimArray ByteArray
arr) = coerce :: forall a b. Coercible a b => a -> b
coerce (ByteArray -> Int
sizeofByteArray ByteArray
arr)

lengthLiftedPrimArray :: LiftedPrim a => LiftedPrimArray a -> ElemCount
lengthLiftedPrimArray :: forall a. LiftedPrim a => LiftedPrimArray a -> ElemCount
lengthLiftedPrimArray pa :: LiftedPrimArray a
pa@(LiftedPrimArray ByteArray
arr) =
  let elemSize :: Int
elemSize = coerce :: forall a b. Coercible a b => a -> b
coerce (forall a. StaticByteSized a => Proxy a -> ByteCount
staticByteSize (forall (f :: * -> *) a. f a -> Proxy a
proxyForF LiftedPrimArray a
pa))
      arrSize :: Int
arrSize = ByteArray -> Int
sizeofByteArray ByteArray
arr
  in  coerce :: forall a b. Coercible a b => a -> b
coerce (forall a. Integral a => a -> a -> a
div Int
arrSize Int
elemSize)

cloneLiftedPrimArray :: LiftedPrim a => LiftedPrimArray a -> ElemCount -> ElemCount -> LiftedPrimArray a
cloneLiftedPrimArray :: forall a.
LiftedPrim a =>
LiftedPrimArray a -> ElemCount -> ElemCount -> LiftedPrimArray a
cloneLiftedPrimArray pa :: LiftedPrimArray a
pa@(LiftedPrimArray ByteArray
arr) ElemCount
off ElemCount
len =
  let elemSize :: ByteCount
elemSize = forall a. StaticByteSized a => Proxy a -> ByteCount
staticByteSize (forall (f :: * -> *) a. f a -> Proxy a
proxyForF LiftedPrimArray a
pa)
      byteOff :: ByteCount
byteOff = coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
off forall a. Num a => a -> a -> a
* ByteCount
elemSize
      byteLen :: ByteCount
byteLen = coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
len forall a. Num a => a -> a -> a
* ByteCount
elemSize
      arr' :: ByteArray
arr' = ByteArray -> Int -> Int -> ByteArray
cloneByteArray ByteArray
arr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
byteOff) (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
byteLen)
  in  forall a. ByteArray -> LiftedPrimArray a
LiftedPrimArray ByteArray
arr'

replicateLiftedPrimArray :: LiftedPrim a => ElemCount -> a -> LiftedPrimArray a
replicateLiftedPrimArray :: forall a. LiftedPrim a => ElemCount -> a -> LiftedPrimArray a
replicateLiftedPrimArray ElemCount
len a
val = forall a. ByteArray -> LiftedPrimArray a
LiftedPrimArray forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutableByteArray s)) -> ByteArray
runByteArray forall a b. (a -> b) -> a -> b
$ do
  let elemSize :: ByteCount
elemSize = forall a. StaticByteSized a => Proxy a -> ByteCount
staticByteSize (forall a. a -> Proxy a
proxyFor a
val)
      byteLen :: ByteCount
byteLen = coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
len forall a. Num a => a -> a -> a
* ByteCount
elemSize
  MutableByteArray s
arr <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
byteLen)
  forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [ElemCount
0 .. ElemCount
len forall a. Num a => a -> a -> a
- ElemCount
1] forall a b. (a -> b) -> a -> b
$ \ElemCount
pos ->
    forall a (m :: * -> *).
(LiftedPrim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> ByteCount -> a -> m ()
writeArrayLiftedInBytes MutableByteArray s
arr (coerce :: forall a b. Coercible a b => a -> b
coerce ElemCount
pos forall a. Num a => a -> a -> a
* ByteCount
elemSize) a
val
  forall (f :: * -> *) a. Applicative f => a -> f a
pure MutableByteArray s
arr