module PrimitiveExtras.PrimArray
where

import PrimitiveExtras.Prelude hiding (replicateM, traverse_)
import PrimitiveExtras.Types
import qualified Data.Serialize as Cereal
import qualified Data.Vector.Unboxed as UnboxedVector
import qualified Data.Vector.Primitive as PrimitiveVector
import qualified PrimitiveExtras.Folds as Folds
import qualified PrimitiveExtras.FoldMs as FoldMs
import qualified Data.ByteString.Short.Internal as ShortByteString


{-|
Construct from a primitive vector.
In case the vector is not a slice, it is an /O(1)/ op.
-}
primitiveVector :: Prim a => PrimitiveVector.Vector a -> PrimArray a
primitiveVector :: Vector a -> PrimArray a
primitiveVector (PrimitiveVector.Vector Int
offset Int
length (ByteArray ByteArray#
unliftedByteArray)) = let
  primArray :: PrimArray a
primArray = ByteArray# -> PrimArray a
forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
unliftedByteArray
  in if Int
offset Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
length Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== PrimArray a -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
primArray
    then PrimArray a
primArray
    else (forall s. ST s (PrimArray a)) -> PrimArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (PrimArray a)) -> PrimArray a)
-> (forall s. ST s (PrimArray a)) -> PrimArray a
forall a b. (a -> b) -> a -> b
$ do
      MutablePrimArray s a
ma <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
length
      MutablePrimArray (PrimState (ST s)) a
-> Int -> PrimArray a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
ma Int
0 PrimArray a
primArray Int
offset Int
length
      MutablePrimArray (PrimState (ST s)) a -> ST s (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
ma

oneHot :: Prim a => Int {-^ Size -} -> Int {-^ Index -} -> a -> PrimArray a
oneHot :: Int -> Int -> a -> PrimArray a
oneHot Int
size Int
index a
value =
  (forall s. ST s (PrimArray a)) -> PrimArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (PrimArray a)) -> PrimArray a)
-> (forall s. ST s (PrimArray a)) -> PrimArray a
forall a b. (a -> b) -> a -> b
$ do
    MutablePrimArray s a
marr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
size
    MutablePrimArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
index a
value
    MutablePrimArray (PrimState (ST s)) a -> ST s (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr

generate :: Prim a => Int -> (Int -> IO a) -> IO (PrimArray a)
generate :: Int -> (Int -> IO a) -> IO (PrimArray a)
generate Int
size Int -> IO a
elementIO =
  do
    MutablePrimArray RealWorld a
array <- Int -> IO (MutablePrimArray (PrimState IO) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
size
    let
      loop :: Int -> IO (PrimArray a)
loop Int
index =
        if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size
          then do
            a
element <- Int -> IO a
elementIO Int
index
            MutablePrimArray (PrimState IO) a -> Int -> a -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld a
MutablePrimArray (PrimState IO) a
array Int
index a
element
            Int -> IO (PrimArray a)
loop (Int -> Int
forall a. Enum a => a -> a
succ Int
index)
          else MutablePrimArray (PrimState IO) a -> IO (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld a
MutablePrimArray (PrimState IO) a
array
      in Int -> IO (PrimArray a)
loop Int
0

replicate :: Prim a => Int -> IO a -> IO (PrimArray a)
replicate :: Int -> IO a -> IO (PrimArray a)
replicate Int
size IO a
elementIO =
  do
    MutablePrimArray RealWorld a
array <- Int -> IO (MutablePrimArray (PrimState IO) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
size
    let
      loop :: Int -> IO (PrimArray a)
loop Int
index =
        if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size
          then do
            a
element <- IO a
elementIO
            MutablePrimArray (PrimState IO) a -> Int -> a -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld a
MutablePrimArray (PrimState IO) a
array Int
index a
element
            Int -> IO (PrimArray a)
loop (Int -> Int
forall a. Enum a => a -> a
succ Int
index)
          else MutablePrimArray (PrimState IO) a -> IO (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld a
MutablePrimArray (PrimState IO) a
array
      in Int -> IO (PrimArray a)
loop Int
0

{-| Please notice that this function is highly untested -}
replicateM :: (Monad m, Prim element) => Int -> m element -> m (PrimArray element)
replicateM :: Int -> m element -> m (PrimArray element)
replicateM Int
size m element
elementM =
  do
    !MutablePrimArray RealWorld element
mutable <- MutablePrimArray RealWorld element
-> m (MutablePrimArray RealWorld element)
forall (m :: * -> *) a. Monad m => a -> m a
return (IO (MutablePrimArray RealWorld element)
-> MutablePrimArray RealWorld element
forall a. IO a -> a
unsafeDupablePerformIO (Int -> IO (MutablePrimArray (PrimState IO) element)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
size))
    let 
      iterate :: Int -> m (PrimArray element)
iterate Int
index =
        if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size
          then do
            element
element <- m element
elementM
            let !() = IO () -> ()
forall a. IO a -> a
unsafeDupablePerformIO (MutablePrimArray (PrimState IO) element -> Int -> element -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld element
MutablePrimArray (PrimState IO) element
mutable Int
index element
element)
            Int -> m (PrimArray element)
iterate (Int -> Int
forall a. Enum a => a -> a
succ Int
index)
          else PrimArray element -> m (PrimArray element)
forall (m :: * -> *) a. Monad m => a -> m a
return (IO (PrimArray element) -> PrimArray element
forall a. IO a -> a
unsafePerformIO (MutablePrimArray (PrimState IO) element -> IO (PrimArray element)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld element
MutablePrimArray (PrimState IO) element
mutable))
      in Int -> m (PrimArray element)
iterate Int
0

traverse_ :: (a -> f b) -> PrimArray a -> f ()
traverse_ = (a -> f b) -> PrimArray a -> f ()
forall (f :: * -> *) a b.
(Applicative f, Prim a) =>
(a -> f b) -> PrimArray a -> f ()
traversePrimArray_

traverseWithIndexInRange_ :: Prim a => PrimArray a -> Int -> Int -> (Int -> a -> IO ()) -> IO ()
traverseWithIndexInRange_ :: PrimArray a -> Int -> Int -> (Int -> a -> IO ()) -> IO ()
traverseWithIndexInRange_ PrimArray a
primArray Int
from Int
to Int -> a -> IO ()
action =
  let iterate :: Int -> IO ()
iterate Int
index = if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
to
        then do
          Int -> a -> IO ()
action Int
index (a -> IO ()) -> a -> IO ()
forall a b. (a -> b) -> a -> b
$! PrimArray a -> Int -> a
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
primArray Int
index
          Int -> IO ()
iterate (Int -> Int
forall a. Enum a => a -> a
succ Int
index)
        else () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      in Int -> IO ()
iterate Int
from

toElementsUnfoldl :: Prim prim => PrimArray prim -> Unfoldl prim
toElementsUnfoldl :: PrimArray prim -> Unfoldl prim
toElementsUnfoldl PrimArray prim
ba = (forall x. (x -> prim -> x) -> x -> x) -> Unfoldl prim
forall a. (forall x. (x -> a -> x) -> x -> x) -> Unfoldl a
Unfoldl ((forall x. (x -> prim -> x) -> x -> x) -> Unfoldl prim)
-> (forall x. (x -> prim -> x) -> x -> x) -> Unfoldl prim
forall a b. (a -> b) -> a -> b
$ \x -> prim -> x
f x
z -> (x -> prim -> x) -> x -> PrimArray prim -> x
forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldlPrimArray' x -> prim -> x
f x
z PrimArray prim
ba

toElementsUnfoldlM :: (Monad m, Prim prim) => PrimArray prim -> UnfoldlM m prim
toElementsUnfoldlM :: PrimArray prim -> UnfoldlM m prim
toElementsUnfoldlM PrimArray prim
ba = (forall x. (x -> prim -> m x) -> x -> m x) -> UnfoldlM m prim
forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM ((forall x. (x -> prim -> m x) -> x -> m x) -> UnfoldlM m prim)
-> (forall x. (x -> prim -> m x) -> x -> m x) -> UnfoldlM m prim
forall a b. (a -> b) -> a -> b
$ \x -> prim -> m x
f x
z -> (x -> prim -> m x) -> x -> PrimArray prim -> m x
forall a (m :: * -> *) b.
(Prim a, Monad m) =>
(b -> a -> m b) -> b -> PrimArray a -> m b
foldlPrimArrayM' x -> prim -> m x
f x
z PrimArray prim
ba

toByteArray :: PrimArray a -> ByteArray
toByteArray :: PrimArray a -> ByteArray
toByteArray (PrimArray ByteArray#
unliftedByteArray) =
  ByteArray# -> ByteArray
ByteArray ByteArray#
unliftedByteArray

toPrimitiveVector :: Prim a => PrimArray a -> PrimitiveVector.Vector a
toPrimitiveVector :: PrimArray a -> Vector a
toPrimitiveVector PrimArray a
primArray =
  Int -> Int -> ByteArray -> Vector a
forall a. Int -> Int -> ByteArray -> Vector a
PrimitiveVector.Vector Int
0 (PrimArray a -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
primArray) (PrimArray a -> ByteArray
forall a. PrimArray a -> ByteArray
toByteArray PrimArray a
primArray)

toUnboxedVector :: Prim a => PrimArray a -> UnboxedVector.Vector a
toUnboxedVector :: PrimArray a -> Vector a
toUnboxedVector PrimArray a
primArray =
  Vector a -> Vector a
forall a b. a -> b
unsafeCoerce (PrimArray a -> Vector a
forall a. Prim a => PrimArray a -> Vector a
toPrimitiveVector PrimArray a
primArray)

cerealGet :: Prim element => Cereal.Get Int -> Cereal.Get element -> Cereal.Get (PrimArray element)
cerealGet :: Get Int -> Get element -> Get (PrimArray element)
cerealGet Get Int
int Get element
element =
  do
    Int
size <- Get Int
int
    Int -> Get element -> Get (PrimArray element)
forall (m :: * -> *) element.
(Monad m, Prim element) =>
Int -> m element -> m (PrimArray element)
replicateM Int
size Get element
element

cerealGetAsInMemory :: Prim element => Cereal.Get Int -> Cereal.Get (PrimArray element)
cerealGetAsInMemory :: Get Int -> Get (PrimArray element)
cerealGetAsInMemory Get Int
int =
  do
    Int
size <- Get Int
int
    ShortByteString.SBS ByteArray#
ba <- Int -> Get ShortByteString
Cereal.getShortByteString Int
size
    PrimArray element -> Get (PrimArray element)
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteArray# -> PrimArray element
forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
ba)

cerealPut :: Prim element => Cereal.Putter Int -> Cereal.Putter element -> Cereal.Putter (PrimArray element)
cerealPut :: Putter Int -> Putter element -> Putter (PrimArray element)
cerealPut Putter Int
int Putter element
element PrimArray element
primArrayValue =
  Put
size Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Put
elements
  where
    size :: Put
size = Putter Int
int (PrimArray element -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray element
primArrayValue)
    elements :: Put
elements = Putter element -> Putter (PrimArray element)
forall (f :: * -> *) a b.
(Applicative f, Prim a) =>
(a -> f b) -> PrimArray a -> f ()
traverse_ Putter element
element PrimArray element
primArrayValue

cerealPutAsInMemory :: Prim element => Cereal.Putter Int -> Cereal.Putter (PrimArray element)
cerealPutAsInMemory :: Putter Int -> Putter (PrimArray element)
cerealPutAsInMemory Putter Int
int primArrayValue :: PrimArray element
primArrayValue@(PrimArray ByteArray#
ba) =
  Put
size Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Put
elements
  where
    size :: Put
size = Putter Int
int (ByteArray -> Int
sizeofByteArray (ByteArray# -> ByteArray
ByteArray ByteArray#
ba))
    elements :: Put
elements = Putter ShortByteString
Cereal.putShortByteString (ByteArray# -> ShortByteString
ShortByteString.SBS ByteArray#
ba)

{-|
Given a size of the array,
construct a fold, which produces an array of index counts.
-}
indexCountsFold :: (Integral count, Prim count) => Int {-^ Array size -} -> Fold Int (PrimArray count)
indexCountsFold :: Int -> Fold Int (PrimArray count)
indexCountsFold = Int -> Fold Int (PrimArray count)
forall count.
(Integral count, Prim count) =>
Int -> Fold Int (PrimArray count)
Folds.indexCounts

{-|
Given a size of the array,
construct a fold, which produces an array of elements.
-}
elementsFoldM :: Prim a => Int {-^ Array size -} -> FoldM IO a (PrimArray a)
elementsFoldM :: Int -> FoldM IO a (PrimArray a)
elementsFoldM = Int -> FoldM IO a (PrimArray a)
forall a. Prim a => Int -> FoldM IO a (PrimArray a)
FoldMs.primArray