{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module PrimitiveExtras.PrimArray where

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

-- |
-- 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 :: forall a. Prim a => 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) =>
  -- | Size
  Int ->
  -- | Index
  Int ->
  a ->
  PrimArray a
oneHot :: forall a. Prim a => 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 :: forall a. Prim a => 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 :: forall a. Prim a => 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 :: forall (m :: * -> *) element.
(Monad m, Prim element) =>
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 a. a -> m a
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 a. a -> m a
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_ :: (Applicative f, Prim a) => (a -> f b) -> PrimArray a -> f ()
traverse_ :: forall (f :: * -> *) a b.
(Applicative f, Prim a) =>
(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_ :: forall a.
Prim a =>
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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   in Int -> IO ()
iterate Int
from

toElementsUnfoldl :: (Prim prim) => PrimArray prim -> Unfoldl prim
toElementsUnfoldl :: forall prim. Prim prim => 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 :: forall (m :: * -> *) prim.
(Monad m, Prim prim) =>
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 :: forall a. PrimArray a -> ByteArray
toByteArray (PrimArray ByteArray#
unliftedByteArray) =
  ByteArray# -> ByteArray
ByteArray ByteArray#
unliftedByteArray

toPrimitiveVector :: (Prim a) => PrimArray a -> PrimitiveVector.Vector a
toPrimitiveVector :: forall a. Prim a => 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 :: forall a. Prim a => 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 :: forall element.
Prim element =>
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 :: forall element. Prim element => 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 a. a -> Get a
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 :: forall element.
Prim element =>
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 -> PrimArray element -> Put
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 :: forall element.
Prim element =>
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) =>
  -- | Array size
  Int ->
  Fold Int (PrimArray count)
indexCountsFold :: forall count.
(Integral count, Prim count) =>
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) =>
  -- | Array size
  Int ->
  FoldM IO a (PrimArray a)
elementsFoldM :: forall a. Prim a => 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