{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE RoleAnnotations #-}

-- |
-- Module      : Data.Primitive.PrimArray
-- Copyright   : (c) Roman Leshchinskiy 2009-2012
-- License     : BSD-style
--
-- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
-- Portability : non-portable
--
-- Arrays of unboxed primitive types. The functions provided by this module
-- match the behavior of those provided by "Data.Primitive.ByteArray", and
-- the underlying types and primops that back them are the same.
-- However, the type constructors 'PrimArray' and 'MutablePrimArray' take one additional
-- argument compared to their respective counterparts 'ByteArray' and 'Data.Primitive.ByteArray.MutableByteArray'.
-- This argument is used to designate the type of element in the array.
-- Consequently, all functions in this module accept length and indices in
-- terms of elements, not bytes.
--
-- @since 0.6.4.0

module Data.Primitive.PrimArray
  ( -- * Types
    PrimArray(..)
  , MutablePrimArray(..)
    -- * Allocation
  , newPrimArray
  , newPinnedPrimArray
  , newAlignedPinnedPrimArray
  , resizeMutablePrimArray
  , shrinkMutablePrimArray
    -- * Element Access
  , readPrimArray
  , writePrimArray
  , indexPrimArray
    -- * Freezing and Thawing
  , freezePrimArray
  , thawPrimArray
  , runPrimArray
  , unsafeFreezePrimArray
  , unsafeThawPrimArray
    -- * Block Operations
  , copyPrimArray
  , copyMutablePrimArray
  , copyPrimArrayToPtr
  , copyMutablePrimArrayToPtr
  , copyPtrToMutablePrimArray
  , clonePrimArray
  , cloneMutablePrimArray
  , setPrimArray
    -- * Information
  , sameMutablePrimArray
  , getSizeofMutablePrimArray
  , sizeofMutablePrimArray
  , sizeofPrimArray
  , primArrayContents
  , withPrimArrayContents
  , mutablePrimArrayContents
  , withMutablePrimArrayContents
#if __GLASGOW_HASKELL__ >= 802
  , isPrimArrayPinned
  , isMutablePrimArrayPinned
#endif
    -- * List Conversion
  , primArrayToList
  , primArrayFromList
  , primArrayFromListN
    -- * Folding
  , foldrPrimArray
  , foldrPrimArray'
  , foldlPrimArray
  , foldlPrimArray'
  , foldlPrimArrayM'
    -- * Effectful Folding
  , traversePrimArray_
  , itraversePrimArray_
    -- * Map/Create
  , emptyPrimArray
  , mapPrimArray
  , imapPrimArray
  , generatePrimArray
  , replicatePrimArray
  , filterPrimArray
  , mapMaybePrimArray
    -- * Effectful Map/Create
    -- $effectfulMapCreate

    -- ** Lazy Applicative
  , traversePrimArray
  , itraversePrimArray
  , generatePrimArrayA
  , replicatePrimArrayA
  , filterPrimArrayA
  , mapMaybePrimArrayA
    -- ** Strict Primitive Monadic
  , traversePrimArrayP
  , itraversePrimArrayP
  , generatePrimArrayP
  , replicatePrimArrayP
  , filterPrimArrayP
  , mapMaybePrimArrayP
  ) where

import GHC.Exts
import Data.Primitive.Types
import Data.Primitive.ByteArray (ByteArray(..))
import Data.Proxy
#if !MIN_VERSION_base(4,18,0)
import Control.Applicative (liftA2)
#endif
import Control.DeepSeq
import Control.Monad (when)
import Control.Monad.Primitive
import Control.Monad.ST
import qualified Data.List as L
import qualified Data.Primitive.ByteArray as PB
import qualified Data.Primitive.Types as PT
#if MIN_VERSION_base(4,10,0)
import qualified GHC.ST as GHCST
#endif
import Language.Haskell.TH.Syntax (Lift (..))

import Data.Semigroup

#if __GLASGOW_HASKELL__ >= 802
import qualified GHC.Exts as Exts
#endif

import Data.Primitive.Internal.Operations (mutableByteArrayContentsShim)

-- | Arrays of unboxed elements. This accepts types like 'Double', 'Char',
-- 'Int' and 'Word', as well as their fixed-length variants ('Word8',
-- 'Word16', etc.). Since the elements are unboxed, a 'PrimArray' is strict
-- in its elements. This differs from the behavior of 'Data.Primitive.Array.Array',
-- which is lazy in its elements.
data PrimArray a = PrimArray ByteArray#

type role PrimArray nominal

instance Lift (PrimArray a) where
#if MIN_VERSION_template_haskell(2,16,0)
  liftTyped :: forall (m :: * -> *).
Quote m =>
PrimArray a -> Code m (PrimArray a)
liftTyped PrimArray a
ary = [|| byteArrayToPrimArray ba ||]
#else
  lift ary = [| byteArrayToPrimArray ba |]
#endif
    where
      ba :: ByteArray
ba = forall a. PrimArray a -> ByteArray
primArrayToByteArray PrimArray a
ary

instance NFData (PrimArray a) where
  rnf :: PrimArray a -> ()
rnf (PrimArray ByteArray#
_) = ()

-- | Mutable primitive arrays associated with a primitive state token.
-- These can be written to and read from in a monadic context that supports
-- sequencing, such as 'IO' or 'ST'. Typically, a mutable primitive array will
-- be built and then converted to an immutable primitive array using
-- 'unsafeFreezePrimArray'. However, it is also acceptable to simply discard
-- a mutable primitive array since it lives in managed memory and will be
-- garbage collected when no longer referenced.
data MutablePrimArray s a = MutablePrimArray (MutableByteArray# s)

instance Eq (MutablePrimArray s a) where
  == :: MutablePrimArray s a -> MutablePrimArray s a -> Bool
(==) = forall s a. MutablePrimArray s a -> MutablePrimArray s a -> Bool
sameMutablePrimArray

instance NFData (MutablePrimArray s a) where
  rnf :: MutablePrimArray s a -> ()
rnf (MutablePrimArray MutableByteArray# s
_) = ()

sameByteArray :: ByteArray# -> ByteArray# -> Bool
sameByteArray :: ByteArray# -> ByteArray# -> Bool
sameByteArray ByteArray#
ba1 ByteArray#
ba2 =
    case forall a. a -> a -> Int#
reallyUnsafePtrEquality# (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# ByteArray#
ba1 :: ()) (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# ByteArray#
ba2 :: ()) of
      Int#
r -> Int# -> Bool
isTrue# Int#
r

-- | @since 0.6.4.0
instance (Eq a, Prim a) => Eq (PrimArray a) where
  a1 :: PrimArray a
a1@(PrimArray ByteArray#
ba1#) == :: PrimArray a -> PrimArray a -> Bool
== a2 :: PrimArray a
a2@(PrimArray ByteArray#
ba2#)
    | ByteArray# -> ByteArray# -> Bool
sameByteArray ByteArray#
ba1# ByteArray#
ba2# = Bool
True
    | Int
sz1 forall a. Eq a => a -> a -> Bool
/= Int
sz2 = Bool
False
    | Bool
otherwise = Int -> Bool
loop (forall a. Integral a => a -> a -> a
quot Int
sz1 (forall a. Prim a => Int
sizeOfType @a) forall a. Num a => a -> a -> a
- Int
1)
    where
    -- Here, we take the size in bytes, not in elements. We do this
    -- since it allows us to defer performing the division to
    -- calculate the size in elements.
    sz1 :: Int
sz1 = ByteArray -> Int
PB.sizeofByteArray (ByteArray# -> ByteArray
ByteArray ByteArray#
ba1#)
    sz2 :: Int
sz2 = ByteArray -> Int
PB.sizeofByteArray (ByteArray# -> ByteArray
ByteArray ByteArray#
ba2#)
    loop :: Int -> Bool
loop !Int
i
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = Bool
True
      | Bool
otherwise = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
a1 Int
i forall a. Eq a => a -> a -> Bool
== forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
a2 Int
i Bool -> Bool -> Bool
&& Int -> Bool
loop (Int
i forall a. Num a => a -> a -> a
- Int
1)
  {-# INLINE (==) #-}

-- | Lexicographic ordering. Subject to change between major versions.
--
-- @since 0.6.4.0
instance (Ord a, Prim a) => Ord (PrimArray a) where
  compare :: PrimArray a -> PrimArray a -> Ordering
compare a1 :: PrimArray a
a1@(PrimArray ByteArray#
ba1#) a2 :: PrimArray a
a2@(PrimArray ByteArray#
ba2#)
    | ByteArray# -> ByteArray# -> Bool
sameByteArray ByteArray#
ba1# ByteArray#
ba2# = Ordering
EQ
    | Bool
otherwise = Int -> Ordering
loop Int
0
    where
    sz1 :: Int
sz1 = ByteArray -> Int
PB.sizeofByteArray (ByteArray# -> ByteArray
ByteArray ByteArray#
ba1#)
    sz2 :: Int
sz2 = ByteArray -> Int
PB.sizeofByteArray (ByteArray# -> ByteArray
ByteArray ByteArray#
ba2#)
    sz :: Int
sz = forall a. Integral a => a -> a -> a
quot (forall a. Ord a => a -> a -> a
min Int
sz1 Int
sz2) (forall a. Prim a => Int
sizeOfType @a)
    loop :: Int -> Ordering
loop !Int
i
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
sz = forall a. Ord a => a -> a -> Ordering
compare (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
a1 Int
i) (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
a2 Int
i) forall a. Semigroup a => a -> a -> a
<> Int -> Ordering
loop (Int
i forall a. Num a => a -> a -> a
+ Int
1)
      | Bool
otherwise = forall a. Ord a => a -> a -> Ordering
compare Int
sz1 Int
sz2
  {-# INLINE compare #-}

-- | @since 0.6.4.0
instance Prim a => IsList (PrimArray a) where
  type Item (PrimArray a) = a
  fromList :: [Item (PrimArray a)] -> PrimArray a
fromList = forall a. Prim a => [a] -> PrimArray a
primArrayFromList
  fromListN :: Int -> [Item (PrimArray a)] -> PrimArray a
fromListN = forall a. Prim a => Int -> [a] -> PrimArray a
primArrayFromListN
  toList :: PrimArray a -> [Item (PrimArray a)]
toList = forall a. Prim a => PrimArray a -> [a]
primArrayToList

-- | @since 0.6.4.0
instance (Show a, Prim a) => Show (PrimArray a) where
  showsPrec :: Int -> PrimArray a -> ShowS
showsPrec Int
_ PrimArray a
a = forall a. Show a => a -> ShowS
shows (forall a. Prim a => PrimArray a -> [a]
primArrayToList PrimArray a
a)

die :: String -> String -> a
die :: forall a. String -> String -> a
die String
fun String
problem = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Data.Primitive.PrimArray." forall a. [a] -> [a] -> [a]
++ String
fun forall a. [a] -> [a] -> [a]
++ String
": " forall a. [a] -> [a] -> [a]
++ String
problem

-- | Create a 'PrimArray' from a list.
--
-- @primArrayFromList vs = `primArrayFromListN` (length vs) vs@
primArrayFromList :: Prim a => [a] -> PrimArray a
primArrayFromList :: forall a. Prim a => [a] -> PrimArray a
primArrayFromList [a]
vs = forall a. Prim a => Int -> [a] -> PrimArray a
primArrayFromListN (forall (t :: * -> *) a. Foldable t => t a -> Int
L.length [a]
vs) [a]
vs

-- | Create a 'PrimArray' from a list of a known length. If the length
-- of the list does not match the given length, this throws an exception.
primArrayFromListN :: forall a. Prim a => Int -> [a] -> PrimArray a
primArrayFromListN :: forall a. Prim a => Int -> [a] -> PrimArray a
primArrayFromListN Int
len [a]
vs = forall a. (forall s. ST s a) -> a
runST forall s. ST s (PrimArray a)
run where
  run :: forall s. ST s (PrimArray a)
  run :: forall s. ST s (PrimArray a)
run = do
    MutablePrimArray s a
arr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
    let go :: [a] -> Int -> ST s ()
        go :: [a] -> Int -> ST s ()
go [] !Int
ix = if Int
ix forall a. Eq a => a -> a -> Bool
== Int
len
          then forall (m :: * -> *) a. Monad m => a -> m a
return ()
          else forall a. String -> String -> a
die String
"fromListN" String
"list length less than specified size"
        go (a
a : [a]
as) !Int
ix = if Int
ix forall a. Ord a => a -> a -> Bool
< Int
len
          then do
            forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s a
arr Int
ix a
a
            [a] -> Int -> ST s ()
go [a]
as (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
          else forall a. String -> String -> a
die String
"fromListN" String
"list length greater than specified size"
    [a] -> Int -> ST s ()
go [a]
vs Int
0
    forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
arr

-- | Convert a 'PrimArray' to a list.
{-# INLINE primArrayToList #-}
primArrayToList :: forall a. Prim a => PrimArray a -> [a]
primArrayToList :: forall a. Prim a => PrimArray a -> [a]
primArrayToList PrimArray a
xs = forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldrPrimArray a -> b -> b
c b
n PrimArray a
xs)

primArrayToByteArray :: PrimArray a -> PB.ByteArray
primArrayToByteArray :: forall a. PrimArray a -> ByteArray
primArrayToByteArray (PrimArray ByteArray#
x) = ByteArray# -> ByteArray
PB.ByteArray ByteArray#
x

byteArrayToPrimArray :: ByteArray -> PrimArray a
byteArrayToPrimArray :: forall a. ByteArray -> PrimArray a
byteArrayToPrimArray (PB.ByteArray ByteArray#
x) = forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
x

-- | @since 0.6.4.0
instance Semigroup (PrimArray a) where
  PrimArray a
x <> :: PrimArray a -> PrimArray a -> PrimArray a
<> PrimArray a
y = forall a. ByteArray -> PrimArray a
byteArrayToPrimArray (forall a. PrimArray a -> ByteArray
primArrayToByteArray PrimArray a
x forall a. Semigroup a => a -> a -> a
<> forall a. PrimArray a -> ByteArray
primArrayToByteArray PrimArray a
y)
  sconcat :: NonEmpty (PrimArray a) -> PrimArray a
sconcat = forall a. ByteArray -> PrimArray a
byteArrayToPrimArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Semigroup a => NonEmpty a -> a
sconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. PrimArray a -> ByteArray
primArrayToByteArray
  stimes :: forall b. Integral b => b -> PrimArray a -> PrimArray a
stimes b
i PrimArray a
arr = forall a. ByteArray -> PrimArray a
byteArrayToPrimArray (forall a b. (Semigroup a, Integral b) => b -> a -> a
stimes b
i (forall a. PrimArray a -> ByteArray
primArrayToByteArray PrimArray a
arr))

-- | @since 0.6.4.0
instance Monoid (PrimArray a) where
  mempty :: PrimArray a
mempty = forall a. PrimArray a
emptyPrimArray
#if !(MIN_VERSION_base(4,11,0))
  mappend = (<>)
#endif
  mconcat :: [PrimArray a] -> PrimArray a
mconcat = forall a. ByteArray -> PrimArray a
byteArrayToPrimArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. PrimArray a -> ByteArray
primArrayToByteArray

-- | The empty 'PrimArray'.
emptyPrimArray :: PrimArray a
{-# NOINLINE emptyPrimArray #-}
emptyPrimArray :: forall a. PrimArray a
emptyPrimArray = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState (ST s))
s0# -> case forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
0# State# (PrimState (ST s))
s0# of
  (# State# s
s1#, MutableByteArray# s
arr# #) -> case forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# s
arr# State# s
s1# of
    (# State# s
s2#, ByteArray#
arr'# #) -> (# State# s
s2#, forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
arr'# #)

-- | Create a new mutable primitive array of the given length. The
-- underlying memory is left uninitialized.
--
-- /Note:/ this function does not check if the input is non-negative.
newPrimArray :: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
{-# INLINE newPrimArray #-}
newPrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (I# Int#
n#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\State# (PrimState m)
s# ->
      case forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# (Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) State# (PrimState m)
s# of
        (# State# (PrimState m)
s'#, MutableByteArray# (PrimState m)
arr# #) -> (# State# (PrimState m)
s'#, forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# (PrimState m)
arr# #)
    )

-- | Resize a mutable primitive array. The new size is given in elements.
--
-- This will either resize the array in-place or, if not possible, allocate the
-- contents into a new, unpinned array and copy the original array\'s contents.
--
-- To avoid undefined behaviour, the original 'MutablePrimArray' shall not be
-- accessed anymore after a 'resizeMutablePrimArray' has been performed.
-- Moreover, no reference to the old one should be kept in order to allow
-- garbage collection of the original 'MutablePrimArray' in case a new
-- 'MutablePrimArray' had to be allocated.
resizeMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a
  -> Int -- ^ new size
  -> m (MutablePrimArray (PrimState m) a)
{-# INLINE resizeMutablePrimArray #-}
resizeMutablePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray (MutablePrimArray MutableByteArray# (PrimState m)
arr#) (I# Int#
n#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\State# (PrimState m)
s# -> case forall d.
MutableByteArray# d
-> Int# -> State# d -> (# State# d, MutableByteArray# d #)
resizeMutableByteArray# MutableByteArray# (PrimState m)
arr# (Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) State# (PrimState m)
s# of
                        (# State# (PrimState m)
s'#, MutableByteArray# (PrimState m)
arr'# #) -> (# State# (PrimState m)
s'#, forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# (PrimState m)
arr'# #))

-- | Shrink a mutable primitive array. The new size is given in elements.
-- It must be smaller than the old size. The array will be resized in place.
shrinkMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a
  -> Int -- ^ new size
  -> m ()
{-# INLINE shrinkMutablePrimArray #-}
shrinkMutablePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> m ()
shrinkMutablePrimArray (MutablePrimArray MutableByteArray# (PrimState m)
arr#) (I# Int#
n#)
  = forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (forall d. MutableByteArray# d -> Int# -> State# d -> State# d
shrinkMutableByteArray# MutableByteArray# (PrimState m)
arr# (Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)))

-- | Read a value from the array at the given index.
--
-- /Note:/ this function does not do bounds checking.
readPrimArray :: (Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> Int -> m a
{-# INLINE readPrimArray #-}
readPrimArray :: forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray (MutablePrimArray MutableByteArray# (PrimState m)
arr#) (I# Int#
i#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (forall a s.
Prim a =>
MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
readByteArray# MutableByteArray# (PrimState m)
arr# Int#
i#)

-- | Write an element to the given index.
--
-- /Note:/ this function does not do bounds checking.
writePrimArray
  :: (Prim a, PrimMonad m)
  => MutablePrimArray (PrimState m) a -- ^ array
  -> Int -- ^ index
  -> a -- ^ element
  -> m ()
{-# INLINE writePrimArray #-}
writePrimArray :: forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray (MutablePrimArray MutableByteArray# (PrimState m)
arr#) (I# Int#
i#) a
x
  = forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (forall a s.
Prim a =>
MutableByteArray# s -> Int# -> a -> State# s -> State# s
writeByteArray# MutableByteArray# (PrimState m)
arr# Int#
i# a
x)

-- | Copy part of a mutable array into another mutable array.
-- In the case that the destination and
-- source arrays are the same, the regions may overlap.
--
-- /Note:/ this function does not do bounds or overlap checking.
copyMutablePrimArray :: forall m a.
     (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a -- ^ destination array
  -> Int -- ^ offset into destination array
  -> MutablePrimArray (PrimState m) a -- ^ source array
  -> Int -- ^ offset into source array
  -> Int -- ^ number of elements to copy
  -> m ()
{-# INLINE copyMutablePrimArray #-}
copyMutablePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray (MutablePrimArray MutableByteArray# (PrimState m)
dst#) (I# Int#
doff#) (MutablePrimArray MutableByteArray# (PrimState m)
src#) (I# Int#
soff#) (I# Int#
n#)
  = forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray#
      MutableByteArray# (PrimState m)
src#
      (Int#
soff# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a))
      MutableByteArray# (PrimState m)
dst#
      (Int#
doff# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a))
      (Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a))
    )

-- | Copy part of an array into another mutable array.
--
-- /Note:/ this function does not do bounds or overlap checking.
copyPrimArray :: forall m a.
     (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a -- ^ destination array
  -> Int -- ^ offset into destination array
  -> PrimArray a -- ^ source array
  -> Int -- ^ offset into source array
  -> Int -- ^ number of elements to copy
  -> m ()
{-# INLINE copyPrimArray #-}
copyPrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray (MutablePrimArray MutableByteArray# (PrimState m)
dst#) (I# Int#
doff#) (PrimArray ByteArray#
src#) (I# Int#
soff#) (I# Int#
n#)
  = forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (forall d.
ByteArray#
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyByteArray#
      ByteArray#
src#
      (Int#
soff# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a))
      MutableByteArray# (PrimState m)
dst#
      (Int#
doff# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a))
      (Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a))
    )

-- | Copy a slice of an immutable primitive array to a pointer.
-- The offset and length are given in elements of type @a@.
-- This function assumes that the 'Prim' instance of @a@
-- agrees with the 'Storable' instance.
--
-- /Note:/ this function does not do bounds or overlap checking.
copyPrimArrayToPtr :: forall m a. (PrimMonad m, Prim a)
  => Ptr a -- ^ destination pointer
  -> PrimArray a -- ^ source array
  -> Int -- ^ offset into source array
  -> Int -- ^ number of elements to copy
  -> m ()
{-# INLINE copyPrimArrayToPtr #-}
copyPrimArrayToPtr :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Ptr a -> PrimArray a -> Int -> Int -> m ()
copyPrimArrayToPtr (Ptr Addr#
addr#) (PrimArray ByteArray#
ba#) (I# Int#
soff#) (I# Int#
n#) =
    forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\ State# (PrimState m)
s# ->
        let s'# :: State# (PrimState m)
s'# = forall d.
ByteArray# -> Int# -> Addr# -> Int# -> State# d -> State# d
copyByteArrayToAddr# ByteArray#
ba# (Int#
soff# Int# -> Int# -> Int#
*# Int#
siz#) Addr#
addr# (Int#
n# Int# -> Int# -> Int#
*# Int#
siz#) State# (PrimState m)
s#
        in (# State# (PrimState m)
s'#, () #))
  where siz# :: Int#
siz# = forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)

-- | Copy a slice of a mutable primitive array to a pointer.
-- The offset and length are given in elements of type @a@.
-- This function assumes that the 'Prim' instance of @a@
-- agrees with the 'Storable' instance.
--
-- /Note:/ this function does not do bounds or overlap checking.
copyMutablePrimArrayToPtr :: forall m a. (PrimMonad m, Prim a)
  => Ptr a -- ^ destination pointer
  -> MutablePrimArray (PrimState m) a -- ^ source array
  -> Int -- ^ offset into source array
  -> Int -- ^ number of elements to copy
  -> m ()
{-# INLINE copyMutablePrimArrayToPtr #-}
copyMutablePrimArrayToPtr :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Ptr a -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArrayToPtr (Ptr Addr#
addr#) (MutablePrimArray MutableByteArray# (PrimState m)
mba#) (I# Int#
soff#) (I# Int#
n#) =
    forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\ State# (PrimState m)
s# ->
        let s'# :: State# (PrimState m)
s'# = forall d.
MutableByteArray# d
-> Int# -> Addr# -> Int# -> State# d -> State# d
copyMutableByteArrayToAddr# MutableByteArray# (PrimState m)
mba# (Int#
soff# Int# -> Int# -> Int#
*# Int#
siz#) Addr#
addr# (Int#
n# Int# -> Int# -> Int#
*# Int#
siz#) State# (PrimState m)
s#
        in (# State# (PrimState m)
s'#, () #))
  where siz# :: Int#
siz# = forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)

-- | Copy from a pointer to a mutable primitive array.
-- The offset and length are given in elements of type @a@.
-- This function assumes that the 'Prim' instance of @a@
-- agrees with the 'Storable' instance.
--
-- /Note:/ this function does not do bounds or overlap checking.
copyPtrToMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a -- ^ destination array
  -> Int -- ^ destination offset
  -> Ptr a -- ^ source pointer
  -> Int -- ^ number of elements
  -> m ()
{-# INLINE copyPtrToMutablePrimArray #-}
copyPtrToMutablePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> Ptr a -> Int -> m ()
copyPtrToMutablePrimArray (MutablePrimArray MutableByteArray# (PrimState m)
ba#) (I# Int#
doff#) (Ptr Addr#
addr#) (I# Int#
n#) =
  forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (forall d.
Addr#
-> MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
copyAddrToByteArray# Addr#
addr# MutableByteArray# (PrimState m)
ba# (Int#
doff# Int# -> Int# -> Int#
*# Int#
siz#) (Int#
n# Int# -> Int# -> Int#
*# Int#
siz#))
  where
  siz# :: Int#
siz# = forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)

-- | Fill a slice of a mutable primitive array with a value.
--
-- /Note:/ this function does not do bounds checking.
setPrimArray
  :: (Prim a, PrimMonad m)
  => MutablePrimArray (PrimState m) a -- ^ array to fill
  -> Int -- ^ offset into array
  -> Int -- ^ number of values to fill
  -> a -- ^ value to fill with
  -> m ()
{-# INLINE setPrimArray #-}
setPrimArray :: forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
setPrimArray (MutablePrimArray MutableByteArray# (PrimState m)
dst#) (I# Int#
doff#) (I# Int#
sz#) a
x
  = forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (forall a s.
Prim a =>
MutableByteArray# s -> Int# -> Int# -> a -> State# s -> State# s
PT.setByteArray# MutableByteArray# (PrimState m)
dst# Int#
doff# Int#
sz# a
x)

-- | Get the size of a mutable primitive array in elements. Unlike 'sizeofMutablePrimArray',
-- this function ensures sequencing in the presence of resizing.
getSizeofMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a -- ^ array
  -> m Int
{-# INLINE getSizeofMutablePrimArray #-}
#if __GLASGOW_HASKELL__ >= 801
getSizeofMutablePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray (MutablePrimArray MutableByteArray# (PrimState m)
arr#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\State# (PrimState m)
s# ->
      case forall d. MutableByteArray# d -> State# d -> (# State# d, Int# #)
getSizeofMutableByteArray# MutableByteArray# (PrimState m)
arr# State# (PrimState m)
s# of
        (# State# (PrimState m)
s'#, Int#
sz# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# (Int# -> Int# -> Int#
quotInt# Int#
sz# (forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a))) #)
    )
#else
-- On older GHCs, it is not possible to resize a byte array, so
-- this provides behavior consistent with the implementation for
-- newer GHCs.
getSizeofMutablePrimArray arr
  = return (sizeofMutablePrimArray arr)
#endif

-- | Size of the mutable primitive array in elements. This function shall not
-- be used on primitive arrays that are an argument to or a result of
-- 'resizeMutablePrimArray' or 'shrinkMutablePrimArray'.
--
-- This function is deprecated and will be removed.
sizeofMutablePrimArray :: forall s a. Prim a => MutablePrimArray s a -> Int
{-# INLINE sizeofMutablePrimArray #-}
{-# DEPRECATED sizeofMutablePrimArray "use getSizeofMutablePrimArray instead" #-}
sizeofMutablePrimArray :: forall s a. Prim a => MutablePrimArray s a -> Int
sizeofMutablePrimArray (MutablePrimArray MutableByteArray# s
arr#) =
  Int# -> Int
I# (Int# -> Int# -> Int#
quotInt# (forall d. MutableByteArray# d -> Int#
sizeofMutableByteArray# MutableByteArray# s
arr#) (forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)))

-- | Check if the two arrays refer to the same memory block.
sameMutablePrimArray :: MutablePrimArray s a -> MutablePrimArray s a -> Bool
{-# INLINE sameMutablePrimArray #-}
sameMutablePrimArray :: forall s a. MutablePrimArray s a -> MutablePrimArray s a -> Bool
sameMutablePrimArray (MutablePrimArray MutableByteArray# s
arr#) (MutablePrimArray MutableByteArray# s
brr#)
  = Int# -> Bool
isTrue# (forall d. MutableByteArray# d -> MutableByteArray# d -> Int#
sameMutableByteArray# MutableByteArray# s
arr# MutableByteArray# s
brr#)

-- | Create an immutable copy of a slice of a primitive array. The offset and
-- length are given in elements.
--
-- This operation makes a copy of the specified section, so it is safe to
-- continue using the mutable array afterward.
--
-- /Note:/ The provided array should contain the full subrange
-- specified by the two Ints, but this is not checked.
freezePrimArray
  :: (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a -- ^ source
  -> Int                              -- ^ offset in elements
  -> Int                              -- ^ length in elements
  -> m (PrimArray a)
{-# INLINE freezePrimArray #-}
freezePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> m (PrimArray a)
freezePrimArray !MutablePrimArray (PrimState m) a
src !Int
off !Int
len = do
  MutablePrimArray (PrimState m) a
dst <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray MutablePrimArray (PrimState m) a
dst Int
0 MutablePrimArray (PrimState m) a
src Int
off Int
len
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) a
dst

-- | Create a mutable primitive array from a slice of an immutable primitive array.
-- The offset and length are given in elements.
--
-- This operation makes a copy of the specified slice, so it is safe to
-- use the immutable array afterward.
--
-- /Note:/ The provided array should contain the full subrange
-- specified by the two Ints, but this is not checked.
--
-- @since 0.7.2.0
thawPrimArray
  :: (PrimMonad m, Prim a)
  => PrimArray a -- ^ source
  -> Int         -- ^ offset in elements
  -> Int         -- ^ length in elements
  -> m (MutablePrimArray (PrimState m) a)
{-# INLINE thawPrimArray #-}
thawPrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimArray a -> Int -> Int -> m (MutablePrimArray (PrimState m) a)
thawPrimArray !PrimArray a
src !Int
off !Int
len = do
  MutablePrimArray (PrimState m) a
dst <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray (PrimState m) a
dst Int
0 PrimArray a
src Int
off Int
len
  forall (m :: * -> *) a. Monad m => a -> m a
return MutablePrimArray (PrimState m) a
dst

-- | Convert a mutable primitive array to an immutable one without copying. The
-- array should not be modified after the conversion.
unsafeFreezePrimArray
  :: PrimMonad m => MutablePrimArray (PrimState m) a -> m (PrimArray a)
{-# INLINE unsafeFreezePrimArray #-}
unsafeFreezePrimArray :: forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray (MutablePrimArray MutableByteArray# (PrimState m)
arr#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\State# (PrimState m)
s# -> case forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# (PrimState m)
arr# State# (PrimState m)
s# of
                        (# State# (PrimState m)
s'#, ByteArray#
arr'# #) -> (# State# (PrimState m)
s'#, forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
arr'# #))

-- | Convert an immutable array to a mutable one without copying. The
-- original array should not be used after the conversion.
unsafeThawPrimArray
  :: PrimMonad m => PrimArray a -> m (MutablePrimArray (PrimState m) a)
{-# INLINE unsafeThawPrimArray #-}
unsafeThawPrimArray :: forall (m :: * -> *) a.
PrimMonad m =>
PrimArray a -> m (MutablePrimArray (PrimState m) a)
unsafeThawPrimArray (PrimArray ByteArray#
arr#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\State# (PrimState m)
s# -> (# State# (PrimState m)
s#, forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# ByteArray#
arr#) #))

-- | Read a primitive value from the primitive array.
--
-- /Note:/ this function does not do bounds checking.
indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
{-# INLINE indexPrimArray #-}
indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray (PrimArray ByteArray#
arr#) (I# Int#
i#) = forall a. Prim a => ByteArray# -> Int# -> a
indexByteArray# ByteArray#
arr# Int#
i#

-- | Get the size, in elements, of the primitive array.
sizeofPrimArray :: forall a. Prim a => PrimArray a -> Int
{-# INLINE sizeofPrimArray #-}
sizeofPrimArray :: forall a. Prim a => PrimArray a -> Int
sizeofPrimArray (PrimArray ByteArray#
arr#) = Int# -> Int
I# (Int# -> Int# -> Int#
quotInt# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
arr#) (forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)))

#if __GLASGOW_HASKELL__ >= 802
-- | Check whether or not the primitive array is pinned. Pinned primitive arrays cannot
-- be moved by the garbage collector. It is safe to use 'primArrayContents'
-- on such arrays. This function is only available when compiling with
-- GHC 8.2 or newer.
--
-- @since 0.7.1.0
isPrimArrayPinned :: PrimArray a -> Bool
{-# INLINE isPrimArrayPinned #-}
isPrimArrayPinned :: forall a. PrimArray a -> Bool
isPrimArrayPinned (PrimArray ByteArray#
arr#) = Int# -> Bool
isTrue# (ByteArray# -> Int#
Exts.isByteArrayPinned# ByteArray#
arr#)

-- | Check whether or not the mutable primitive array is pinned. This function is
-- only available when compiling with GHC 8.2 or newer.
--
-- @since 0.7.1.0
isMutablePrimArrayPinned :: MutablePrimArray s a -> Bool
{-# INLINE isMutablePrimArrayPinned #-}
isMutablePrimArrayPinned :: forall s a. MutablePrimArray s a -> Bool
isMutablePrimArrayPinned (MutablePrimArray MutableByteArray# s
marr#) = Int# -> Bool
isTrue# (forall d. MutableByteArray# d -> Int#
Exts.isMutableByteArrayPinned# MutableByteArray# s
marr#)
#endif

-- | Lazy right-associated fold over the elements of a 'PrimArray'.
{-# INLINE foldrPrimArray #-}
foldrPrimArray :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldrPrimArray :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldrPrimArray a -> b -> b
f b
z PrimArray a
arr = Int -> b
go Int
0
  where
    !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
    go :: Int -> b
go !Int
i
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
sz = a -> b -> b
f (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i) (Int -> b
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))
      | Bool
otherwise = b
z

-- | Strict right-associated fold over the elements of a 'PrimArray'.
{-# INLINE foldrPrimArray' #-}
foldrPrimArray' :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldrPrimArray' :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldrPrimArray' a -> b -> b
f b
z0 PrimArray a
arr = Int -> b -> b
go (forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr forall a. Num a => a -> a -> a
- Int
1) b
z0
  where
    go :: Int -> b -> b
go !Int
i !b
acc
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = b
acc
      | Bool
otherwise = Int -> b -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1) (a -> b -> b
f (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i) b
acc)

-- | Lazy left-associated fold over the elements of a 'PrimArray'.
{-# INLINE foldlPrimArray #-}
foldlPrimArray :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldlPrimArray :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldlPrimArray b -> a -> b
f b
z PrimArray a
arr = Int -> b
go (forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr forall a. Num a => a -> a -> a
- Int
1)
  where
    go :: Int -> b
go !Int
i
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = b
z
      | Bool
otherwise = b -> a -> b
f (Int -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1)) (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i)

-- | Strict left-associated fold over the elements of a 'PrimArray'.
{-# INLINE foldlPrimArray' #-}
foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldlPrimArray' b -> a -> b
f b
z0 PrimArray a
arr = Int -> b -> b
go Int
0 b
z0
  where
    !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
    go :: Int -> b -> b
go !Int
i !b
acc
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
sz = Int -> b -> b
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) (b -> a -> b
f b
acc (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i))
      | Bool
otherwise = b
acc

-- | Strict left-associated fold over the elements of a 'PrimArray'.
{-# INLINE foldlPrimArrayM' #-}
foldlPrimArrayM' :: (Prim a, Monad m) => (b -> a -> m b) -> b -> PrimArray a -> m b
foldlPrimArrayM' :: forall a (m :: * -> *) b.
(Prim a, Monad m) =>
(b -> a -> m b) -> b -> PrimArray a -> m b
foldlPrimArrayM' b -> a -> m b
f b
z0 PrimArray a
arr = Int -> b -> m b
go Int
0 b
z0
  where
    !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
    go :: Int -> b -> m b
go !Int
i !b
acc1
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
sz = do
          b
acc2 <- b -> a -> m b
f b
acc1 (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i)
          Int -> b -> m b
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) b
acc2
      | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return b
acc1

-- | Traverse a primitive array. The traversal forces the resulting values and
-- writes them to the new primitive array as it performs the monadic effects.
-- Consequently:
--
-- >>> traversePrimArrayP (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
-- 1
-- 2
-- *** Exception: Prelude.undefined
--
-- In many situations, 'traversePrimArrayP' can replace 'traversePrimArray',
-- changing the strictness characteristics of the traversal but typically improving
-- the performance. Consider the following short-circuiting traversal:
--
-- > incrPositiveA :: PrimArray Int -> Maybe (PrimArray Int)
-- > incrPositiveA xs = traversePrimArray (\x -> bool Nothing (Just (x + 1)) (x > 0)) xs
--
-- This can be rewritten using 'traversePrimArrayP'. To do this, we must
-- change the traversal context to @MaybeT (ST s)@, which has a 'PrimMonad'
-- instance:
--
-- > incrPositiveB :: PrimArray Int -> Maybe (PrimArray Int)
-- > incrPositiveB xs = runST $ runMaybeT $ traversePrimArrayP
-- >   (\x -> bool (MaybeT (return Nothing)) (MaybeT (return (Just (x + 1)))) (x > 0))
-- >   xs
--
-- Benchmarks demonstrate that the second implementation runs 150 times
-- faster than the first. It also results in fewer allocations.
{-# INLINE traversePrimArrayP #-}
traversePrimArrayP :: (PrimMonad m, Prim a, Prim b)
  => (a -> m b)
  -> PrimArray a
  -> m (PrimArray b)
traversePrimArrayP :: forall (m :: * -> *) a b.
(PrimMonad m, Prim a, Prim b) =>
(a -> m b) -> PrimArray a -> m (PrimArray b)
traversePrimArrayP a -> m b
f PrimArray a
arr = do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray (PrimState m) b
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> m ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$ do
        b
b <- a -> m b
f (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ix)
        forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) b
marr Int
ix b
b
        Int -> m ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) b
marr

-- | Filter the primitive array, keeping the elements for which the monadic
-- predicate evaluates to true.
{-# INLINE filterPrimArrayP #-}
filterPrimArrayP :: (PrimMonad m, Prim a)
  => (a -> m Bool)
  -> PrimArray a
  -> m (PrimArray a)
filterPrimArrayP :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
(a -> m Bool) -> PrimArray a -> m (PrimArray a)
filterPrimArrayP a -> m Bool
f PrimArray a
arr = do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray (PrimState m) a
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> Int -> m Int
go !Int
ixSrc !Int
ixDst = if Int
ixSrc forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          let a :: a
a = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ixSrc
          Bool
b <- a -> m Bool
f a
a
          if Bool
b
            then do
              forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
marr Int
ixDst a
a
              Int -> Int -> m Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) (Int
ixDst forall a. Num a => a -> a -> a
+ Int
1)
            else Int -> Int -> m Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) Int
ixDst
        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
ixDst
  Int
lenDst <- Int -> Int -> m Int
go Int
0 Int
0
  MutablePrimArray (PrimState m) a
marr' <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray (PrimState m) a
marr Int
lenDst
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) a
marr'

-- | Map over the primitive array, keeping the elements for which the monadic
-- predicate provides a 'Just'.
{-# INLINE mapMaybePrimArrayP #-}
mapMaybePrimArrayP :: (PrimMonad m, Prim a, Prim b)
  => (a -> m (Maybe b))
  -> PrimArray a
  -> m (PrimArray b)
mapMaybePrimArrayP :: forall (m :: * -> *) a b.
(PrimMonad m, Prim a, Prim b) =>
(a -> m (Maybe b)) -> PrimArray a -> m (PrimArray b)
mapMaybePrimArrayP a -> m (Maybe b)
f PrimArray a
arr = do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray (PrimState m) b
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> Int -> m Int
go !Int
ixSrc !Int
ixDst = if Int
ixSrc forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          let a :: a
a = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ixSrc
          Maybe b
mb <- a -> m (Maybe b)
f a
a
          case Maybe b
mb of
            Just b
b -> do
              forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) b
marr Int
ixDst b
b
              Int -> Int -> m Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) (Int
ixDst forall a. Num a => a -> a -> a
+ Int
1)
            Maybe b
Nothing -> Int -> Int -> m Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) Int
ixDst
        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
ixDst
  Int
lenDst <- Int -> Int -> m Int
go Int
0 Int
0
  MutablePrimArray (PrimState m) b
marr' <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray (PrimState m) b
marr Int
lenDst
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) b
marr'

-- | Generate a primitive array by evaluating the monadic generator function
-- at each index.
{-# INLINE generatePrimArrayP #-}
generatePrimArrayP :: (PrimMonad m, Prim a)
  => Int -- ^ length
  -> (Int -> m a) -- ^ generator
  -> m (PrimArray a)
generatePrimArrayP :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> (Int -> m a) -> m (PrimArray a)
generatePrimArrayP Int
sz Int -> m a
f = do
  MutablePrimArray (PrimState m) a
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> m ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$ do
        a
b <- Int -> m a
f Int
ix
        forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
marr Int
ix a
b
        Int -> m ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) a
marr

-- | Execute the monadic action the given number of times and store the
-- results in a primitive array.
{-# INLINE replicatePrimArrayP #-}
replicatePrimArrayP :: (PrimMonad m, Prim a)
  => Int
  -> m a
  -> m (PrimArray a)
replicatePrimArrayP :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m a -> m (PrimArray a)
replicatePrimArrayP Int
sz m a
f = do
  MutablePrimArray (PrimState m) a
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> m ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$ do
        a
b <- m a
f
        forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
marr Int
ix a
b
        Int -> m ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) a
marr

-- | Map over the elements of a primitive array.
{-# INLINE mapPrimArray #-}
mapPrimArray :: (Prim a, Prim b)
  => (a -> b)
  -> PrimArray a
  -> PrimArray b
mapPrimArray :: forall a b.
(Prim a, Prim b) =>
(a -> b) -> PrimArray a -> PrimArray b
mapPrimArray a -> b
f PrimArray a
arr = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray s b
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> ST s ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$ do
        let b :: b
b = a -> b
f (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ix)
        forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s b
marr Int
ix b
b
        Int -> ST s ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go Int
0
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s b
marr

-- | Indexed map over the elements of a primitive array.
{-# INLINE imapPrimArray #-}
imapPrimArray :: (Prim a, Prim b)
  => (Int -> a -> b)
  -> PrimArray a
  -> PrimArray b
imapPrimArray :: forall a b.
(Prim a, Prim b) =>
(Int -> a -> b) -> PrimArray a -> PrimArray b
imapPrimArray Int -> a -> b
f PrimArray a
arr = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray s b
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> ST s ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$ do
        let b :: b
b = Int -> a -> b
f Int
ix (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ix)
        forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s b
marr Int
ix b
b
        Int -> ST s ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go Int
0
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s b
marr

-- | Filter elements of a primitive array according to a predicate.
{-# INLINE filterPrimArray #-}
filterPrimArray :: Prim a
  => (a -> Bool)
  -> PrimArray a
  -> PrimArray a
filterPrimArray :: forall a. Prim a => (a -> Bool) -> PrimArray a -> PrimArray a
filterPrimArray a -> Bool
p PrimArray a
arr = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray s a
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> Int -> ST s Int
go !Int
ixSrc !Int
ixDst = if Int
ixSrc forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          let !a :: a
a = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ixSrc
          if a -> Bool
p a
a
            then do
              forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s a
marr Int
ixDst a
a
              Int -> Int -> ST s Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) (Int
ixDst forall a. Num a => a -> a -> a
+ Int
1)
            else Int -> Int -> ST s Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) Int
ixDst
        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
ixDst
  Int
dstLen <- Int -> Int -> ST s Int
go Int
0 Int
0
  MutablePrimArray s a
marr' <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray s a
marr Int
dstLen
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
marr'

-- | Filter the primitive array, keeping the elements for which the monadic
-- predicate evaluates true.
filterPrimArrayA
  :: (Applicative f, Prim a)
  => (a -> f Bool) -- ^ mapping function
  -> PrimArray a -- ^ primitive array
  -> f (PrimArray a)
filterPrimArrayA :: forall (f :: * -> *) a.
(Applicative f, Prim a) =>
(a -> f Bool) -> PrimArray a -> f (PrimArray a)
filterPrimArrayA a -> f Bool
f = \ !PrimArray a
ary ->
  let
    !len :: Int
len = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
ary
    go :: Int -> f (IxSTA a)
go !Int
ixSrc
      | Int
ixSrc forall a. Eq a => a -> a -> Bool
== Int
len = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. Int -> MutableByteArray# s -> ST s Int) -> IxSTA a
IxSTA forall a b. (a -> b) -> a -> b
$ \Int
ixDst MutableByteArray# s
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Int
ixDst
      | Bool
otherwise = let x :: a
x = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
ary Int
ixSrc in
          forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
            (\Bool
keep (IxSTA forall s. Int -> MutableByteArray# s -> ST s Int
m) -> forall a.
(forall s. Int -> MutableByteArray# s -> ST s Int) -> IxSTA a
IxSTA forall a b. (a -> b) -> a -> b
$ \Int
ixDst MutableByteArray# s
mary -> if Bool
keep
              then forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary) Int
ixDst a
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. Int -> MutableByteArray# s -> ST s Int
m (Int
ixDst forall a. Num a => a -> a -> a
+ Int
1) MutableByteArray# s
mary
              else forall s. Int -> MutableByteArray# s -> ST s Int
m Int
ixDst MutableByteArray# s
mary
            )
            (a -> f Bool
f a
x)
            (Int -> f (IxSTA a)
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1))
  in if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
     then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. PrimArray a
emptyPrimArray
     else forall a. Prim a => Int -> IxSTA a -> PrimArray a
runIxSTA Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (IxSTA a)
go Int
0

-- | Map over the primitive array, keeping the elements for which the applicative
-- predicate provides a 'Just'.
mapMaybePrimArrayA
  :: (Applicative f, Prim a, Prim b)
  => (a -> f (Maybe b)) -- ^ mapping function
  -> PrimArray a -- ^ primitive array
  -> f (PrimArray b)
mapMaybePrimArrayA :: forall (f :: * -> *) a b.
(Applicative f, Prim a, Prim b) =>
(a -> f (Maybe b)) -> PrimArray a -> f (PrimArray b)
mapMaybePrimArrayA a -> f (Maybe b)
f = \ !PrimArray a
ary ->
  let
    !len :: Int
len = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
ary
    go :: Int -> f (IxSTA b)
go !Int
ixSrc
      | Int
ixSrc forall a. Eq a => a -> a -> Bool
== Int
len = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. Int -> MutableByteArray# s -> ST s Int) -> IxSTA a
IxSTA forall a b. (a -> b) -> a -> b
$ \Int
ixDst MutableByteArray# s
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Int
ixDst
      | Bool
otherwise = let x :: a
x = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
ary Int
ixSrc in
          forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
            (\Maybe b
mb (IxSTA forall s. Int -> MutableByteArray# s -> ST s Int
m) -> forall a.
(forall s. Int -> MutableByteArray# s -> ST s Int) -> IxSTA a
IxSTA forall a b. (a -> b) -> a -> b
$ \Int
ixDst MutableByteArray# s
mary -> case Maybe b
mb of
              Just b
b -> forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary) Int
ixDst b
b forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. Int -> MutableByteArray# s -> ST s Int
m (Int
ixDst forall a. Num a => a -> a -> a
+ Int
1) MutableByteArray# s
mary
              Maybe b
Nothing -> forall s. Int -> MutableByteArray# s -> ST s Int
m Int
ixDst MutableByteArray# s
mary
            )
            (a -> f (Maybe b)
f a
x)
            (Int -> f (IxSTA b)
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1))
  in if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
     then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. PrimArray a
emptyPrimArray
     else forall a. Prim a => Int -> IxSTA a -> PrimArray a
runIxSTA Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (IxSTA b)
go Int
0

-- | Map over a primitive array, optionally discarding some elements. This
--   has the same behavior as @Data.Maybe.mapMaybe@.
{-# INLINE mapMaybePrimArray #-}
mapMaybePrimArray :: (Prim a, Prim b)
  => (a -> Maybe b)
  -> PrimArray a
  -> PrimArray b
mapMaybePrimArray :: forall a b.
(Prim a, Prim b) =>
(a -> Maybe b) -> PrimArray a -> PrimArray b
mapMaybePrimArray a -> Maybe b
p PrimArray a
arr = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray s b
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> Int -> ST s Int
go !Int
ixSrc !Int
ixDst = if Int
ixSrc forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          let !a :: a
a = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ixSrc
          case a -> Maybe b
p a
a of
            Just b
b -> do
              forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s b
marr Int
ixDst b
b
              Int -> Int -> ST s Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) (Int
ixDst forall a. Num a => a -> a -> a
+ Int
1)
            Maybe b
Nothing -> Int -> Int -> ST s Int
go (Int
ixSrc forall a. Num a => a -> a -> a
+ Int
1) Int
ixDst
        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
ixDst
  Int
dstLen <- Int -> Int -> ST s Int
go Int
0 Int
0
  MutablePrimArray s b
marr' <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray s b
marr Int
dstLen
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s b
marr'

-- | Traverse a primitive array. The traversal performs all of the applicative
-- effects /before/ forcing the resulting values and writing them to the new
-- primitive array. Consequently:
--
-- >>> traversePrimArray (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
-- 1
-- 2
-- 3
-- *** Exception: Prelude.undefined
--
-- The function 'traversePrimArrayP' always outperforms this function, but it
-- requires a 'PrimMonad' constraint, and it forces the values as
-- it performs the effects.
traversePrimArray
  :: (Applicative f, Prim a, Prim b)
  => (a -> f b) -- ^ mapping function
  -> PrimArray a -- ^ primitive array
  -> f (PrimArray b)
traversePrimArray :: forall (f :: * -> *) a b.
(Applicative f, Prim a, Prim b) =>
(a -> f b) -> PrimArray a -> f (PrimArray b)
traversePrimArray a -> f b
f = \ !PrimArray a
ary ->
  let
    !len :: Int
len = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
ary
    go :: Int -> f (STA b)
go !Int
i
      | Int
i forall a. Eq a => a -> a -> Bool
== Int
len = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary -> forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary)
      | a
x <- forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
ary Int
i
      = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\b
b (STA forall s. MutableByteArray# s -> ST s (PrimArray b)
m) -> forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary ->
                  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary) Int
i b
b forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. MutableByteArray# s -> ST s (PrimArray b)
m MutableByteArray# s
mary)
               (a -> f b
f a
x) (Int -> f (STA b)
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))
  in if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
     then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. PrimArray a
emptyPrimArray
     else forall a. Prim a => Int -> STA a -> PrimArray a
runSTA Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA b)
go Int
0

-- | Traverse a primitive array with the index of each element.
itraversePrimArray
  :: (Applicative f, Prim a, Prim b)
  => (Int -> a -> f b)
  -> PrimArray a
  -> f (PrimArray b)
itraversePrimArray :: forall (f :: * -> *) a b.
(Applicative f, Prim a, Prim b) =>
(Int -> a -> f b) -> PrimArray a -> f (PrimArray b)
itraversePrimArray Int -> a -> f b
f = \ !PrimArray a
ary ->
  let
    !len :: Int
len = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
ary
    go :: Int -> f (STA b)
go !Int
i
      | Int
i forall a. Eq a => a -> a -> Bool
== Int
len = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary -> forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary)
      | a
x <- forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
ary Int
i
      = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\b
b (STA forall s. MutableByteArray# s -> ST s (PrimArray b)
m) -> forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary ->
                  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary) Int
i b
b forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. MutableByteArray# s -> ST s (PrimArray b)
m MutableByteArray# s
mary)
               (Int -> a -> f b
f Int
i a
x) (Int -> f (STA b)
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))
  in if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
     then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. PrimArray a
emptyPrimArray
     else forall a. Prim a => Int -> STA a -> PrimArray a
runSTA Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA b)
go Int
0

-- | Traverse a primitive array with the indices. The traversal forces the
-- resulting values and writes them to the new primitive array as it performs
-- the monadic effects.
{-# INLINE itraversePrimArrayP #-}
itraversePrimArrayP :: (Prim a, Prim b, PrimMonad m)
  => (Int -> a -> m b)
  -> PrimArray a
  -> m (PrimArray b)
itraversePrimArrayP :: forall a b (m :: * -> *).
(Prim a, Prim b, PrimMonad m) =>
(Int -> a -> m b) -> PrimArray a -> m (PrimArray b)
itraversePrimArrayP Int -> a -> m b
f PrimArray a
arr = do
  let !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr
  MutablePrimArray (PrimState m) b
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go :: Int -> m ()
go !Int
ix
        | Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz = do
            forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) b
marr Int
ix forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> a -> m b
f Int
ix (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
ix)
            Int -> m ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return ()
  Int -> m ()
go Int
0
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) b
marr

-- | Generate a primitive array.
{-# INLINE generatePrimArray #-}
generatePrimArray :: Prim a
  => Int -- ^ length
  -> (Int -> a) -- ^ element from index
  -> PrimArray a
generatePrimArray :: forall a. Prim a => Int -> (Int -> a) -> PrimArray a
generatePrimArray Int
len Int -> a
f = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MutablePrimArray s a
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  let go :: Int -> ST s ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
        forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s a
marr Int
ix (Int -> a
f Int
ix)
        Int -> ST s ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go Int
0
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
marr

-- | Create a primitive array by copying the element the given
-- number of times.
{-# INLINE replicatePrimArray #-}
replicatePrimArray :: Prim a
  => Int -- ^ length
  -> a -- ^ element
  -> PrimArray a
replicatePrimArray :: forall a. Prim a => Int -> a -> PrimArray a
replicatePrimArray Int
len a
a = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MutablePrimArray s a
marr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
setPrimArray MutablePrimArray s a
marr Int
0 Int
len a
a
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
marr

-- | Generate a primitive array by evaluating the applicative generator
-- function at each index.
{-# INLINE generatePrimArrayA #-}
generatePrimArrayA
  :: (Applicative f, Prim a)
  => Int -- ^ length
  -> (Int -> f a) -- ^ element from index
  -> f (PrimArray a)
generatePrimArrayA :: forall (f :: * -> *) a.
(Applicative f, Prim a) =>
Int -> (Int -> f a) -> f (PrimArray a)
generatePrimArrayA Int
len Int -> f a
f =
  let
    go :: Int -> f (STA a)
go !Int
i
      | Int
i forall a. Eq a => a -> a -> Bool
== Int
len = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary -> forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary)
      | Bool
otherwise
      = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\a
b (STA forall s. MutableByteArray# s -> ST s (PrimArray a)
m) -> forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary ->
                  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary) Int
i a
b forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. MutableByteArray# s -> ST s (PrimArray a)
m MutableByteArray# s
mary)
               (Int -> f a
f Int
i) (Int -> f (STA a)
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))
  in if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
     then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. PrimArray a
emptyPrimArray
     else forall a. Prim a => Int -> STA a -> PrimArray a
runSTA Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA a)
go Int
0

-- | Execute the applicative action the given number of times and store the
-- results in a 'PrimArray'.
{-# INLINE replicatePrimArrayA #-}
replicatePrimArrayA
  :: (Applicative f, Prim a)
  => Int -- ^ length
  -> f a -- ^ applicative element producer
  -> f (PrimArray a)
replicatePrimArrayA :: forall (f :: * -> *) a.
(Applicative f, Prim a) =>
Int -> f a -> f (PrimArray a)
replicatePrimArrayA Int
len f a
f =
  let
    go :: Int -> f (STA a)
go !Int
i
      | Int
i forall a. Eq a => a -> a -> Bool
== Int
len = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary -> forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary)
      | Bool
otherwise
      = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\a
b (STA forall s. MutableByteArray# s -> ST s (PrimArray a)
m) -> forall a.
(forall s. MutableByteArray# s -> ST s (PrimArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \MutableByteArray# s
mary ->
                  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# s
mary) Int
i a
b forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. MutableByteArray# s -> ST s (PrimArray a)
m MutableByteArray# s
mary)
               f a
f (Int -> f (STA a)
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))
  in if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
     then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. PrimArray a
emptyPrimArray
     else forall a. Prim a => Int -> STA a -> PrimArray a
runSTA Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA a)
go Int
0

-- | Traverse the primitive array, discarding the results. There
-- is no 'PrimMonad' variant of this function, since it would not provide
-- any performance benefit.
traversePrimArray_
  :: (Applicative f, Prim a)
  => (a -> f b)
  -> PrimArray a
  -> f ()
traversePrimArray_ :: forall (f :: * -> *) a b.
(Applicative f, Prim a) =>
(a -> f b) -> PrimArray a -> f ()
traversePrimArray_ a -> f b
f PrimArray a
a = Int -> f ()
go Int
0 where
  !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
a
  go :: Int -> f ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$
    a -> f b
f (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
a Int
ix) forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> f ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)

-- | Traverse the primitive array with the indices, discarding the results.
-- There is no 'PrimMonad' variant of this function, since it would not
-- provide any performance benefit.
itraversePrimArray_
  :: (Applicative f, Prim a)
  => (Int -> a -> f b)
  -> PrimArray a
  -> f ()
itraversePrimArray_ :: forall (f :: * -> *) a b.
(Applicative f, Prim a) =>
(Int -> a -> f b) -> PrimArray a -> f ()
itraversePrimArray_ Int -> a -> f b
f PrimArray a
a = Int -> f ()
go Int
0 where
  !sz :: Int
sz = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
a
  go :: Int -> f ()
go !Int
ix = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$
    Int -> a -> f b
f Int
ix (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
a Int
ix) forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> f ()
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1)

newtype IxSTA a = IxSTA {forall a.
IxSTA a -> forall s. Int -> MutableByteArray# s -> ST s Int
_runIxSTA :: forall s. Int -> MutableByteArray# s -> ST s Int}

runIxSTA :: forall a. Prim a
  => Int -- maximum possible size
  -> IxSTA a
  -> PrimArray a
runIxSTA :: forall a. Prim a => Int -> IxSTA a -> PrimArray a
runIxSTA !Int
szUpper = \ (IxSTA forall s. Int -> MutableByteArray# s -> ST s Int
m) -> forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MutablePrimArray s a
ar :: MutablePrimArray s a <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
szUpper
  Int
sz <- forall s. Int -> MutableByteArray# s -> ST s Int
m Int
0 (forall s a. MutablePrimArray s a -> MutableByteArray# s
unMutablePrimArray MutablePrimArray s a
ar)
  MutablePrimArray s a
ar' <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray s a
ar Int
sz
  forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
ar'
{-# INLINE runIxSTA #-}

newtype STA a = STA {forall a.
STA a -> forall s. MutableByteArray# s -> ST s (PrimArray a)
_runSTA :: forall s. MutableByteArray# s -> ST s (PrimArray a)}

runSTA :: forall a. Prim a => Int -> STA a -> PrimArray a
runSTA :: forall a. Prim a => Int -> STA a -> PrimArray a
runSTA !Int
sz = \ (STA forall s. MutableByteArray# s -> ST s (PrimArray a)
m) -> forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ (MutablePrimArray s a
ar :: MutablePrimArray s a) -> forall s. MutableByteArray# s -> ST s (PrimArray a)
m (forall s a. MutablePrimArray s a -> MutableByteArray# s
unMutablePrimArray MutablePrimArray s a
ar)
{-# INLINE runSTA #-}

unMutablePrimArray :: MutablePrimArray s a -> MutableByteArray# s
unMutablePrimArray :: forall s a. MutablePrimArray s a -> MutableByteArray# s
unMutablePrimArray (MutablePrimArray MutableByteArray# s
m) = MutableByteArray# s
m

{- $effectfulMapCreate
The naming conventions adopted in this section are explained in the
documentation of the @Data.Primitive@ module.
-}

-- | Create a /pinned/ primitive array of the specified size (in elements). The garbage
-- collector is guaranteed not to move it. The underlying memory is left uninitialized.
--
-- @since 0.7.1.0
newPinnedPrimArray :: forall m a. (PrimMonad m, Prim a)
  => Int -> m (MutablePrimArray (PrimState m) a)
{-# INLINE newPinnedPrimArray #-}
newPinnedPrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPinnedPrimArray (I# Int#
n#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\State# (PrimState m)
s# -> case forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newPinnedByteArray# (Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) State# (PrimState m)
s# of
                        (# State# (PrimState m)
s'#, MutableByteArray# (PrimState m)
arr# #) -> (# State# (PrimState m)
s'#, forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# (PrimState m)
arr# #))

-- | Create a /pinned/ primitive array of the specified size (in elements) and
-- with the alignment given by its 'Prim' instance. The garbage collector is
-- guaranteed not to move it. The underlying memory is left uninitialized.
--
-- @since 0.7.0.0
newAlignedPinnedPrimArray :: forall m a. (PrimMonad m, Prim a)
  => Int -> m (MutablePrimArray (PrimState m) a)
{-# INLINE newAlignedPinnedPrimArray #-}
newAlignedPinnedPrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newAlignedPinnedPrimArray (I# Int#
n#)
  = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (\State# (PrimState m)
s# -> case forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# (Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => Proxy a -> Int#
sizeOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) (forall a. Prim a => Proxy a -> Int#
alignmentOfType# (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) State# (PrimState m)
s# of
                        (# State# (PrimState m)
s'#, MutableByteArray# (PrimState m)
arr# #) -> (# State# (PrimState m)
s'#, forall s a. MutableByteArray# s -> MutablePrimArray s a
MutablePrimArray MutableByteArray# (PrimState m)
arr# #))

-- | Yield a pointer to the array's data. This operation is only safe on
-- /pinned/ prim arrays allocated by 'newPinnedByteArray' or
-- 'newAlignedPinnedByteArray'.
--
-- @since 0.7.1.0
primArrayContents :: PrimArray a -> Ptr a
{-# INLINE primArrayContents #-}
primArrayContents :: forall a. PrimArray a -> Ptr a
primArrayContents (PrimArray ByteArray#
arr#) = forall a. Addr# -> Ptr a
Ptr (ByteArray# -> Addr#
byteArrayContents# ByteArray#
arr#)

-- | Yield a pointer to the array's data. This operation is only safe on
-- /pinned/ byte arrays allocated by 'newPinnedByteArray' or
-- 'newAlignedPinnedByteArray'.
--
-- @since 0.7.1.0
mutablePrimArrayContents :: MutablePrimArray s a -> Ptr a
{-# INLINE mutablePrimArrayContents #-}
mutablePrimArrayContents :: forall s a. MutablePrimArray s a -> Ptr a
mutablePrimArrayContents (MutablePrimArray MutableByteArray# s
arr#) =
  forall a. Addr# -> Ptr a
Ptr (forall s. MutableByteArray# s -> Addr#
mutableByteArrayContentsShim MutableByteArray# s
arr#)

-- | Return a newly allocated array with the specified subrange of the
-- provided array. The provided array should contain the full subrange
-- specified by the two Ints, but this is not checked.
clonePrimArray :: Prim a
  => PrimArray a -- ^ source array
  -> Int     -- ^ offset into destination array
  -> Int     -- ^ number of elements to copy
  -> PrimArray a
{-# INLINE clonePrimArray #-}
clonePrimArray :: forall a. Prim a => PrimArray a -> Int -> Int -> PrimArray a
clonePrimArray PrimArray a
src Int
off Int
n = forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray forall a b. (a -> b) -> a -> b
$ do
  MutablePrimArray s a
dst <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
n
  forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
dst Int
0 PrimArray a
src Int
off Int
n
  forall (m :: * -> *) a. Monad m => a -> m a
return MutablePrimArray s a
dst

-- | Return a newly allocated mutable array with the specified subrange of
-- the provided mutable array. The provided mutable array should contain the
-- full subrange specified by the two Ints, but this is not checked.
cloneMutablePrimArray :: (PrimMonad m, Prim a)
  => MutablePrimArray (PrimState m) a -- ^ source array
  -> Int -- ^ offset into destination array
  -> Int -- ^ number of elements to copy
  -> m (MutablePrimArray (PrimState m) a)
{-# INLINE cloneMutablePrimArray #-}
cloneMutablePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArray MutablePrimArray (PrimState m) a
src Int
off Int
n = do
  MutablePrimArray (PrimState m) a
dst <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
n
  forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray MutablePrimArray (PrimState m) a
dst Int
0 MutablePrimArray (PrimState m) a
src Int
off Int
n
  forall (m :: * -> *) a. Monad m => a -> m a
return MutablePrimArray (PrimState m) a
dst

-- | Execute the monadic action and freeze the resulting array.
--
-- > runPrimArray m = runST $ m >>= unsafeFreezePrimArray
runPrimArray
  :: (forall s. ST s (MutablePrimArray s a))
  -> PrimArray a
#if MIN_VERSION_base(4,10,0) /* In new GHCs, runRW# is available. */
runPrimArray :: forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray forall s. ST s (MutablePrimArray s a)
m = forall a. ByteArray# -> PrimArray a
PrimArray (forall a. (forall s. ST s (MutablePrimArray s a)) -> ByteArray#
runPrimArray# forall s. ST s (MutablePrimArray s a)
m)

runPrimArray#
  :: (forall s. ST s (MutablePrimArray s a))
  -> ByteArray#
runPrimArray# :: forall a. (forall s. ST s (MutablePrimArray s a)) -> ByteArray#
runPrimArray# forall s. ST s (MutablePrimArray s a)
m = case forall o. (State# RealWorld -> o) -> o
runRW# forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case forall s a. ST s a -> State# s -> (# State# s, a #)
unST forall s. ST s (MutablePrimArray s a)
m State# RealWorld
s of { (# State# RealWorld
s', MutablePrimArray MutableByteArray# RealWorld
mary# #) ->
  forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# RealWorld
mary# State# RealWorld
s'} of (# State# RealWorld
_, ByteArray#
ary# #) -> ByteArray#
ary#

unST :: ST s a -> State# s -> (# State# s, a #)
unST :: forall s a. ST s a -> State# s -> (# State# s, a #)
unST (GHCST.ST STRep s a
f) = STRep s a
f
#else /* In older GHCs, runRW# is not available. */
runPrimArray m = runST $ m >>= unsafeFreezePrimArray
#endif

-- | A composition of 'primArrayContents' and 'keepAliveUnlifted'.
-- The callback function must not return the pointer. The argument
-- array must be /pinned/. See 'primArrayContents' for an explanation
-- of which primitive arrays are pinned.
--
-- Note: This could be implemented with 'keepAlive' instead of
-- 'keepAliveUnlifted', but 'keepAlive' here would cause GHC to materialize
-- the wrapper data constructor on the heap.
withPrimArrayContents :: PrimBase m => PrimArray a -> (Ptr a -> m a) -> m a
{-# INLINE withPrimArrayContents #-}
withPrimArrayContents :: forall (m :: * -> *) a.
PrimBase m =>
PrimArray a -> (Ptr a -> m a) -> m a
withPrimArrayContents (PrimArray ByteArray#
arr#) Ptr a -> m a
f =
  forall (m :: * -> *) (a :: UnliftedType) r.
PrimBase m =>
a -> m r -> m r
keepAliveUnlifted ByteArray#
arr# (Ptr a -> m a
f (forall a. Addr# -> Ptr a
Ptr (ByteArray# -> Addr#
byteArrayContents# ByteArray#
arr#)))

-- | A composition of 'mutablePrimArrayContents' and 'keepAliveUnlifted'.
-- The callback function must not return the pointer. The argument
-- array must be /pinned/. See 'primArrayContents' for an explanation
-- of which primitive arrays are pinned.
withMutablePrimArrayContents :: PrimBase m => MutablePrimArray (PrimState m) a -> (Ptr a -> m a) -> m a
{-# INLINE withMutablePrimArrayContents #-}
withMutablePrimArrayContents :: forall (m :: * -> *) a.
PrimBase m =>
MutablePrimArray (PrimState m) a -> (Ptr a -> m a) -> m a
withMutablePrimArrayContents (MutablePrimArray MutableByteArray# (PrimState m)
arr#) Ptr a -> m a
f =
  forall (m :: * -> *) (a :: UnliftedType) r.
PrimBase m =>
a -> m r -> m r
keepAliveUnlifted MutableByteArray# (PrimState m)
arr# (Ptr a -> m a
f (forall a. Addr# -> Ptr a
Ptr (forall s. MutableByteArray# s -> Addr#
mutableByteArrayContentsShim MutableByteArray# (PrimState m)
arr#)))