{-# LANGUAGE CPP #-}
{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module      : Streamly.Internal.Data.MutArray.Type
-- Copyright   : (c) 2020 Composewell Technologies
-- License     : BSD3-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Pinned and unpinned mutable array for 'Unboxed' types. Fulfils the following
-- goals:
--
-- * Random access (array)
-- * Efficient storage (unboxed)
-- * Performance (unboxed access)
-- * Performance - in-place operations (mutable)
-- * Performance - GC (pinned, mutable)
-- * interfacing with OS (pinned)
--
-- Stream and Fold APIs allow easy, efficient and convenient operations on
-- arrays.
--
-- Mutable arrays and file system files are quite similar, they can grow and
-- their content is mutable. Therefore, both have similar APIs as well. We
-- strive to keep the API consistent for both. Ideally, you should be able to
-- replace one with another with little changes to the code.

module Streamly.Internal.Data.MutArray.Type
    (
    -- ** Type
    -- $arrayNotes
      MutArray (..)
    , MutByteArray
    , MutableByteArray
    , pin
    , unpin
    , isPinned

    -- -- * Constructing and Writing
    -- ** Construction
    , nil

    -- *** Uninitialized Arrays
    , pinnedNew
    , pinnedNewBytes
    , pinnedNewAligned
    , new
    , newArrayWith

    -- *** From streams
    , ArrayUnsafe (..)
    , writeNWithUnsafe
    , writeNWith
    , writeNUnsafe
    , pinnedWriteNUnsafe
    , writeN
    , pinnedWriteN
    , pinnedWriteNAligned

    , writeWith
    , write
    , pinnedWrite

    , writeRevN
    -- , writeRev

    -- *** From containers
    , fromListN
    , pinnedFromListN
    , fromList
    , pinnedFromList
    , fromListRevN
    , fromListRev
    , fromStreamDN
    , fromStreamD
    , fromPureStream

    -- ** Random writes
    , putIndex
    , putIndexUnsafe
    , putIndices
    -- , putFromThenTo
    -- , putFrom -- start writing at the given position
    -- , putUpto -- write from beginning up to the given position
    -- , putFromTo
    -- , putFromRev
    -- , putUptoRev
    , modifyIndexUnsafe
    , modifyIndex
    , modifyIndices
    , modify
    , swapIndices
    , unsafeSwapIndices

    -- ** Growing and Shrinking
    -- | Arrays grow only at the end, though it is possible to grow on both sides
    -- and therefore have a cons as well as snoc. But that will require two
    -- bounds in the array representation.

    -- *** Appending elements
    , snocWith
    , snoc
    , snocLinear
    , snocMay
    , snocUnsafe

    -- *** Appending streams
    , writeAppendNUnsafe
    , writeAppendN
    , writeAppendWith
    , writeAppend

    -- ** Eliminating and Reading

    -- *** To streams
    , reader
    , readerRevWith
    , readerRev

    -- *** To containers
    , toStreamDWith
    , toStreamDRevWith
    , toStreamKWith
    , toStreamKRevWith
    , read
    , readRev
    , toStreamK
    , toStreamKRev
    , toList

    -- experimental
    , producerWith
    , producer

    -- *** Random reads
    , getIndex
    , getIndexUnsafe
    , getIndices
    , getIndicesD
    -- , getFromThenTo
    , getIndexRev

    -- ** Memory Management
    , blockSize
    , arrayChunkBytes
    , allocBytesToElemCount
    , realloc
    , resize
    , resizeExp
    , rightSize

    -- ** Size
    , length
    , byteLength
    -- , capacity
    , byteCapacity
    , bytesFree

    -- ** In-place Mutation Algorithms
    , strip
    , reverse
    , permute
    , partitionBy
    , shuffleBy
    , divideBy
    , mergeBy
    , bubble

    -- ** Casting
    , cast
    , castUnsafe
    , asBytes
    , asPtrUnsafe

    -- ** Folding
    , foldl'
    , foldr
    , cmp

    -- ** Arrays of arrays
    --  We can add dimensionality parameter to the array type to get
    --  multidimensional arrays. Multidimensional arrays would just be a
    --  convenience wrapper on top of single dimensional arrays.

    -- | Operations dealing with multiple arrays, streams of arrays or
    -- multidimensional array representations.

    -- *** Construct from streams
    , chunksOf
    , pinnedChunksOf
    , writeChunks

    -- *** Eliminate to streams
    , flattenArrays
    , flattenArraysRev
    , fromArrayStreamK

    -- *** Construct from arrays
    -- get chunks without copying
    , getSliceUnsafe
    , getSlice
    -- , getSlicesFromLenN
    , splitAt -- XXX should be able to express using getSlice
    , breakOn

    -- ** Cloning arrays
    , clone
    , pinnedClone

    -- ** Appending arrays
    , spliceCopy
    , spliceWith
    , splice
    , spliceExp
    , spliceUnsafe
    -- , putSlice
    -- , appendSlice
    -- , appendSliceFrom

    -- ** Utilities
    , roundUpToPower2
    , memcpy
    , memcmp
    , c_memchr
    )
where

#include "assert.hs"
#include "inline.hs"
#include "ArrayMacros.h"
#include "MachDeps.h"

import Control.Monad (when, void)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Bits (shiftR, (.|.), (.&.))
import Data.Functor.Identity (Identity(..))
import Data.Proxy (Proxy(..))
import Data.Word (Word8)
import Foreign.C.Types (CSize(..), CInt(..))
import Foreign.Ptr (plusPtr, minusPtr, nullPtr)
import Streamly.Internal.Data.MutByteArray.Type
    ( MutByteArray(..)
    , MutableByteArray
    , PinnedState(..)
    , getMutableByteArray#
    , putSliceUnsafe
    )
import Streamly.Internal.Data.Unbox (Unbox(..))
import GHC.Base
    ( IO(..)
    , Int(..)
    , compareByteArrays#
    , copyMutableByteArray#
    )
import GHC.Base (noinline)
import GHC.Exts (unsafeCoerce#)
import GHC.Ptr (Ptr(..))

import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Producer.Type (Producer (..))
import Streamly.Internal.Data.Stream.Type (Stream)
import Streamly.Internal.Data.StreamK.Type (StreamK)
import Streamly.Internal.Data.SVar.Type (adaptState, defState)
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.System.IO (arrayPayloadSize, defaultChunkSize)

import qualified Streamly.Internal.Data.Fold.Type as FL
import qualified Streamly.Internal.Data.MutByteArray.Type as Unboxed
import qualified Streamly.Internal.Data.Producer as Producer
import qualified Streamly.Internal.Data.Stream.Type as D
import qualified Streamly.Internal.Data.Stream.Lift as D
import qualified Streamly.Internal.Data.StreamK.Type as K
import qualified Prelude

import Prelude hiding
    (Foldable(..), read, unlines, splitAt, reverse, truncate)

#include "DocTestDataMutArray.hs"

-------------------------------------------------------------------------------
-- Foreign helpers
-------------------------------------------------------------------------------

foreign import ccall unsafe "string.h memcpy" c_memcpy
    :: Ptr Word8 -> Ptr Word8 -> CSize -> IO (Ptr Word8)

foreign import ccall unsafe "string.h memchr" c_memchr
    :: Ptr Word8 -> Word8 -> CSize -> IO (Ptr Word8)

foreign import ccall unsafe "string.h memcmp" c_memcmp
    :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt

-- | Given an 'Unboxed' type (unused first arg) and a number of bytes, return
-- how many elements of that type will completely fit in those bytes.
--
{-# INLINE bytesToElemCount #-}
bytesToElemCount :: forall a. Unbox a => a -> Int -> Int
bytesToElemCount :: forall a. Unbox a => a -> Int -> Int
bytesToElemCount a
_ Int
n = Int
n forall a. Integral a => a -> a -> a
`div` SIZE_OF(a)

-- XXX we are converting Int to CSize
memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy Ptr Word8
dst Ptr Word8
src Int
len = forall (f :: * -> *) a. Functor f => f a -> f ()
void (Ptr Word8 -> Ptr Word8 -> CSize -> IO (Ptr Word8)
c_memcpy Ptr Word8
dst Ptr Word8
src (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len))

-- XXX we are converting Int to CSize
-- return True if the memory locations have identical contents
{-# INLINE memcmp #-}
memcmp :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
memcmp :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
memcmp Ptr Word8
p1 Ptr Word8
p2 Int
len = do
    CInt
r <- Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt
c_memcmp Ptr Word8
p1 Ptr Word8
p2 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ CInt
r forall a. Eq a => a -> a -> Bool
== CInt
0

-------------------------------------------------------------------------------
-- MutArray Data Type
-------------------------------------------------------------------------------

-- $arrayNotes
--
-- We can use an 'Unboxed' constraint in the MutArray type and the constraint
-- can be automatically provided to a function that pattern matches on the
-- MutArray type. However, it has huge performance cost, so we do not use it.
-- Investigate a GHC improvement possiblity.

-- | An unboxed mutable array. An array is created with a given length
-- and capacity. Length is the number of valid elements in the array.  Capacity
-- is the maximum number of elements that the array can be expanded to without
-- having to reallocate the memory.
--
-- The elements in the array can be mutated in-place without changing the
-- reference (constructor). However, the length of the array cannot be mutated
-- in-place.  A new array reference is generated when the length changes.  When
-- the length is increased (upto the maximum reserved capacity of the array),
-- the array is not reallocated and the new reference uses the same underlying
-- memory as the old one.
--
-- Several routines in this module allow the programmer to control the capacity
-- of the array. The programmer can control the trade-off between memory usage
-- and performance impact due to reallocations when growing or shrinking the
-- array.
--
data MutArray a =
#ifdef DEVBUILD
    Unbox a =>
#endif
    -- The array is a range into arrContents. arrContents may be a superset of
    -- the slice represented by the array. All offsets are in bytes.
    MutArray
    { forall a. MutArray a -> MutByteArray
arrContents :: {-# UNPACK #-} !MutByteArray
    , forall a. MutArray a -> Int
arrStart :: {-# UNPACK #-} !Int  -- ^ index into arrContents
    , forall a. MutArray a -> Int
arrEnd   :: {-# UNPACK #-} !Int    -- ^ index into arrContents
                                       -- Represents the first invalid index of
                                       -- the array.
    , forall a. MutArray a -> Int
arrBound :: {-# UNPACK #-} !Int    -- ^ first invalid index of arrContents.
    }

-------------------------------------------------------------------------------
-- Pinning & Unpinning
-------------------------------------------------------------------------------

-- | Return a copy of the array in pinned memory if unpinned, else return the
-- original array.
{-# INLINE pin #-}
pin :: MutArray a -> IO (MutArray a)
pin :: forall a. MutArray a -> IO (MutArray a)
pin arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    if MutByteArray -> Bool
Unboxed.isPinned MutByteArray
arrContents
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure MutArray a
arr
    else forall (m :: * -> *) a. MonadIO m => MutArray a -> m (MutArray a)
pinnedClone MutArray a
arr

-- | Return a copy of the array in unpinned memory if pinned, else return the
-- original array.
{-# INLINE unpin #-}
unpin :: MutArray a -> IO (MutArray a)
unpin :: forall a. MutArray a -> IO (MutArray a)
unpin arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    if MutByteArray -> Bool
Unboxed.isPinned MutByteArray
arrContents
    then forall (m :: * -> *) a. MonadIO m => MutArray a -> m (MutArray a)
clone MutArray a
arr
    else forall (f :: * -> *) a. Applicative f => a -> f a
pure MutArray a
arr

-- | Return 'True' if the array is allocated in pinned memory.
{-# INLINE isPinned #-}
isPinned :: MutArray a -> Bool
isPinned :: forall a. MutArray a -> Bool
isPinned MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = MutByteArray -> Bool
Unboxed.isPinned MutByteArray
arrContents

-------------------------------------------------------------------------------
-- Construction
-------------------------------------------------------------------------------

-- XXX Change the names to use "new" instead of "newArray". That way we can use
-- the same names for managed file system objects as well. For unmanaged ones
-- we can use open/create etc as usual.
--
-- A new array is similar to "touch" creating a zero length file. An mmapped
-- array would be similar to a sparse file with holes. TBD: support mmapped
-- files and arrays.

-- GHC always guarantees word-aligned memory, alignment is important only when
-- we need more than that.  See stg_pinnedNewAlignedByteArrayzh and
-- allocatePinned in GHC source.

-- | @newArrayWith allocator alignment count@ allocates a new array of zero
-- length and with a capacity to hold @count@ elements, using @allocator
-- size alignment@ as the memory allocator function.
--
-- Alignment must be greater than or equal to machine word size and a power of
-- 2.
--
-- Alignment is ignored if the allocator allocates unpinned memory.
--
-- /Pre-release/
{-# INLINE newArrayWith #-}
newArrayWith :: forall m a. (MonadIO m, Unbox a)
    => (Int -> Int -> m MutByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> m MutByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith Int -> Int -> m MutByteArray
alloc Int
alignSize Int
count = do
    let size :: Int
size = forall a. Ord a => a -> a -> a
max (Int
count forall a. Num a => a -> a -> a
* SIZE_OF(a)) 0
    MutByteArray
contents <- Int -> Int -> m MutByteArray
alloc Int
size Int
alignSize
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray
        { arrContents :: MutByteArray
arrContents = MutByteArray
contents
        , arrStart :: Int
arrStart = Int
0
        , arrEnd :: Int
arrEnd   = Int
0
        , arrBound :: Int
arrBound = Int
size
        }

nil ::
#ifdef DEVBUILD
    Unbox a =>
#endif
    MutArray a
nil :: forall a. MutArray a
nil = forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
Unboxed.nil Int
0 Int
0 Int
0

{-# INLINE newBytesAs #-}
newBytesAs :: MonadIO m =>
#ifdef DEVBUILD
    Unbox a =>
#endif
    PinnedState -> Int -> m (MutArray a)
newBytesAs :: forall (m :: * -> *) a.
MonadIO m =>
PinnedState -> Int -> m (MutArray a)
newBytesAs PinnedState
ps Int
bytes = do
    MutByteArray
contents <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ PinnedState -> Int -> IO MutByteArray
Unboxed.newBytesAs PinnedState
ps Int
bytes
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray
        { arrContents :: MutByteArray
arrContents = MutByteArray
contents
        , arrStart :: Int
arrStart = Int
0
        , arrEnd :: Int
arrEnd   = Int
0
        , arrBound :: Int
arrBound = Int
bytes
        }

-- | Allocates a pinned empty array that can hold 'count' items.  The memory of
-- the array is uninitialized and the allocation is aligned as per the
-- 'Unboxed' instance of the type.
--
-- /Pre-release/
{-# INLINE pinnedNewBytes #-}
pinnedNewBytes :: MonadIO m =>
#ifdef DEVBUILD
    Unbox a =>
#endif
    Int -> m (MutArray a)
pinnedNewBytes :: forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
pinnedNewBytes = forall (m :: * -> *) a.
MonadIO m =>
PinnedState -> Int -> m (MutArray a)
newBytesAs PinnedState
Pinned

-- | Like 'newArrayWith' but using an allocator is a pinned memory allocator and
-- the alignment is dictated by the 'Unboxed' instance of the type.
--
-- /Internal/
{-# INLINE pinnedNewAligned #-}
pinnedNewAligned :: (MonadIO m, Unbox a) => Int -> Int -> m (MutArray a)
pinnedNewAligned :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Int -> m (MutArray a)
pinnedNewAligned =
    forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> m MutByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith (\Int
s Int
a -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Int -> Int -> IO MutByteArray
Unboxed.pinnedNewAlignedBytes Int
s Int
a)

{-# INLINE newAs #-}
newAs :: (MonadIO m, Unbox a) => PinnedState -> Int -> m (MutArray a)
newAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
ps =
    forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> m MutByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith
        (\Int
s Int
_ -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ PinnedState -> Int -> IO MutByteArray
Unboxed.newBytesAs PinnedState
ps Int
s)
        (forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"new: alignment is not used in unpinned arrays.")

-- XXX can unaligned allocation be more efficient when alignment is not needed?
--
-- | Allocates an empty pinned array that can hold 'count' items.  The memory of
-- the array is uninitialized and the allocation is aligned as per the 'Unboxed'
-- instance of the type.
--
{-# INLINE pinnedNew #-}
pinnedNew :: forall m a. (MonadIO m, Unbox a) => Int -> m (MutArray a)
pinnedNew :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
pinnedNew = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
Pinned

-- | Allocates an empty unpinned array that can hold 'count' items.  The memory
-- of the array is uninitialized.
--
{-# INLINE new #-}
new :: (MonadIO m, Unbox a) => Int -> m (MutArray a)
new :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
new = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
Unpinned

-------------------------------------------------------------------------------
-- Random writes
-------------------------------------------------------------------------------

-- | Write the given element to the given index of the array. Does not check if
-- the index is out of bounds of the array.
--
-- /Pre-release/
{-# INLINE putIndexUnsafe #-}
putIndexUnsafe :: forall m a. (MonadIO m, Unbox a)
    => Int -> MutArray a -> a -> m ()
putIndexUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndexUnsafe Int
i MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} a
x = do
    let index :: Int
index = Int
INDEX_OF(arrStart, i, a)
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index, arrEnd, a)) (return ())
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
index MutByteArray
arrContents  a
x

invalidIndex :: String -> Int -> a
invalidIndex :: forall a. [Char] -> Int -> a
invalidIndex [Char]
label Int
i =
    forall a. (?callStack::CallStack) => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
label forall a. [a] -> [a] -> [a]
++ [Char]
": invalid array index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
i

-- | /O(1)/ Write the given element at the given index in the array.
-- Performs in-place mutation of the array.
--
-- >>> putIndex ix arr val = MutArray.modifyIndex ix arr (const (val, ()))
-- >>> f = MutArray.putIndices
-- >>> putIndex ix arr val = Stream.fold (f arr) (Stream.fromPure (ix, val))
--
{-# INLINE putIndex #-}
putIndex :: forall m a. (MonadIO m, Unbox a) => Int -> MutArray a -> a -> m ()
putIndex :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndex Int
i MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} a
x = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)
    then forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
index MutByteArray
arrContents  a
x
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"putIndex" Int
i

-- | Write an input stream of (index, value) pairs to an array. Throws an
-- error if any index is out of bounds.
--
-- /Pre-release/
{-# INLINE putIndices #-}
putIndices :: forall m a. (MonadIO m, Unbox a)
    => MutArray a -> Fold m (Int, a) ()
putIndices :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> Fold m (Int, a) ()
putIndices MutArray a
arr = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' forall {m :: * -> *}. MonadIO m => () -> (Int, a) -> m ()
step (forall (m :: * -> *) a. Monad m => a -> m a
return ())

    where

    step :: () -> (Int, a) -> m ()
step () (Int
i, a
x) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndex Int
i MutArray a
arr a
x)

-- | Modify a given index of an array using a modifier function.
--
-- Unsafe because it does not check the bounds of the array.
--
-- /Pre-release/
modifyIndexUnsafe :: forall m a b. (MonadIO m, Unbox a) =>
    Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndexUnsafe :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndexUnsafe Int
i MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} a -> (a, b)
f = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
        forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_NEXT(index,a) <= arrEnd) (return ())
        a
r <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
index MutByteArray
arrContents
        let (a
x, b
res) = a -> (a, b)
f a
r
        forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
index MutByteArray
arrContents  a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return b
res

-- | Modify a given index of an array using a modifier function.
--
-- /Pre-release/
modifyIndex :: forall m a b. (MonadIO m, Unbox a) =>
    Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndex :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndex Int
i MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} a -> (a, b)
f = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)
    then forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        a
r <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
index MutByteArray
arrContents
        let (a
x, b
res) = a -> (a, b)
f a
r
        forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
index MutByteArray
arrContents  a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return b
res
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"modifyIndex" Int
i

-- | Modify the array indices generated by the supplied stream.
--
-- /Pre-release/
{-# INLINE modifyIndices #-}
modifyIndices :: forall m a . (MonadIO m, Unbox a)
    => MutArray a -> (Int -> a -> a) -> Fold m Int ()
modifyIndices :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> (Int -> a -> a) -> Fold m Int ()
modifyIndices MutArray a
arr Int -> a -> a
f = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' forall {m :: * -> *}. MonadIO m => () -> Int -> m ()
step m ()
initial

    where

    initial :: m ()
initial = forall (m :: * -> *) a. Monad m => a -> m a
return ()

    step :: () -> Int -> m ()
step () Int
i =
        let f1 :: a -> (a, ())
f1 a
x = (Int -> a -> a
f Int
i a
x, ())
         in forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndex Int
i MutArray a
arr a -> (a, ())
f1

-- | Modify each element of an array using the supplied modifier function.
--
-- This is an in-place equivalent of an immutable map operation.
--
-- /Pre-release/
modify :: forall m a. (MonadIO m, Unbox a)
    => MutArray a -> (a -> a) -> m ()
modify :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> (a -> a) -> m ()
modify MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} a -> a
f = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$
    Int -> IO ()
go Int
arrStart

    where

    go :: Int -> IO ()
go Int
i =
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (INDEX_VALID(i,arrEnd,a)) $ do
            r <- peekAt i arrContents
            pokeAt i arrContents (f r)
            go (INDEX_NEXT(i,a))

-- XXX We could specify the number of bytes to swap instead of Proxy. Need
-- to ensure that the memory does not overlap.
{-# INLINE swapArrayByteIndices #-}
swapArrayByteIndices ::
       forall a. Unbox a
    => Proxy a
    -> MutByteArray
    -> Int
    -> Int
    -> IO ()
swapArrayByteIndices :: forall a. Unbox a => Proxy a -> MutByteArray -> Int -> Int -> IO ()
swapArrayByteIndices Proxy a
_ MutByteArray
arrContents Int
i1 Int
i2 = do
    a
r1 <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
i1 MutByteArray
arrContents
    a
r2 <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
i2 MutByteArray
arrContents
    forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
i1 MutByteArray
arrContents (a
r2 :: a)
    forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
i2 MutByteArray
arrContents (a
r1 :: a)

-- | Swap the elements at two indices without validating the indices.
--
-- /Unsafe/: This could result in memory corruption if indices are not valid.
--
-- /Pre-release/
{-# INLINE unsafeSwapIndices #-}
unsafeSwapIndices :: forall m a. (MonadIO m, Unbox a)
    => Int -> Int -> MutArray a -> m ()
unsafeSwapIndices :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Int -> MutArray a -> m ()
unsafeSwapIndices Int
i1 Int
i2 MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let t1 :: Int
t1 = Int
INDEX_OF(arrStart,i1,a)
            t2 :: Int
t2 = Int
INDEX_OF(arrStart,i2,a)
        forall a. Unbox a => Proxy a -> MutByteArray -> Int -> Int -> IO ()
swapArrayByteIndices (forall {k} (t :: k). Proxy t
Proxy :: Proxy a) MutByteArray
arrContents Int
t1 Int
t2

-- | Swap the elements at two indices.
--
-- /Pre-release/
swapIndices :: forall m a. (MonadIO m, Unbox a)
    => Int -> Int -> MutArray a -> m ()
swapIndices :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Int -> MutArray a -> m ()
swapIndices Int
i1 Int
i2 MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let t1 :: Int
t1 = Int
INDEX_OF(arrStart,i1,a)
            t2 :: Int
t2 = Int
INDEX_OF(arrStart,i2,a)
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i1 forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| INDEX_INVALID(t1,arrEnd,a))
            forall a b. (a -> b) -> a -> b
$ forall a. [Char] -> Int -> a
invalidIndex [Char]
"swapIndices" Int
i1
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i2 forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| INDEX_INVALID(t2,arrEnd,a))
            forall a b. (a -> b) -> a -> b
$ forall a. [Char] -> Int -> a
invalidIndex [Char]
"swapIndices" Int
i2
        forall a. Unbox a => Proxy a -> MutByteArray -> Int -> Int -> IO ()
swapArrayByteIndices (forall {k} (t :: k). Proxy t
Proxy :: Proxy a) MutByteArray
arrContents Int
t1 Int
t2

-------------------------------------------------------------------------------
-- Rounding
-------------------------------------------------------------------------------

-- XXX Should we use bitshifts in calculations or it gets optimized by the
-- compiler/processor itself?
--
-- | The page or block size used by the GHC allocator. Allocator allocates at
-- least a block and then allocates smaller allocations from within a block.
blockSize :: Int
blockSize :: Int
blockSize = Int
4 forall a. Num a => a -> a -> a
* Int
1024

-- | Allocations larger than 'largeObjectThreshold' are in multiples of block
-- size and are always pinned. The space beyond the end of a large object up to
-- the end of the block is unused.
largeObjectThreshold :: Int
largeObjectThreshold :: Int
largeObjectThreshold = (Int
blockSize forall a. Num a => a -> a -> a
* Int
8) forall a. Integral a => a -> a -> a
`div` Int
10

-- XXX Should be done only when we are using the GHC allocator.
-- | Round up an array larger than 'largeObjectThreshold' to use the whole
-- block.
{-# INLINE roundUpLargeArray #-}
roundUpLargeArray :: Int -> Int
roundUpLargeArray :: Int -> Int
roundUpLargeArray Int
size =
    if Int
size forall a. Ord a => a -> a -> Bool
>= Int
largeObjectThreshold
    then
        forall a. (?callStack::CallStack) => Bool -> a -> a
assert
            (Int
blockSize forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& ((Int
blockSize forall a. Bits a => a -> a -> a
.&. (Int
blockSize forall a. Num a => a -> a -> a
- Int
1)) forall a. Eq a => a -> a -> Bool
== Int
0))
            ((Int
size forall a. Num a => a -> a -> a
+ Int
blockSize forall a. Num a => a -> a -> a
- Int
1) forall a. Bits a => a -> a -> a
.&. forall a. Num a => a -> a
negate Int
blockSize)
    else Int
size

{-# INLINE isPower2 #-}
isPower2 :: Int -> Bool
isPower2 :: Int -> Bool
isPower2 Int
n = Int
n forall a. Bits a => a -> a -> a
.&. (Int
n forall a. Num a => a -> a -> a
- Int
1) forall a. Eq a => a -> a -> Bool
== Int
0

{-# INLINE roundUpToPower2 #-}
roundUpToPower2 :: Int -> Int
roundUpToPower2 :: Int -> Int
roundUpToPower2 Int
n =
#if WORD_SIZE_IN_BITS == 64
    Int
1 forall a. Num a => a -> a -> a
+ Int
z6
#else
    1 + z5
#endif

    where

    z0 :: Int
z0 = Int
n forall a. Num a => a -> a -> a
- Int
1
    z1 :: Int
z1 = Int
z0 forall a. Bits a => a -> a -> a
.|. Int
z0 forall a. Bits a => a -> Int -> a
`shiftR` Int
1
    z2 :: Int
z2 = Int
z1 forall a. Bits a => a -> a -> a
.|. Int
z1 forall a. Bits a => a -> Int -> a
`shiftR` Int
2
    z3 :: Int
z3 = Int
z2 forall a. Bits a => a -> a -> a
.|. Int
z2 forall a. Bits a => a -> Int -> a
`shiftR` Int
4
    z4 :: Int
z4 = Int
z3 forall a. Bits a => a -> a -> a
.|. Int
z3 forall a. Bits a => a -> Int -> a
`shiftR` Int
8
    z5 :: Int
z5 = Int
z4 forall a. Bits a => a -> a -> a
.|. Int
z4 forall a. Bits a => a -> Int -> a
`shiftR` Int
16
    z6 :: Int
z6 = Int
z5 forall a. Bits a => a -> a -> a
.|. Int
z5 forall a. Bits a => a -> Int -> a
`shiftR` Int
32

-- | @allocBytesToBytes elem allocatedBytes@ returns the array size in bytes
-- such that the real allocation is less than or equal to @allocatedBytes@,
-- unless @allocatedBytes@ is less than the size of one array element in which
-- case it returns one element's size.
--
{-# INLINE allocBytesToBytes #-}
allocBytesToBytes :: forall a. Unbox a => a -> Int -> Int
allocBytesToBytes :: forall a. Unbox a => a -> Int -> Int
allocBytesToBytes a
_ Int
n = forall a. Ord a => a -> a -> a
max (Int -> Int
arrayPayloadSize Int
n) (SIZE_OF(a))

-- | Given an 'Unboxed' type (unused first arg) and real allocation size
-- (including overhead), return how many elements of that type will completely
-- fit in it, returns at least 1.
--
{-# INLINE allocBytesToElemCount #-}
allocBytesToElemCount :: Unbox a => a -> Int -> Int
allocBytesToElemCount :: forall a. Unbox a => a -> Int -> Int
allocBytesToElemCount a
x Int
bytes =
    let n :: Int
n = forall a. Unbox a => a -> Int -> Int
bytesToElemCount a
x (forall a. Unbox a => a -> Int -> Int
allocBytesToBytes a
x Int
bytes)
     in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
>= Int
1) Int
n

-- | The default chunk size by which the array creation routines increase the
-- size of the array when the array is grown linearly.
arrayChunkBytes :: Int
arrayChunkBytes :: Int
arrayChunkBytes = Int
1024

-------------------------------------------------------------------------------
-- Resizing
-------------------------------------------------------------------------------

-- | Round the second argument down to multiples of the first argument.
{-# INLINE roundDownTo #-}
roundDownTo :: Int -> Int -> Int
roundDownTo :: Int -> Int -> Int
roundDownTo Int
elemSize Int
size = Int
size forall a. Num a => a -> a -> a
- (Int
size forall a. Integral a => a -> a -> a
`mod` Int
elemSize)

-- XXX See if resizing can be implemented by reading the old array as a stream
-- and then using writeN to the new array.
--
-- NOTE: we are passing elemSize explicitly to avoid an Unboxed constraint.
-- Since this is not inlined Unboxed consrraint leads to dictionary passing
-- which complicates some inspection tests.
--
{-# NOINLINE reallocExplicit #-}
reallocExplicit :: Int -> Int -> MutArray a -> IO (MutArray a)
reallocExplicit :: forall a. Int -> Int -> MutArray a -> IO (MutArray a)
reallocExplicit Int
elemSize Int
newCapacityInBytes MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = do
    assertM(Int
arrEnd forall a. Ord a => a -> a -> Bool
<= Int
arrBound)

    -- Allocate new array
    let newCapMaxInBytes :: Int
newCapMaxInBytes = Int -> Int
roundUpLargeArray Int
newCapacityInBytes
    MutByteArray
contents <-
        if MutByteArray -> Bool
Unboxed.isPinned MutByteArray
arrContents
        then Int -> IO MutByteArray
Unboxed.pinnedNew Int
newCapMaxInBytes
        else Int -> IO MutByteArray
Unboxed.new Int
newCapMaxInBytes
    let !(MutByteArray MutableByteArray# RealWorld
mbarrFrom#) = MutByteArray
arrContents
        !(MutByteArray MutableByteArray# RealWorld
mbarrTo#) = MutByteArray
contents

    -- Copy old data
    let oldStart :: Int
oldStart = Int
arrStart
        !(I# Int#
oldStartInBytes#) = Int
oldStart
        oldSizeInBytes :: Int
oldSizeInBytes = Int
arrEnd forall a. Num a => a -> a -> a
- Int
oldStart
        newCapInBytes :: Int
newCapInBytes = Int -> Int -> Int
roundDownTo Int
elemSize Int
newCapMaxInBytes
        !newLenInBytes :: Int
newLenInBytes@(I# Int#
newLenInBytes#) = forall a. Ord a => a -> a -> a
min Int
oldSizeInBytes Int
newCapInBytes
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
oldSizeInBytes forall a. Integral a => a -> a -> a
`mod` Int
elemSize forall a. Eq a => a -> a -> Bool
== Int
0) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
newLenInBytes forall a. Ord a => a -> a -> Bool
>= Int
0) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
newLenInBytes forall a. Integral a => a -> a -> a
`mod` Int
elemSize forall a. Eq a => a -> a -> Bool
== Int
0) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# -> (# forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray# MutableByteArray# RealWorld
mbarrFrom# Int#
oldStartInBytes#
                        MutableByteArray# RealWorld
mbarrTo# Int#
0# Int#
newLenInBytes# State# RealWorld
s#, () #)

    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray
        { arrStart :: Int
arrStart = Int
0
        , arrContents :: MutByteArray
arrContents = MutByteArray
contents
        , arrEnd :: Int
arrEnd   = Int
newLenInBytes
        , arrBound :: Int
arrBound = Int
newCapInBytes
        }

-- | @realloc newCapacity array@ reallocates the array to the specified
-- capacity in bytes.
--
-- If the new size is less than the original array the array gets truncated.
-- If the new size is not a multiple of array element size then it is rounded
-- down to multiples of array size.  If the new size is more than
-- 'largeObjectThreshold' then it is rounded up to the block size (4K).
--
-- If the original array is pinned, the newly allocated array is also pinned.
{-# INLINABLE realloc #-}
realloc :: forall m a. (MonadIO m, Unbox a) => Int -> MutArray a -> m (MutArray a)
realloc :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
bytes MutArray a
arr = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> MutArray a -> IO (MutArray a)
reallocExplicit (SIZE_OF(a)) bytes arr

-- | @reallocWith label capSizer minIncrBytes array@. The label is used
-- in error messages and the capSizer is used to determine the capacity of the
-- new array in bytes given the current byte length of the array.
reallocWith :: forall m a. (MonadIO m , Unbox a) =>
       String
    -> (Int -> Int)
    -> Int
    -> MutArray a
    -> m (MutArray a)
reallocWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
[Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
reallocWith [Char]
label Int -> Int
capSizer Int
minIncrBytes MutArray a
arr = do
    let oldSizeBytes :: Int
oldSizeBytes = forall a. MutArray a -> Int
arrEnd MutArray a
arr forall a. Num a => a -> a -> a
- forall a. MutArray a -> Int
arrStart MutArray a
arr
        newCapBytes :: Int
newCapBytes = Int -> Int
capSizer Int
oldSizeBytes
        newSizeBytes :: Int
newSizeBytes = Int
oldSizeBytes forall a. Num a => a -> a -> a
+ Int
minIncrBytes
        safeCapBytes :: Int
safeCapBytes = forall a. Ord a => a -> a -> a
max Int
newCapBytes Int
newSizeBytes
    assertM(Int
safeCapBytes forall a. Ord a => a -> a -> Bool
>= Int
newSizeBytes Bool -> Bool -> Bool
|| forall a. (?callStack::CallStack) => [Char] -> a
error (forall a. Show a => a -> [Char]
badSize Int
newSizeBytes))

    forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
safeCapBytes MutArray a
arr

    where

    badSize :: a -> [Char]
badSize a
newSize =
        forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
            [ [Char]
label
            , [Char]
": new array size (in bytes) is less than required size "
            , forall a. Show a => a -> [Char]
show a
newSize
            , [Char]
". Please check the sizing function passed."
            ]

-- | @resize newCapacity array@ changes the total capacity of the array so that
-- it is enough to hold the specified number of elements.  Nothing is done if
-- the specified capacity is less than the length of the array.
--
-- If the capacity is more than 'largeObjectThreshold' then it is rounded up to
-- the block size (4K).
--
-- /Pre-release/
{-# INLINE resize #-}
resize :: forall m a. (MonadIO m, Unbox a) =>
    Int -> MutArray a -> m (MutArray a)
resize :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
resize Int
nElems arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = do
    let req :: Int
req = SIZE_OF(a) * nElems
        len :: Int
len = Int
arrEnd forall a. Num a => a -> a -> a
- Int
arrStart
    if Int
req forall a. Ord a => a -> a -> Bool
< Int
len
    then forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr
    else forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
req MutArray a
arr

-- | Like 'resize' but if the byte capacity is more than 'largeObjectThreshold'
-- then it is rounded up to the closest power of 2.
--
-- /Pre-release/
{-# INLINE resizeExp #-}
resizeExp :: forall m a. (MonadIO m, Unbox a) =>
    Int -> MutArray a -> m (MutArray a)
resizeExp :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
resizeExp Int
nElems arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = do
    let req :: Int
req = Int -> Int
roundUpLargeArray (SIZE_OF(a) * nElems)
        req1 :: Int
req1 =
            if Int
req forall a. Ord a => a -> a -> Bool
> Int
largeObjectThreshold
            then Int -> Int
roundUpToPower2 Int
req
            else Int
req
        len :: Int
len = Int
arrEnd forall a. Num a => a -> a -> a
- Int
arrStart
    if Int
req1 forall a. Ord a => a -> a -> Bool
< Int
len
    then forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr
    else forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
req1 MutArray a
arr

-- | Resize the allocated memory to drop any reserved free space at the end of
-- the array and reallocate it to reduce wastage.
--
-- Up to 25% wastage is allowed to avoid reallocations.  If the capacity is
-- more than 'largeObjectThreshold' then free space up to the 'blockSize' is
-- retained.
--
-- /Pre-release/
{-# INLINE rightSize #-}
rightSize :: forall m a. (MonadIO m, Unbox a) => MutArray a -> m (MutArray a)
rightSize :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> m (MutArray a)
rightSize arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = do
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
arrEnd forall a. Ord a => a -> a -> Bool
<= Int
arrBound) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let start :: Int
start = Int
arrStart
        len :: Int
len = Int
arrEnd forall a. Num a => a -> a -> a
- Int
start
        capacity :: Int
capacity = Int
arrBound forall a. Num a => a -> a -> a
- Int
start
        target :: Int
target = Int -> Int
roundUpLargeArray Int
len
        waste :: Int
waste = Int
arrBound forall a. Num a => a -> a -> a
- Int
arrEnd
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
target forall a. Ord a => a -> a -> Bool
>= Int
len) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
len forall a. Integral a => a -> a -> a
`mod` SIZE_OF(a) == 0) (return ())
    -- We trade off some wastage (25%) to avoid reallocations and copying.
    if Int
target forall a. Ord a => a -> a -> Bool
< Int
capacity Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
< Int
3 forall a. Num a => a -> a -> a
* Int
waste
    then forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
target MutArray a
arr
    else forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr

-------------------------------------------------------------------------------
-- Snoc
-------------------------------------------------------------------------------

-- XXX We can possibly use a smallMutableByteArray to hold the start, end,
-- bound pointers.  Using fully mutable handle will ensure that we do not have
-- multiple references to the same array of different lengths lying around and
-- potentially misused. In that case "snoc" need not return a new array (snoc
-- :: MutArray a -> a -> m ()), it will just modify the old reference.  The array
-- length will be mutable.  This means the length function would also be
-- monadic.  Mutable arrays would behave more like files that grow in that
-- case.

-- | Snoc using a 'Ptr'. Low level reusable function.
--
-- /Internal/
{-# INLINE snocNewEnd #-}
snocNewEnd :: (MonadIO m, Unbox a) => Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd Int
newEnd arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} a
x = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
newEnd forall a. Ord a => a -> a -> Bool
<= Int
arrBound) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
arrEnd MutByteArray
arrContents a
x
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrEnd :: Int
arrEnd = Int
newEnd}

-- | Really really unsafe, appends the element into the first array, may
-- cause silent data corruption or if you are lucky a segfault if the first
-- array does not have enough space to append the element.
--
-- /Internal/
{-# INLINE snocUnsafe #-}
snocUnsafe :: forall m a. (MonadIO m, Unbox a) =>
    MutArray a -> a -> m (MutArray a)
snocUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
snocUnsafe arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd (INDEX_NEXT(arrEnd,a)) arr

-- | Like 'snoc' but does not reallocate when pre-allocated array capacity
-- becomes full.
--
-- /Internal/
{-# INLINE snocMay #-}
snocMay :: forall m a. (MonadIO m, Unbox a) =>
    MutArray a -> a -> m (Maybe (MutArray a))
snocMay :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (Maybe (MutArray a))
snocMay arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} a
x = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let newEnd :: Int
newEnd = INDEX_NEXT(arrEnd,a)
    if Int
newEnd forall a. Ord a => a -> a -> Bool
<= Int
arrBound
    then forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd Int
newEnd MutArray a
arr a
x
    else forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

-- NOINLINE to move it out of the way and not pollute the instruction cache.
{-# NOINLINE snocWithRealloc #-}
snocWithRealloc :: forall m a. (MonadIO m, Unbox a) =>
       (Int -> Int)
    -> MutArray a
    -> a
    -> m (MutArray a)
snocWithRealloc :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWithRealloc Int -> Int
sizer MutArray a
arr a
x = do
    MutArray a
arr1 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
[Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
reallocWith [Char]
"snocWith" Int -> Int
sizer (SIZE_OF(a)) arr
    forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
snocUnsafe MutArray a
arr1 a
x

-- | @snocWith sizer arr elem@ mutates @arr@ to append @elem@. The length of
-- the array increases by 1.
--
-- If there is no reserved space available in @arr@ it is reallocated to a size
-- in bytes determined by the @sizer oldSizeBytes@ function, where
-- @oldSizeBytes@ is the original size of the array in bytes.
--
-- If the new array size is more than 'largeObjectThreshold' we automatically
-- round it up to 'blockSize'.
--
-- Note that the returned array may be a mutated version of the original array.
--
-- /Pre-release/
{-# INLINE snocWith #-}
snocWith :: forall m a. (MonadIO m, Unbox a) =>
       (Int -> Int)
    -> MutArray a
    -> a
    -> m (MutArray a)
snocWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith Int -> Int
allocSize MutArray a
arr a
x = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let newEnd :: Int
newEnd = INDEX_NEXT(arrEnd arr,a)
    if Int
newEnd forall a. Ord a => a -> a -> Bool
<= forall a. MutArray a -> Int
arrBound MutArray a
arr
    then forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd Int
newEnd MutArray a
arr a
x
    else forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWithRealloc Int -> Int
allocSize MutArray a
arr a
x

-- | The array is mutated to append an additional element to it. If there
-- is no reserved space available in the array then it is reallocated to grow
-- it by 'arrayChunkBytes' rounded up to 'blockSize' when the size becomes more
-- than 'largeObjectThreshold'.
--
-- Note that the returned array may be a mutated version of the original array.
--
-- Performs O(n^2) copies to grow but is thrifty on memory.
--
-- /Pre-release/
{-# INLINE snocLinear #-}
snocLinear :: forall m a. (MonadIO m, Unbox a) => MutArray a -> a -> m (MutArray a)
snocLinear :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
snocLinear = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith (forall a. Num a => a -> a -> a
+ forall a. Unbox a => a -> Int -> Int
allocBytesToBytes (forall a. (?callStack::CallStack) => a
undefined :: a) Int
arrayChunkBytes)

-- | The array is mutated to append an additional element to it. If there is no
-- reserved space available in the array then it is reallocated to double the
-- original size.
--
-- This is useful to reduce allocations when appending unknown number of
-- elements.
--
-- Note that the returned array may be a mutated version of the original array.
--
-- >>> snoc = MutArray.snocWith (* 2)
--
-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
--
{-# INLINE snoc #-}
snoc :: forall m a. (MonadIO m, Unbox a) => MutArray a -> a -> m (MutArray a)
snoc :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
snoc = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith Int -> Int
f

    where

    f :: Int -> Int
f Int
oldSize =
        if Int -> Bool
isPower2 Int
oldSize
        then Int
oldSize forall a. Num a => a -> a -> a
* Int
2
        else Int -> Int
roundUpToPower2 Int
oldSize forall a. Num a => a -> a -> a
* Int
2

-------------------------------------------------------------------------------
-- Random reads
-------------------------------------------------------------------------------

-- XXX Can this be deduplicated with array/foreign

-- | Return the element at the specified index without checking the bounds.
--
-- Unsafe because it does not check the bounds of the array.
{-# INLINE_NORMAL getIndexUnsafe #-}
getIndexUnsafe :: forall m a. (MonadIO m, Unbox a) => Int -> MutArray a -> m a
getIndexUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)) (return ())
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
index MutByteArray
arrContents

-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
--
{-# INLINE getIndex #-}
getIndex :: forall m a. (MonadIO m, Unbox a) => Int -> MutArray a -> m (Maybe a)
getIndex :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (Maybe a)
getIndex Int
i MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)
    then forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
index MutByteArray
arrContents
    else forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

-- | /O(1)/ Lookup the element at the given index from the end of the array.
-- Index starts from 0.
--
-- Slightly faster than computing the forward index and using getIndex.
--
{-# INLINE getIndexRev #-}
getIndexRev :: forall m a. (MonadIO m, Unbox a) => Int -> MutArray a -> m a
getIndexRev :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m a
getIndexRev Int
i MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = do
    let index :: Int
index = RINDEX_OF(forall a. Unbox a => Proxy a -> Int
arrEnd,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
index forall a. Ord a => a -> a -> Bool
>= Int
arrStart
    then forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
index MutByteArray
arrContents
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"getIndexRev" Int
i

data GetIndicesState contents start end st =
    GetIndicesState contents start end st

-- | Given an unfold that generates array indices, read the elements on those
-- indices from the supplied MutArray. An error is thrown if an index is out of
-- bounds.
--
-- /Pre-release/
{-# INLINE getIndicesD #-}
getIndicesD :: (Monad m, Unbox a) =>
    (forall b. IO b -> m b) -> D.Stream m Int -> Unfold m (MutArray a) a
getIndicesD :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Stream m Int -> Unfold m (MutArray a) a
getIndicesD forall b. IO b -> m b
liftio (D.Stream State StreamK m Int -> s -> m (Step s Int)
stepi s
sti) = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {a}.
Unbox a =>
GetIndicesState MutByteArray Int Int s
-> m (Step (GetIndicesState MutByteArray Int Int s) a)
step forall {m :: * -> *} {a}.
Monad m =>
MutArray a -> m (GetIndicesState MutByteArray Int Int s)
inject

    where

    inject :: MutArray a -> m (GetIndicesState MutByteArray Int Int s)
inject (MutArray MutByteArray
contents Int
start Int
end Int
_) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState MutByteArray
contents Int
start Int
end s
sti

    {-# INLINE_LATE step #-}
    step :: GetIndicesState MutByteArray Int Int s
-> m (Step (GetIndicesState MutByteArray Int Int s) a)
step (GetIndicesState MutByteArray
contents Int
start Int
end s
st) = do
        Step s Int
r <- State StreamK m Int -> s -> m (Step s Int)
stepi forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s Int
r of
            D.Yield Int
i s
s -> do
                Maybe a
x <- forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (Maybe a)
getIndex Int
i (forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
contents Int
start Int
end forall a. (?callStack::CallStack) => a
undefined)
                case Maybe a
x of
                    Just a
v -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
v (forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState MutByteArray
contents Int
start Int
end s
s)
                    Maybe a
Nothing -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Invalid Index"
            D.Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState MutByteArray
contents Int
start Int
end s
s)
            Step s Int
D.Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop

{-# INLINE getIndices #-}
getIndices :: (MonadIO m, Unbox a) => Stream m Int -> Unfold m (MutArray a) a
getIndices :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m Int -> Unfold m (MutArray a) a
getIndices = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Stream m Int -> Unfold m (MutArray a) a
getIndicesD forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

-------------------------------------------------------------------------------
-- Subarrays
-------------------------------------------------------------------------------

-- XXX We can also get immutable slices.

-- | /O(1)/ Slice an array in constant time.
--
-- Unsafe: The bounds of the slice are not checked.
--
-- /Unsafe/
--
-- /Pre-release/
{-# INLINE getSliceUnsafe #-}
getSliceUnsafe :: forall a. Unbox a
    => Int -- ^ from index
    -> Int -- ^ length of the slice
    -> MutArray a
    -> MutArray a
getSliceUnsafe :: forall a. Unbox a => Int -> Int -> MutArray a -> MutArray a
getSliceUnsafe Int
index Int
len (MutArray MutByteArray
contents Int
start Int
e Int
_) =
    let fp1 :: Int
fp1 = INDEX_OF(start,index,a)
        end :: Int
end = Int
fp1 forall a. Num a => a -> a -> a
+ (Int
len forall a. Num a => a -> a -> a
* SIZE_OF(a))
     in forall a. (?callStack::CallStack) => Bool -> a -> a
assert
            (Int
index forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
end forall a. Ord a => a -> a -> Bool
<= Int
e)
            -- Note: In a slice we always use bound = end so that the slice
            -- user cannot overwrite elements beyond the end of the slice.
            (forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
contents Int
fp1 Int
end Int
end)

-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
-- extends out of the array bounds.
--
-- /Pre-release/
{-# INLINE getSlice #-}
getSlice :: forall a. Unbox a =>
       Int -- ^ from index
    -> Int -- ^ length of the slice
    -> MutArray a
    -> MutArray a
getSlice :: forall a. Unbox a => Int -> Int -> MutArray a -> MutArray a
getSlice Int
index Int
len (MutArray MutByteArray
contents Int
start Int
e Int
_) =
    let fp1 :: Int
fp1 = INDEX_OF(start,index,a)
        end :: Int
end = Int
fp1 forall a. Num a => a -> a -> a
+ (Int
len forall a. Num a => a -> a -> a
* SIZE_OF(a))
     in if Int
index forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
end forall a. Ord a => a -> a -> Bool
<= Int
e
        -- Note: In a slice we always use bound = end so that the slice user
        -- cannot overwrite elements beyond the end of the slice.
        then forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
contents Int
fp1 Int
end Int
end
        else forall a. (?callStack::CallStack) => [Char] -> a
error
                forall a b. (a -> b) -> a -> b
$ [Char]
"getSlice: invalid slice, index "
                forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
index forall a. [a] -> [a] -> [a]
++ [Char]
" length " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
len

-------------------------------------------------------------------------------
-- In-place mutation algorithms
-------------------------------------------------------------------------------

-- XXX consider the bulk update/accumulation/permutation APIs from vector.

-- | You may not need to reverse an array because you can consume it in reverse
-- using 'readerRev'. To reverse large arrays you can read in reverse and write
-- to another array. However, in-place reverse can be useful to take adavantage
-- of cache locality and when you do not want to allocate additional memory.
--
{-# INLINE reverse #-}
reverse :: forall m a. (MonadIO m, Unbox a) => MutArray a -> m ()
reverse :: forall (m :: * -> *) a. (MonadIO m, Unbox a) => MutArray a -> m ()
reverse MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let l :: Int
l = Int
arrStart
        h :: Int
h = INDEX_PREV(arrEnd,a)
     in Int -> Int -> IO ()
swap Int
l Int
h

    where

    swap :: Int -> Int -> IO ()
swap Int
l Int
h = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
l forall a. Ord a => a -> a -> Bool
< Int
h) forall a b. (a -> b) -> a -> b
$ do
            forall a. Unbox a => Proxy a -> MutByteArray -> Int -> Int -> IO ()
swapArrayByteIndices (forall {k} (t :: k). Proxy t
Proxy :: Proxy a) MutByteArray
arrContents Int
l Int
h
            Int -> Int -> IO ()
swap (INDEX_NEXT(l,a)) (INDEX_PREV(h,aInt
))

-- | Generate the next permutation of the sequence, returns False if this is
-- the last permutation.
--
-- /Unimplemented/
{-# INLINE permute #-}
permute :: MutArray a -> m Bool
permute :: forall a (m :: * -> *). MutArray a -> m Bool
permute = forall a. (?callStack::CallStack) => a
undefined

-- | Partition an array into two halves using a partitioning predicate. The
-- first half retains values where the predicate is 'False' and the second half
-- retains values where the predicate is 'True'.
--
-- /Pre-release/
{-# INLINE partitionBy #-}
partitionBy :: forall m a. (MonadIO m, Unbox a)
    => (a -> Bool) -> MutArray a -> m (MutArray a, MutArray a)
partitionBy :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(a -> Bool) -> MutArray a -> m (MutArray a, MutArray a)
partitionBy a -> Bool
f arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    if Int
arrStart forall a. Ord a => a -> a -> Bool
>= Int
arrEnd
    then forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a
arr, MutArray a
arr)
    else do
        Int
ptr <- Int -> Int -> IO Int
go Int
arrStart (INDEX_PREV(arrEnd,a))
        let pl :: MutArray a
pl = forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
arrContents Int
arrStart Int
ptr Int
ptr
            pr :: MutArray a
pr = forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
arrContents Int
ptr Int
arrEnd Int
arrEnd
        forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. MutArray a
pl, forall a. MutArray a
pr)

    where

    -- Invariant low < high on entry, and on return as well
    moveHigh :: Int -> Int -> IO (Maybe (Int, a))
moveHigh Int
low Int
high = do
        a
h <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
high MutByteArray
arrContents
        if a -> Bool
f a
h
        then
            -- Correctly classified, continue the loop
            let high1 :: Int
high1 = INDEX_PREV(high,a)
             in if Int
low forall a. Eq a => a -> a -> Bool
== Int
high1
                then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
                else Int -> Int -> IO (Maybe (Int, a))
moveHigh Int
low Int
high1
        else forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just (Int
high, a
h)) -- incorrectly classified

    -- Keep a low pointer starting at the start of the array (first partition)
    -- and a high pointer starting at the end of the array (second partition).
    -- Keep incrementing the low ptr and decrementing the high ptr until both
    -- are wrongly classified, at that point swap the two and continue until
    -- the two pointer cross each other.
    --
    -- Invariants when entering this loop:
    -- low <= high
    -- Both low and high are valid locations within the array
    go :: Int -> Int -> IO Int
go Int
low Int
high = do
        a
l <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
low MutByteArray
arrContents
        if a -> Bool
f a
l
        then
            -- low is wrongly classified
            if Int
low forall a. Eq a => a -> a -> Bool
== Int
high
            then forall (m :: * -> *) a. Monad m => a -> m a
return Int
low
            else do -- low < high
                Maybe (Int, a)
r <- Int -> Int -> IO (Maybe (Int, a))
moveHigh Int
low Int
high
                case Maybe (Int, a)
r of
                    Maybe (Int, a)
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return Int
low
                    Just (Int
high1, a
h) -> do -- low < high1
                        forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
low MutByteArray
arrContents a
h
                        forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
high1 MutByteArray
arrContents a
l
                        let low1 :: Int
low1 = INDEX_NEXT(low,a)
                            high2 :: Int
high2 = INDEX_PREV(high1,a)
                        if Int
low1 forall a. Ord a => a -> a -> Bool
<= Int
high2
                        then Int -> Int -> IO Int
go Int
low1 Int
high2
                        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
low1 -- low1 > high2

        else do
            -- low is correctly classified
            let low1 :: Int
low1 = INDEX_NEXT(low,a)
            if Int
low forall a. Eq a => a -> a -> Bool
== Int
high
            then forall (m :: * -> *) a. Monad m => a -> m a
return Int
low1
            else Int -> Int -> IO Int
go Int
low1 Int
high

-- | Shuffle corresponding elements from two arrays using a shuffle function.
-- If the shuffle function returns 'False' then do nothing otherwise swap the
-- elements. This can be used in a bottom up fold to shuffle or reorder the
-- elements.
--
-- /Unimplemented/
{-# INLINE shuffleBy #-}
shuffleBy :: (a -> a -> m Bool) -> MutArray a -> MutArray a -> m ()
shuffleBy :: forall a (m :: * -> *).
(a -> a -> m Bool) -> MutArray a -> MutArray a -> m ()
shuffleBy = forall a. (?callStack::CallStack) => a
undefined

-- XXX we can also make the folds partial by stopping at a certain level.
--
-- | @divideBy level partition array@  performs a top down hierarchical
-- recursive partitioning fold of items in the container using the given
-- function as the partition function.  Level indicates the level in the tree
-- where the fold would stop.
--
-- This performs a quick sort if the partition function is
-- 'partitionBy (< pivot)'.
--
-- /Unimplemented/
{-# INLINABLE divideBy #-}
divideBy ::
    Int -> (MutArray a -> m (MutArray a, MutArray a)) -> MutArray a -> m ()
divideBy :: forall a (m :: * -> *).
Int
-> (MutArray a -> m (MutArray a, MutArray a)) -> MutArray a -> m ()
divideBy = forall a. (?callStack::CallStack) => a
undefined

-- | @mergeBy level merge array@ performs a pairwise bottom up fold recursively
-- merging the pairs using the supplied merge function. Level indicates the
-- level in the tree where the fold would stop.
--
-- This performs a random shuffle if the merge function is random.  If we
-- stop at level 0 and repeatedly apply the function then we can do a bubble
-- sort.
--
-- /Unimplemented/
mergeBy :: Int -> (MutArray a -> MutArray a -> m ()) -> MutArray a -> m ()
mergeBy :: forall a (m :: * -> *).
Int -> (MutArray a -> MutArray a -> m ()) -> MutArray a -> m ()
mergeBy = forall a. (?callStack::CallStack) => a
undefined

-------------------------------------------------------------------------------
-- Size
-------------------------------------------------------------------------------

-- | /O(1)/ Get the byte length of the array.
--
{-# INLINE byteLength #-}
byteLength :: MutArray a -> Int
byteLength :: forall a. MutArray a -> Int
byteLength MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    let len :: Int
len = Int
arrEnd forall a. Num a => a -> a -> a
- Int
arrStart
    in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
len forall a. Ord a => a -> a -> Bool
>= Int
0) Int
len

-- Note: try to avoid the use of length in performance sensitive internal
-- routines as it involves a costly 'div' operation. Instead use the end ptr
-- in the array to check the bounds etc.
--
-- | /O(1)/ Get the length of the array i.e. the number of elements in the
-- array.
--
-- Note that 'byteLength' is less expensive than this operation, as 'length'
-- involves a costly division operation.
--
{-# INLINE length #-}
length :: forall a. Unbox a => MutArray a -> Int
length :: forall a. Unbox a => MutArray a -> Int
length MutArray a
arr =
    let elemSize :: Int
elemSize = SIZE_OF(a)
        blen :: Int
blen = forall a. MutArray a -> Int
byteLength MutArray a
arr
     in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
blen forall a. Integral a => a -> a -> a
`mod` Int
elemSize forall a. Eq a => a -> a -> Bool
== Int
0) (Int
blen forall a. Integral a => a -> a -> a
`div` Int
elemSize)

-- | Get the total capacity of an array. An array may have space reserved
-- beyond the current used length of the array.
--
-- /Pre-release/
{-# INLINE byteCapacity #-}
byteCapacity :: MutArray a -> Int
byteCapacity :: forall a. MutArray a -> Int
byteCapacity MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    let len :: Int
len = Int
arrBound forall a. Num a => a -> a -> a
- Int
arrStart
    in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
len forall a. Ord a => a -> a -> Bool
>= Int
0) Int
len

-- | The remaining capacity in the array for appending more elements without
-- reallocation.
--
-- /Pre-release/
{-# INLINE bytesFree #-}
bytesFree :: MutArray a -> Int
bytesFree :: forall a. MutArray a -> Int
bytesFree MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    let n :: Int
n = Int
arrBound forall a. Num a => a -> a -> a
- Int
arrEnd
    in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
>= Int
0) Int
n

-------------------------------------------------------------------------------
-- Streams of arrays - Creation
-------------------------------------------------------------------------------

data GroupState s contents start end bound
    = GroupStart s
    | GroupBuffer s contents start end bound
    | GroupYield
        contents start end bound (GroupState s contents start end bound)
    | GroupFinish

{-# INLINE_NORMAL chunksOfAs #-}
chunksOfAs :: forall m a. (MonadIO m, Unbox a)
    => PinnedState -> Int -> D.Stream m a -> D.Stream m (MutArray a)
chunksOfAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Stream m a -> Stream m (MutArray a)
chunksOfAs PinnedState
ps Int
n (D.Stream State StreamK m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {a} {a}.
State StreamK m a
-> GroupState s MutByteArray Int Int Int
-> m (Step (GroupState s MutByteArray Int Int Int) (MutArray a))
step' (forall s contents start end bound.
s -> GroupState s contents start end bound
GroupStart s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State StreamK m a
-> GroupState s MutByteArray Int Int Int
-> m (Step (GroupState s MutByteArray Int Int Int) (MutArray a))
step' State StreamK m a
_ (GroupStart s
st) = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n forall a. Ord a => a -> a -> Bool
<= Int
0) forall a b. (a -> b) -> a -> b
$
            -- XXX we can pass the module string from the higher level API
            forall a. (?callStack::CallStack) => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Streamly.Internal.Data.MutArray.Mut.Type.chunksOf: "
                    forall a. [a] -> [a] -> [a]
++ [Char]
"the size of arrays [" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
n
                    forall a. [a] -> [a] -> [a]
++ [Char]
"] must be a natural number"
        (MutArray MutByteArray
contents Int
start Int
end Int
bound :: MutArray a) <- forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
ps Int
n
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
st MutByteArray
contents Int
start Int
end Int
bound)

    step' State StreamK m a
gst (GroupBuffer s
st MutByteArray
contents Int
start Int
end Int
bound) = do
        Step s a
r <- State StreamK m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State StreamK m a
gst) s
st
        case Step s a
r of
            D.Yield a
x s
s -> do
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
end MutByteArray
contents  a
x
                let end1 :: Int
end1 = INDEX_NEXT(end,a)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
                    if Int
end1 forall a. Ord a => a -> a -> Bool
>= Int
bound
                    then forall s a. s -> Step s a
D.Skip
                            (forall s contents start end bound.
contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
-> GroupState s contents start end bound
GroupYield
                                MutByteArray
contents Int
start Int
end1 Int
bound (forall s contents start end bound.
s -> GroupState s contents start end bound
GroupStart s
s))
                    else forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
s MutByteArray
contents Int
start Int
end1 Int
bound)
            D.Skip s
s ->
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
s MutByteArray
contents Int
start Int
end Int
bound)
            Step s a
D.Stop ->
                forall (m :: * -> *) a. Monad m => a -> m a
return
                    forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
-> GroupState s contents start end bound
GroupYield MutByteArray
contents Int
start Int
end Int
bound forall s contents start end bound.
GroupState s contents start end bound
GroupFinish)

    step' State StreamK m a
_ (GroupYield MutByteArray
contents Int
start Int
end Int
bound GroupState s MutByteArray Int Int Int
next) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield (forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
contents Int
start Int
end Int
bound) GroupState s MutByteArray Int Int Int
next

    step' State StreamK m a
_ GroupState s MutByteArray Int Int Int
GroupFinish = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop

-- | @chunksOf n stream@ groups the elements in the input stream into arrays of
-- @n@ elements each.
--
-- Same as the following but may be more efficient:
--
-- >>> chunksOf n = Stream.foldMany (MutArray.writeN n)
--
-- /Pre-release/
{-# INLINE_NORMAL chunksOf #-}
chunksOf :: forall m a. (MonadIO m, Unbox a)
    => Int -> D.Stream m a -> D.Stream m (MutArray a)
-- XXX the idiomatic implementation leads to large regression in the D.reverse'
-- benchmark. It seems it has difficulty producing optimized code when
-- converting to StreamK. Investigate GHC optimizations.
-- chunksOf n = D.foldMany (writeN n)
chunksOf :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> Stream m (MutArray a)
chunksOf = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Stream m a -> Stream m (MutArray a)
chunksOfAs PinnedState
Unpinned

-- | Like 'chunksOf' but creates pinned arrays.
{-# INLINE_NORMAL pinnedChunksOf #-}
pinnedChunksOf :: forall m a. (MonadIO m, Unbox a)
    => Int -> D.Stream m a -> D.Stream m (MutArray a)
-- pinnedChunksOf n = D.foldMany (pinnedWriteN n)
pinnedChunksOf :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> Stream m (MutArray a)
pinnedChunksOf = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Stream m a -> Stream m (MutArray a)
chunksOfAs PinnedState
Pinned

-- XXX This should take a PinnedState
-- XXX buffer to a list instead?
-- | Buffer the stream into arrays in memory.
{-# INLINE arrayStreamKFromStreamDAs #-}
arrayStreamKFromStreamDAs :: forall m a. (MonadIO m, Unbox a) =>
    PinnedState -> D.Stream m a -> m (StreamK m (MutArray a))
arrayStreamKFromStreamDAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Stream m a -> m (StreamK m (MutArray a))
arrayStreamKFromStreamDAs PinnedState
ps =
    let n :: Int
n = forall a. Unbox a => a -> Int -> Int
allocBytesToElemCount (forall a. (?callStack::CallStack) => a
undefined :: a) Int
defaultChunkSize
     in forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr forall a (m :: * -> *). a -> StreamK m a -> StreamK m a
K.cons forall (m :: * -> *) a. StreamK m a
K.nil forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Stream m a -> Stream m (MutArray a)
chunksOfAs PinnedState
ps Int
n

-------------------------------------------------------------------------------
-- Streams of arrays - Flattening
-------------------------------------------------------------------------------

data FlattenState s contents a =
      OuterLoop s
    | InnerLoop s contents !Int !Int

-- | Use the "reader" unfold instead.
--
-- @flattenArrays = unfoldMany reader@
--
-- We can try this if there are any fusion issues in the unfold.
--
{-# INLINE_NORMAL flattenArrays #-}
flattenArrays :: forall m a. (MonadIO m, Unbox a)
    => D.Stream m (MutArray a) -> D.Stream m a
flattenArrays :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m (MutArray a) -> Stream m a
flattenArrays (D.Stream State StreamK m (MutArray a) -> s -> m (Step s (MutArray a))
step s
state) = forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {a} {m :: * -> *} {a} {a} {a}.
Unbox a =>
State StreamK m a
-> FlattenState s MutByteArray a
-> m (Step (FlattenState s MutByteArray a) a)
step' (forall s contents a. s -> FlattenState s contents a
OuterLoop s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State StreamK m a
-> FlattenState s MutByteArray a
-> m (Step (FlattenState s MutByteArray a) a)
step' State StreamK m a
gst (OuterLoop s
st) = do
        Step s (MutArray a)
r <- State StreamK m (MutArray a) -> s -> m (Step s (MutArray a))
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State StreamK m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s (MutArray a)
r of
            D.Yield MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} s
s ->
                forall s a. s -> Step s a
D.Skip (forall s contents a.
s -> contents -> Int -> Int -> FlattenState s contents a
InnerLoop s
s MutByteArray
arrContents Int
arrStart Int
arrEnd)
            D.Skip s
s -> forall s a. s -> Step s a
D.Skip (forall s contents a. s -> FlattenState s contents a
OuterLoop s
s)
            Step s (MutArray a)
D.Stop -> forall s a. Step s a
D.Stop

    step' State StreamK m a
_ (InnerLoop s
st MutByteArray
_ Int
p Int
end) | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
p forall a. Ord a => a -> a -> Bool
<= Int
end) (Int
p forall a. Eq a => a -> a -> Bool
== Int
end) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip forall a b. (a -> b) -> a -> b
$ forall s contents a. s -> FlattenState s contents a
OuterLoop s
st

    step' State StreamK m a
_ (InnerLoop s
st MutByteArray
contents Int
p Int
end) = do
        a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
contents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall s contents a.
s -> contents -> Int -> Int -> FlattenState s contents a
InnerLoop s
st MutByteArray
contents (INDEX_NEXT(p,a)) end)

-- | Use the "readerRev" unfold instead.
--
-- @flattenArrays = unfoldMany readerRev@
--
-- We can try this if there are any fusion issues in the unfold.
--
{-# INLINE_NORMAL flattenArraysRev #-}
flattenArraysRev :: forall m a. (MonadIO m, Unbox a)
    => D.Stream m (MutArray a) -> D.Stream m a
flattenArraysRev :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m (MutArray a) -> Stream m a
flattenArraysRev (D.Stream State StreamK m (MutArray a) -> s -> m (Step s (MutArray a))
step s
state) = forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {a} {m :: * -> *} {a} {a} {a}.
Unbox a =>
State StreamK m a
-> FlattenState s MutByteArray a
-> m (Step (FlattenState s MutByteArray a) a)
step' (forall s contents a. s -> FlattenState s contents a
OuterLoop s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State StreamK m a
-> FlattenState s MutByteArray a
-> m (Step (FlattenState s MutByteArray a) a)
step' State StreamK m a
gst (OuterLoop s
st) = do
        Step s (MutArray a)
r <- State StreamK m (MutArray a) -> s -> m (Step s (MutArray a))
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State StreamK m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s (MutArray a)
r of
            D.Yield MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} s
s ->
                let p :: Int
p = INDEX_PREV(arrEnd,a)
                 in forall s a. s -> Step s a
D.Skip (forall s contents a.
s -> contents -> Int -> Int -> FlattenState s contents a
InnerLoop s
s MutByteArray
arrContents Int
p Int
arrStart)
            D.Skip s
s -> forall s a. s -> Step s a
D.Skip (forall s contents a. s -> FlattenState s contents a
OuterLoop s
s)
            Step s (MutArray a)
D.Stop -> forall s a. Step s a
D.Stop

    step' State StreamK m a
_ (InnerLoop s
st MutByteArray
_ Int
p Int
start) | Int
p forall a. Ord a => a -> a -> Bool
< Int
start =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip forall a b. (a -> b) -> a -> b
$ forall s contents a. s -> FlattenState s contents a
OuterLoop s
st

    step' State StreamK m a
_ (InnerLoop s
st MutByteArray
contents Int
p Int
start) = do
        a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
contents
        let cur :: Int
cur = INDEX_PREV(p,a)
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall s contents a.
s -> contents -> Int -> Int -> FlattenState s contents a
InnerLoop s
st MutByteArray
contents Int
cur Int
start)

-------------------------------------------------------------------------------
-- Unfolds
-------------------------------------------------------------------------------

data ArrayUnsafe a = ArrayUnsafe
    {-# UNPACK #-} !MutByteArray   -- contents
    {-# UNPACK #-} !Int                -- index 1
    {-# UNPACK #-} !Int                -- index 2

toArrayUnsafe :: MutArray a -> ArrayUnsafe a
toArrayUnsafe :: forall a. MutArray a -> ArrayUnsafe a
toArrayUnsafe (MutArray MutByteArray
contents Int
start Int
end Int
_) = forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents Int
start Int
end

fromArrayUnsafe ::
#ifdef DEVBUILD
    Unbox a =>
#endif
    ArrayUnsafe a -> MutArray a
fromArrayUnsafe :: forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe (ArrayUnsafe MutByteArray
contents Int
start Int
end) =
         forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
contents Int
start Int
end Int
end

{-# INLINE_NORMAL producerWith #-}
producerWith ::
       forall m a. (Monad m, Unbox a)
    => (forall b. IO b -> m b) -> Producer m (MutArray a) a
producerWith :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Producer m (MutArray a) a
producerWith forall b. IO b -> m b
liftio = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer forall {a} {a} {a}.
Unbox a =>
ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. MutArray a -> ArrayUnsafe a
toArrayUnsafe) forall {a}. ArrayUnsafe a -> m (MutArray a)
extract
    where

    {-# INLINE_LATE step #-}
    step :: ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step (ArrayUnsafe MutByteArray
_ Int
cur Int
end)
        | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
cur forall a. Ord a => a -> a -> Bool
<= Int
end) (Int
cur forall a. Eq a => a -> a -> Bool
== Int
end) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (ArrayUnsafe MutByteArray
contents Int
cur Int
end) = do
            -- When we use a purely lazy Monad like Identity, we need to force a
            -- few actions for correctness and execution order sanity. We want
            -- the peek to occur right here and not lazily at some later point
            -- because we want the peek to be ordered with respect to the touch.
            !a
x <- forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
cur MutByteArray
contents
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents (INDEX_NEXT(cur,a)) end)

    extract :: ArrayUnsafe a -> m (MutArray a)
extract = forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe

-- | Resumable unfold of an array.
--
{-# INLINE_NORMAL producer #-}
producer :: forall m a. (MonadIO m, Unbox a) => Producer m (MutArray a) a
producer :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Producer m (MutArray a) a
producer = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Producer m (MutArray a) a
producerWith forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

-- | Unfold an array into a stream.
--
{-# INLINE_NORMAL reader #-}
reader :: forall m a. (MonadIO m, Unbox a) => Unfold m (MutArray a) a
reader :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Unfold m (MutArray a) a
reader = forall (m :: * -> *) a b. Producer m a b -> Unfold m a b
Producer.simplify forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Producer m (MutArray a) a
producer

{-# INLINE_NORMAL readerRevWith #-}
readerRevWith ::
       forall m a. (Monad m, Unbox a)
    => (forall b. IO b -> m b) -> Unfold m (MutArray a) a
readerRevWith :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Unfold m (MutArray a) a
readerRevWith forall b. IO b -> m b
liftio = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {a} {a} {a}.
Unbox a =>
ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step forall {m :: * -> *} {a} {a}.
Monad m =>
MutArray a -> m (ArrayUnsafe a)
inject
    where

    inject :: MutArray a -> m (ArrayUnsafe a)
inject (MutArray MutByteArray
contents Int
start Int
end Int
_) =
        let p :: Int
p = INDEX_PREV(end,a)
         in forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents Int
start Int
p

    {-# INLINE_LATE step #-}
    step :: ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step (ArrayUnsafe MutByteArray
_ Int
start Int
p) | Int
p forall a. Ord a => a -> a -> Bool
< Int
start = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (ArrayUnsafe MutByteArray
contents Int
start Int
p) = do
        !a
x <- forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
contents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents Int
start (INDEX_PREV(p,a)))

-- | Unfold an array into a stream in reverse order.
--
{-# INLINE_NORMAL readerRev #-}
readerRev :: forall m a. (MonadIO m, Unbox a) => Unfold m (MutArray a) a
readerRev :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Unfold m (MutArray a) a
readerRev = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Unfold m (MutArray a) a
readerRevWith forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

-------------------------------------------------------------------------------
-- to Lists and streams
-------------------------------------------------------------------------------

{-
-- Use foldr/build fusion to fuse with list consumers
-- This can be useful when using the IsList instance
{-# INLINE_LATE toListFB #-}
toListFB :: forall a b. Unbox a => (a -> b -> b) -> b -> MutArray a -> b
toListFB c n MutArray{..} = go arrStart
    where

    go p | assert (p <= arrEnd) (p == arrEnd) = n
    go p =
        -- unsafeInlineIO allows us to run this in Identity monad for pure
        -- toList/foldr case which makes them much faster due to not
        -- accumulating the list and fusing better with the pure consumers.
        --
        -- This should be safe as the array contents are guaranteed to be
        -- evaluated/written to before we peek at them.
        -- XXX
        let !x = unsafeInlineIO $ do
                    r <- peekAt arrContents p
                    return r
        in c x (go (PTR_NEXT(p,a)))
-}

-- XXX Monadic foldr/build fusion?
-- Reference: https://www.researchgate.net/publication/220676509_Monadic_augment_and_generalised_short_cut_fusion

-- | Convert a 'MutArray' into a list.
--
{-# INLINE toList #-}
toList :: forall m a. (MonadIO m, Unbox a) => MutArray a -> m [a]
toList :: forall (m :: * -> *) a. (MonadIO m, Unbox a) => MutArray a -> m [a]
toList MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall {a}. Unbox a => Int -> IO [a]
go Int
arrStart
    where

    go :: Int -> IO [a]
go Int
p | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
p forall a. Ord a => a -> a -> Bool
<= Int
arrEnd) (Int
p forall a. Eq a => a -> a -> Bool
== Int
arrEnd) = forall (m :: * -> *) a. Monad m => a -> m a
return []
    go Int
p = do
        a
x <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
arrContents
        (:) a
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO [a]
go (INDEX_NEXT(p,a))

{-# INLINE_NORMAL toStreamDWith #-}
toStreamDWith ::
       forall m a. (Monad m, Unbox a)
    => (forall b. IO b -> m b) -> MutArray a -> D.Stream m a
toStreamDWith :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> Stream m a
toStreamDWith forall b. IO b -> m b
liftio MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {a} {p}. Unbox a => p -> Int -> m (Step Int a)
step Int
arrStart

    where

    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
p | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
p forall a. Ord a => a -> a -> Bool
<= Int
arrEnd) (Int
p forall a. Eq a => a -> a -> Bool
== Int
arrEnd) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step p
_ Int
p = forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ do
        a
r <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
arrContents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
r (INDEX_NEXT(p,a))

-- | Convert a 'MutArray' into a stream.
--
-- >>> read = Stream.unfold MutArray.reader
--
{-# INLINE_NORMAL read #-}
read :: forall m a. (MonadIO m, Unbox a) => MutArray a -> D.Stream m a
read :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> Stream m a
read = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> Stream m a
toStreamDWith forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

{-# INLINE toStreamKWith #-}
toStreamKWith ::
       forall m a. (Monad m, Unbox a)
    => (forall b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKWith :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKWith forall b. IO b -> m b
liftio MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall {a}. Unbox a => Int -> StreamK m a
go Int
arrStart

    where

    go :: Int -> StreamK m a
go Int
p | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
p forall a. Ord a => a -> a -> Bool
<= Int
arrEnd) (Int
p forall a. Eq a => a -> a -> Bool
== Int
arrEnd) = forall (m :: * -> *) a. StreamK m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
arrContents
        in forall b. IO b -> m b
liftio IO a
elemM forall (m :: * -> *) a.
Monad m =>
m a -> StreamK m a -> StreamK m a
`K.consM` Int -> StreamK m a
go (INDEX_NEXT(p,a))

{-# INLINE toStreamK #-}
toStreamK :: forall m a. (MonadIO m, Unbox a) => MutArray a -> StreamK m a
toStreamK :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> StreamK m a
toStreamK = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKWith forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

{-# INLINE_NORMAL toStreamDRevWith #-}
toStreamDRevWith ::
       forall m a. (Monad m, Unbox a)
    => (forall b. IO b -> m b) -> MutArray a -> D.Stream m a
toStreamDRevWith :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> Stream m a
toStreamDRevWith forall b. IO b -> m b
liftio MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    let p :: Int
p = INDEX_PREV(arrEnd,a)
    in forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {a} {p}. Unbox a => p -> Int -> m (Step Int a)
step Int
p

    where

    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
p | Int
p forall a. Ord a => a -> a -> Bool
< Int
arrStart = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step p
_ Int
p = forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ do
        a
r <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
arrContents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
r (INDEX_PREV(p,a))

-- | Convert a 'MutArray' into a stream in reverse order.
--
-- >>> readRev = Stream.unfold MutArray.readerRev
--
{-# INLINE_NORMAL readRev #-}
readRev :: forall m a. (MonadIO m, Unbox a) => MutArray a -> D.Stream m a
readRev :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> Stream m a
readRev = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> Stream m a
toStreamDRevWith forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

{-# INLINE toStreamKRevWith #-}
toStreamKRevWith ::
       forall m a. (Monad m, Unbox a)
    => (forall b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKRevWith :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKRevWith forall b. IO b -> m b
liftio MutArray {Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    let p :: Int
p = INDEX_PREV(arrEnd,a)
    in forall {a}. Unbox a => Int -> StreamK m a
go Int
p

    where

    go :: Int -> StreamK m a
go Int
p | Int
p forall a. Ord a => a -> a -> Bool
< Int
arrStart = forall (m :: * -> *) a. StreamK m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
p MutByteArray
arrContents
        in forall b. IO b -> m b
liftio IO a
elemM forall (m :: * -> *) a.
Monad m =>
m a -> StreamK m a -> StreamK m a
`K.consM` Int -> StreamK m a
go (INDEX_PREV(p,a))

{-# INLINE toStreamKRev #-}
toStreamKRev :: forall m a. (MonadIO m, Unbox a) => MutArray a -> StreamK m a
toStreamKRev :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> StreamK m a
toStreamKRev = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKRevWith forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

-------------------------------------------------------------------------------
-- Folding
-------------------------------------------------------------------------------

-- XXX Need something like "MutArray m a" enforcing monadic action to avoid the
-- possibility of such APIs.
--
-- | Strict left fold of an array.
{-# INLINE_NORMAL foldl' #-}
foldl' :: (MonadIO m, Unbox a) => (b -> a -> b) -> b -> MutArray a -> m b
foldl' :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
(b -> a -> b) -> b -> MutArray a -> m b
foldl' b -> a -> b
f b
z MutArray a
arr = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
D.foldl' b -> a -> b
f b
z forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> Stream m a
read MutArray a
arr

-- | Right fold of an array.
{-# INLINE_NORMAL foldr #-}
foldr :: (MonadIO m, Unbox a) => (a -> b -> b) -> b -> MutArray a -> m b
foldr :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
(a -> b -> b) -> b -> MutArray a -> m b
foldr a -> b -> b
f b
z MutArray a
arr = forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr a -> b -> b
f b
z forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> Stream m a
read MutArray a
arr

-------------------------------------------------------------------------------
-- Folds
-------------------------------------------------------------------------------

-- Note: Arrays may be allocated with a specific alignment at the beginning of
-- the array. If you need to maintain that alignment on reallocations then you
-- can resize the array manually before append, using an aligned resize
-- operation.

-- XXX Keep the bound intact to not lose any free space? Perf impact?

-- | @writeAppendNUnsafe n arr@ appends up to @n@ input items to the supplied
-- array.
--
-- Unsafe: Do not drive the fold beyond @n@ elements, it will lead to memory
-- corruption or segfault.
--
-- Any free space left in the array after appending @n@ elements is lost.
--
-- /Internal/
{-# INLINE_NORMAL writeAppendNUnsafe #-}
writeAppendNUnsafe :: forall m a. (MonadIO m, Unbox a) =>
       Int
    -> m (MutArray a)
    -> Fold m a (MutArray a)
writeAppendNUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendNUnsafe Int
n m (MutArray a)
action =
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' forall {m :: * -> *} {a} {a} {a}.
(MonadIO m, Unbox a) =>
ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step m (ArrayUnsafe a)
initial

    where

    initial :: m (ArrayUnsafe a)
initial = do
        forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
>= Int
0) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
        arr :: MutArray a
arr@(MutArray MutByteArray
_ Int
_ Int
end Int
bound) <- m (MutArray a)
action
        let free :: Int
free = Int
bound forall a. Num a => a -> a -> a
- Int
end
            needed :: Int
needed = Int
n forall a. Num a => a -> a -> a
* SIZE_OF(a)
        -- XXX We can also reallocate if the array has too much free space,
        -- otherwise we lose that space.
        MutArray a
arr1 <-
            if Int
free forall a. Ord a => a -> a -> Bool
< Int
needed
            then forall a. a -> a
noinline forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
[Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
reallocWith [Char]
"writeAppendNUnsafeWith" (forall a. Num a => a -> a -> a
+ Int
needed) Int
needed MutArray a
arr
            else forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutArray a -> ArrayUnsafe a
toArrayUnsafe MutArray a
arr1

    step :: ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step (ArrayUnsafe MutByteArray
contents Int
start Int
end) a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
end MutByteArray
contents a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents Int
start (INDEX_NEXT(end,a))

-- | Append @n@ elements to an existing array. Any free space left in the array
-- after appending @n@ elements is lost.
--
-- >>> writeAppendN n initial = Fold.take n (MutArray.writeAppendNUnsafe n initial)
--
{-# INLINE_NORMAL writeAppendN #-}
writeAppendN :: forall m a. (MonadIO m, Unbox a) =>
    Int -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendN Int
n m (MutArray a)
initial = forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendNUnsafe Int
n m (MutArray a)
initial)

-- | @writeAppendWith realloc action@ mutates the array generated by @action@ to
-- append the input stream. If there is no reserved space available in the
-- array it is reallocated to a size in bytes  determined by @realloc oldSize@,
-- where @oldSize@ is the current size of the array in bytes.
--
-- Note that the returned array may be a mutated version of original array.
--
-- >>> writeAppendWith sizer = Fold.foldlM' (MutArray.snocWith sizer)
--
-- /Pre-release/
{-# INLINE writeAppendWith #-}
writeAppendWith :: forall m a. (MonadIO m, Unbox a) =>
    (Int -> Int) -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendWith Int -> Int
sizer = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith Int -> Int
sizer)

-- | @append action@ mutates the array generated by @action@ to append the
-- input stream. If there is no reserved space available in the array it is
-- reallocated to double the size.
--
-- Note that the returned array may be a mutated version of original array.
--
-- >>> writeAppend = MutArray.writeAppendWith (* 2)
--
{-# INLINE writeAppend #-}
writeAppend :: forall m a. (MonadIO m, Unbox a) =>
    m (MutArray a) -> Fold m a (MutArray a)
writeAppend :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
m (MutArray a) -> Fold m a (MutArray a)
writeAppend = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendWith (forall a. Num a => a -> a -> a
* Int
2)

-- XXX We can carry bound as well in the state to make sure we do not lose the
-- remaining capacity. Need to check perf impact.
--
-- | Like 'writeNUnsafe' but takes a new array allocator @alloc size@ function
-- as argument.
--
-- >>> writeNWithUnsafe alloc n = MutArray.writeAppendNUnsafe (alloc n) n
--
-- /Pre-release/
{-# INLINE_NORMAL writeNWithUnsafe #-}
writeNWithUnsafe :: forall m a. (MonadIO m, Unbox a)
    => (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWithUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWithUnsafe Int -> m (MutArray a)
alloc Int
n = forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' forall {m :: * -> *} {a} {a} {a}.
(MonadIO m, Unbox a) =>
ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step m (ArrayUnsafe a)
initial

    where

    initial :: m (ArrayUnsafe a)
initial = forall a. MutArray a -> ArrayUnsafe a
toArrayUnsafe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (MutArray a)
alloc (forall a. Ord a => a -> a -> a
max Int
n Int
0)

    step :: ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step (ArrayUnsafe MutByteArray
contents Int
start Int
end) a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
end MutByteArray
contents a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return
          forall a b. (a -> b) -> a -> b
$ forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents Int
start (INDEX_NEXT(end,a))

{-# INLINE_NORMAL writeNUnsafeAs #-}
writeNUnsafeAs :: forall m a. (MonadIO m, Unbox a)
    => PinnedState -> Int -> Fold m a (MutArray a)
writeNUnsafeAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeNUnsafeAs PinnedState
ps = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWithUnsafe (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
ps)

-- | Like 'writeN' but does not check the array bounds when writing. The fold
-- driver must not call the step function more than 'n' times otherwise it will
-- corrupt the memory and crash. This function exists mainly because any
-- conditional in the step function blocks fusion causing 10x performance
-- slowdown.
--
-- >>> writeNUnsafe = MutArray.writeNWithUnsafe MutArray.new
--
{-# INLINE_NORMAL writeNUnsafe #-}
writeNUnsafe :: forall m a. (MonadIO m, Unbox a)
    => Int -> Fold m a (MutArray a)
writeNUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeNUnsafe = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeNUnsafeAs PinnedState
Unpinned

-- | Like 'writeNUnsafe' but creates a pinned array.
{-# INLINE_NORMAL pinnedWriteNUnsafe #-}
pinnedWriteNUnsafe :: forall m a. (MonadIO m, Unbox a)
    => Int -> Fold m a (MutArray a)
pinnedWriteNUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
pinnedWriteNUnsafe = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeNUnsafeAs PinnedState
Pinned

-- | @writeNWith alloc n@ folds a maximum of @n@ elements into an array
-- allocated using the @alloc@ function.
--
-- >>> writeNWith alloc n = Fold.take n (MutArray.writeNWithUnsafe alloc n)
-- >>> writeNWith alloc n = MutArray.writeAppendN (alloc n) n
--
{-# INLINE_NORMAL writeNWith #-}
writeNWith :: forall m a. (MonadIO m, Unbox a)
    => (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWith Int -> m (MutArray a)
alloc Int
n = forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWithUnsafe Int -> m (MutArray a)
alloc Int
n)

{-# INLINE_NORMAL writeNAs #-}
writeNAs ::
       forall m a. (MonadIO m, Unbox a)
    => PinnedState
    -> Int
    -> Fold m a (MutArray a)
writeNAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeNAs PinnedState
ps = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWith (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
ps)

-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
-- 'MutArray'.
--
-- >>> writeN = MutArray.writeNWith MutArray.new
-- >>> writeN n = Fold.take n (MutArray.writeNUnsafe n)
-- >>> writeN n = MutArray.writeAppendN n (MutArray.new n)
--
{-# INLINE_NORMAL writeN #-}
writeN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (MutArray a)
writeN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeN = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeNAs PinnedState
Unpinned

-- | Like 'writeN' but creates a pinned array.
{-# INLINE_NORMAL pinnedWriteN #-}
pinnedWriteN ::
       forall m a. (MonadIO m, Unbox a)
    => Int
    -> Fold m a (MutArray a)
pinnedWriteN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
pinnedWriteN = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeNAs PinnedState
Pinned

-- | Like writeNWithUnsafe but writes the array in reverse order.
--
-- /Internal/
{-# INLINE_NORMAL writeRevNWithUnsafe #-}
writeRevNWithUnsafe :: forall m a. (MonadIO m, Unbox a)
    => (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWithUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWithUnsafe Int -> m (MutArray a)
alloc Int
n = forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' forall {m :: * -> *} {a} {a} {a}.
(MonadIO m, Unbox a) =>
ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step forall {a}. m (ArrayUnsafe a)
initial

    where

    toArrayUnsafeRev :: MutArray a -> ArrayUnsafe a
toArrayUnsafeRev (MutArray MutByteArray
contents Int
_ Int
_ Int
bound) =
         forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents Int
bound Int
bound

    initial :: m (ArrayUnsafe a)
initial = forall {a} {a}. MutArray a -> ArrayUnsafe a
toArrayUnsafeRev forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (MutArray a)
alloc (forall a. Ord a => a -> a -> a
max Int
n Int
0)

    step :: ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step (ArrayUnsafe MutByteArray
contents Int
start Int
end) a
x = do
        let ptr :: Int
ptr = INDEX_PREV(start,a)
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
ptr MutByteArray
contents a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return
          forall a b. (a -> b) -> a -> b
$ forall a. MutByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutByteArray
contents Int
ptr Int
end

-- | Like writeNWith but writes the array in reverse order.
--
-- /Internal/
{-# INLINE_NORMAL writeRevNWith #-}
writeRevNWith :: forall m a. (MonadIO m, Unbox a)
    => (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWith Int -> m (MutArray a)
alloc Int
n = forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWithUnsafe Int -> m (MutArray a)
alloc Int
n)

-- | Like writeN but writes the array in reverse order.
--
-- /Pre-release/
{-# INLINE_NORMAL writeRevN #-}
writeRevN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (MutArray a)
writeRevN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeRevN = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWith forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
new

-- | @pinnedWriteNAligned align n@ folds a maximum of @n@ elements from the
-- input stream to a 'MutArray' aligned to the given size.
--
-- >>> pinnedWriteNAligned align = MutArray.writeNWith (MutArray.pinnedNewAligned align)
-- >>> pinnedWriteNAligned align n = MutArray.writeAppendN n (MutArray.pinnedNewAligned align n)
--
-- /Pre-release/
--
{-# INLINE_NORMAL pinnedWriteNAligned #-}
pinnedWriteNAligned :: forall m a. (MonadIO m, Unbox a)
    => Int -> Int -> Fold m a (MutArray a)
pinnedWriteNAligned :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Int -> Fold m a (MutArray a)
pinnedWriteNAligned Int
align = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWith (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Int -> m (MutArray a)
pinnedNewAligned Int
align)

-- XXX Buffer to a list instead?

-- | Buffer a stream into a stream of arrays.
--
-- >>> writeChunks n = Fold.many (MutArray.writeN n) Fold.toStreamK
--
-- Breaking an array into an array stream  can be useful to consume a large
-- array sequentially such that memory of the array is released incrementatlly.
--
-- See also: 'arrayStreamKFromStreamD'.
--
-- /Unimplemented/
--
{-# INLINE_NORMAL writeChunks #-}
writeChunks :: (MonadIO m, Unbox a) =>
    Int -> Fold m a (StreamK n (MutArray a))
writeChunks :: forall (m :: * -> *) a (n :: * -> *).
(MonadIO m, Unbox a) =>
Int -> Fold m a (StreamK n (MutArray a))
writeChunks Int
n = forall (m :: * -> *) a b c.
Monad m =>
Fold m a b -> Fold m b c -> Fold m a c
FL.many (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeN Int
n) forall (m :: * -> *) a (n :: * -> *).
Monad m =>
Fold m a (StreamK n a)
FL.toStreamK

{-# INLINE_NORMAL writeWithAs #-}
writeWithAs :: forall m a. (MonadIO m, Unbox a)
    => PinnedState -> Int -> Fold m a (MutArray a)
-- writeWithAs ps n = FL.rmapM rightSize $ writeAppendWith (* 2) (newAs ps n)
writeWithAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeWithAs PinnedState
ps Int
elemCount =
    forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> Fold m a b -> Fold m a c
FL.rmapM MutArray a -> m (MutArray a)
extract forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
step m (MutArray a)
initial

    where

    initial :: m (MutArray a)
initial = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
elemCount forall a. Ord a => a -> a -> Bool
< Int
0) forall a b. (a -> b) -> a -> b
$ forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"writeWith: elemCount is negative"
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
ps Int
elemCount

    step :: MutArray a -> a -> m (MutArray a)
step arr :: MutArray a
arr@(MutArray MutByteArray
_ Int
start Int
end Int
bound) a
x
        | INDEX_NEXT(end,a) > bound = do
        let oldSize = end - start
            newSize = max (oldSize * 2) 1
        arr1 <- liftIO $ reallocExplicit (SIZE_OF(a)) newSize arr
        snocUnsafe arr1 x
    step MutArray a
arr a
x = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
snocUnsafe MutArray a
arr a
x

    extract :: MutArray a -> m (MutArray a)
extract = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> m (MutArray a)
rightSize

-- XXX Compare writeWith with fromStreamD which uses an array of streams
-- implementation. We can write this using writeChunks above if that is faster.
-- If writeWith is faster then we should use that to implement
-- fromStreamD.
--
-- XXX The realloc based implementation needs to make one extra copy if we use
-- shrinkToFit.  On the other hand, the stream of arrays implementation may
-- buffer the array chunk pointers in memory but it does not have to shrink as
-- we know the exact size in the end. However, memory copying does not seem to
-- be as expensive as the allocations. Therefore, we need to reduce the number
-- of allocations instead. Also, the size of allocations matters, right sizing
-- an allocation even at the cost of copying sems to help.  Should be measured
-- on a big stream with heavy calls to toArray to see the effect.
--
-- XXX check if GHC's memory allocator is efficient enough. We can try the C
-- malloc to compare against.

-- | @writeWith minCount@ folds the whole input to a single array. The array
-- starts at a size big enough to hold minCount elements, the size is doubled
-- every time the array needs to be grown.
--
-- /Caution! Do not use this on infinite streams./
--
-- >>> f n = MutArray.writeAppendWith (* 2) (MutArray.new n)
-- >>> writeWith n = Fold.rmapM MutArray.rightSize (f n)
-- >>> writeWith n = Fold.rmapM MutArray.fromArrayStreamK (MutArray.writeChunks n)
--
-- /Pre-release/
{-# INLINE_NORMAL writeWith #-}
writeWith :: forall m a. (MonadIO m, Unbox a)
    => Int -> Fold m a (MutArray a)
-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (new n)
writeWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeWith = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeWithAs PinnedState
Unpinned

-- | Fold the whole input to a single array.
--
-- Same as 'writeWith' using an initial array size of 'arrayChunkBytes' bytes
-- rounded up to the element size.
--
-- /Caution! Do not use this on infinite streams./
--
{-# INLINE write #-}
write :: forall m a. (MonadIO m, Unbox a) => Fold m a (MutArray a)
write :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Fold m a (MutArray a)
write = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeWith (forall a. Unbox a => a -> Int -> Int
allocBytesToElemCount (forall a. (?callStack::CallStack) => a
undefined :: a) Int
arrayChunkBytes)

-- | Like 'write' but creates a pinned array.
{-# INLINE pinnedWrite #-}
pinnedWrite :: forall m a. (MonadIO m, Unbox a) => Fold m a (MutArray a)
pinnedWrite :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Fold m a (MutArray a)
pinnedWrite =
    forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Fold m a (MutArray a)
writeWithAs PinnedState
Pinned (forall a. Unbox a => a -> Int -> Int
allocBytesToElemCount (forall a. (?callStack::CallStack) => a
undefined :: a) Int
arrayChunkBytes)

-------------------------------------------------------------------------------
-- construct from streams, known size
-------------------------------------------------------------------------------

{-# INLINE_NORMAL fromStreamDNAs #-}
fromStreamDNAs :: forall m a. (MonadIO m, Unbox a)
    => PinnedState -> Int -> D.Stream m a -> m (MutArray a)
fromStreamDNAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Stream m a -> m (MutArray a)
fromStreamDNAs PinnedState
ps Int
limit Stream m a
str = do
    (MutArray a
arr :: MutArray a) <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> m (MutArray a)
newAs PinnedState
ps Int
limit
    Int
end <- forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> m b
D.foldlM' (forall {m :: * -> *} {a}.
(MonadIO m, Unbox a) =>
MutByteArray -> Int -> a -> m Int
fwrite (forall a. MutArray a -> MutByteArray
arrContents MutArray a
arr)) (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutArray a -> Int
arrEnd MutArray a
arr) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Applicative m =>
Int -> Stream m a -> Stream m a
D.take Int
limit Stream m a
str
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrEnd :: Int
arrEnd = Int
end}

    where

    fwrite :: MutByteArray -> Int -> a -> m Int
fwrite MutByteArray
arrContents Int
ptr a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
pokeAt Int
ptr MutByteArray
arrContents  a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ INDEX_NEXT(ptr,a)

-- | Use the 'writeN' fold instead.
--
-- >>> fromStreamDN n = Stream.fold (MutArray.writeN n)
--
{-# INLINE_NORMAL fromStreamDN #-}
fromStreamDN :: forall m a. (MonadIO m, Unbox a)
    => Int -> D.Stream m a -> m (MutArray a)
-- fromStreamDN n = D.fold (writeN n)
fromStreamDN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> m (MutArray a)
fromStreamDN = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Stream m a -> m (MutArray a)
fromStreamDNAs PinnedState
Unpinned

-- | Create a 'MutArray' from the first N elements of a list. The array is
-- allocated to size N, if the list terminates before N elements then the
-- array may hold less than N elements.
--
{-# INLINABLE fromListN #-}
fromListN :: (MonadIO m, Unbox a) => Int -> [a] -> m (MutArray a)
fromListN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> [a] -> m (MutArray a)
fromListN Int
n [a]
xs = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> m (MutArray a)
fromStreamDN Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

-- | Like 'fromListN' but creates a pinned array.
{-# INLINABLE pinnedFromListN #-}
pinnedFromListN :: (MonadIO m, Unbox a) => Int -> [a] -> m (MutArray a)
pinnedFromListN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> [a] -> m (MutArray a)
pinnedFromListN Int
n [a]
xs = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Int -> Stream m a -> m (MutArray a)
fromStreamDNAs PinnedState
Pinned Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

-- | Like fromListN but writes the array in reverse order.
--
-- /Pre-release/
{-# INLINE fromListRevN #-}
fromListRevN :: (MonadIO m, Unbox a) => Int -> [a] -> m (MutArray a)
fromListRevN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> [a] -> m (MutArray a)
fromListRevN Int
n [a]
xs = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeRevN Int
n) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

-- | Convert a pure stream in Identity monad to a mutable array.
{-# INLINABLE fromPureStream #-}
fromPureStream :: (MonadIO m, Unbox a) => Stream Identity a -> m (MutArray a)
fromPureStream :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream Identity a -> m (MutArray a)
fromPureStream Stream Identity a
xs =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Fold m a (MutArray a)
write forall a b. (a -> b) -> a -> b
$ forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> Stream m a -> Stream n a
D.morphInner (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Identity a -> a
runIdentity) Stream Identity a
xs

-------------------------------------------------------------------------------
-- convert stream to a single array
-------------------------------------------------------------------------------

{-# INLINE arrayStreamKLength #-}
arrayStreamKLength :: (Monad m, Unbox a) => StreamK m (MutArray a) -> m Int
arrayStreamKLength :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
StreamK m (MutArray a) -> m Int
arrayStreamKLength StreamK m (MutArray a)
as = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> StreamK m a -> m b
K.foldl' forall a. Num a => a -> a -> a
(+) Int
0 (forall a b (m :: * -> *). (a -> b) -> StreamK m a -> StreamK m b
K.map forall a. Unbox a => MutArray a -> Int
length StreamK m (MutArray a)
as)

-- | Convert an array stream to an array. Note that this requires peak memory
-- that is double the size of the array stream.
--
{-# INLINE fromArrayStreamK #-}
fromArrayStreamK :: (Unbox a, MonadIO m) =>
    StreamK m (MutArray a) -> m (MutArray a)
fromArrayStreamK :: forall a (m :: * -> *).
(Unbox a, MonadIO m) =>
StreamK m (MutArray a) -> m (MutArray a)
fromArrayStreamK StreamK m (MutArray a)
as = do
    Int
len <- forall (m :: * -> *) a.
(Monad m, Unbox a) =>
StreamK m (MutArray a) -> m Int
arrayStreamKLength StreamK m (MutArray a)
as
    forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> m (MutArray a)
fromStreamDN Int
len forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> Stream m a -> Stream m b
D.unfoldMany forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Unfold m (MutArray a) a
reader forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => StreamK m a -> Stream m a
D.fromStreamK StreamK m (MutArray a)
as

{-# INLINE fromStreamDAs #-}
fromStreamDAs ::
       (MonadIO m, Unbox a) => PinnedState -> D.Stream m a -> m (MutArray a)
fromStreamDAs :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Stream m a -> m (MutArray a)
fromStreamDAs PinnedState
ps Stream m a
m = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Stream m a -> m (StreamK m (MutArray a))
arrayStreamKFromStreamDAs PinnedState
ps Stream m a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a (m :: * -> *).
(Unbox a, MonadIO m) =>
StreamK m (MutArray a) -> m (MutArray a)
fromArrayStreamK

-- CAUTION: a very large number (millions) of arrays can degrade performance
-- due to GC overhead because we need to buffer the arrays before we flatten
-- all the arrays.
--
-- XXX Compare if this is faster or "fold write".
--
-- | We could take the approach of doubling the memory allocation on each
-- overflow. This would result in more or less the same amount of copying as in
-- the chunking approach. However, if we have to shrink in the end then it may
-- result in an extra copy of the entire data.
--
-- >>> fromStreamD = StreamD.fold MutArray.write
--
{-# INLINE fromStreamD #-}
fromStreamD :: (MonadIO m, Unbox a) => D.Stream m a -> m (MutArray a)
fromStreamD :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m a -> m (MutArray a)
fromStreamD = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Stream m a -> m (MutArray a)
fromStreamDAs PinnedState
Unpinned

-- | Create a 'MutArray' from a list. The list must be of finite size.
--
{-# INLINE fromList #-}
fromList :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)
fromList :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
[a] -> m (MutArray a)
fromList [a]
xs = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m a -> m (MutArray a)
fromStreamD forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

-- | Like 'fromList' but creates a pinned array.
{-# INLINE pinnedFromList #-}
pinnedFromList :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)
pinnedFromList :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
[a] -> m (MutArray a)
pinnedFromList [a]
xs = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
PinnedState -> Stream m a -> m (MutArray a)
fromStreamDAs PinnedState
Pinned forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

-- XXX We are materializing the whole list first for getting the length. Check
-- if the 'fromList' like chunked implementation would fare better.

-- | Like 'fromList' but writes the contents of the list in reverse order.
{-# INLINE fromListRev #-}
fromListRev :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)
fromListRev :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
[a] -> m (MutArray a)
fromListRev [a]
xs = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> [a] -> m (MutArray a)
fromListRevN (forall (t :: * -> *) a. Foldable t => t a -> Int
Prelude.length [a]
xs) [a]
xs

-------------------------------------------------------------------------------
-- Cloning
-------------------------------------------------------------------------------

{-# INLINE cloneAs #-}
cloneAs ::
    ( MonadIO m
#ifdef DEVBUILD
    , Unbox a
#endif
    )
    => PinnedState -> MutArray a -> m (MutArray a)
cloneAs :: forall (m :: * -> *) a.
MonadIO m =>
PinnedState -> MutArray a -> m (MutArray a)
cloneAs PinnedState
ps MutArray a
src =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let startSrc :: Int
startSrc = forall a. MutArray a -> Int
arrStart MutArray a
src
            srcLen :: Int
srcLen = forall a. MutArray a -> Int
arrEnd MutArray a
src forall a. Num a => a -> a -> a
- Int
startSrc
        MutByteArray
newArrContents <-
            forall (m :: * -> *).
MonadIO m =>
PinnedState -> Int -> Int -> MutByteArray -> m MutByteArray
Unboxed.cloneSliceUnsafeAs PinnedState
ps Int
startSrc Int
srcLen (forall a. MutArray a -> MutByteArray
arrContents MutArray a
src)
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
newArrContents Int
0 Int
srcLen Int
srcLen

{-# INLINE clone #-}
clone ::
    ( MonadIO m
#ifdef DEVBUILD
    , Unbox a
#endif
    )
    => MutArray a -> m (MutArray a)
clone :: forall (m :: * -> *) a. MonadIO m => MutArray a -> m (MutArray a)
clone = forall (m :: * -> *) a.
MonadIO m =>
PinnedState -> MutArray a -> m (MutArray a)
cloneAs PinnedState
Unpinned

{-# INLINE pinnedClone #-}
pinnedClone ::
    ( MonadIO m
#ifdef DEVBUILD
    , Unbox a
#endif
    )
    => MutArray a -> m (MutArray a)
pinnedClone :: forall (m :: * -> *) a. MonadIO m => MutArray a -> m (MutArray a)
pinnedClone = forall (m :: * -> *) a.
MonadIO m =>
PinnedState -> MutArray a -> m (MutArray a)
cloneAs PinnedState
Pinned

-------------------------------------------------------------------------------
-- Combining
-------------------------------------------------------------------------------

-- | Copy two arrays into a newly allocated array. If the first array is pinned
-- the spliced array is also pinned.
{-# INLINE spliceCopy #-}
spliceCopy :: forall m a. MonadIO m =>
#ifdef DEVBUILD
    Unbox a =>
#endif
    MutArray a -> MutArray a -> m (MutArray a)
spliceCopy :: forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> MutArray a -> m (MutArray a)
spliceCopy MutArray a
arr1 MutArray a
arr2 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let start1 :: Int
start1 = forall a. MutArray a -> Int
arrStart MutArray a
arr1
        start2 :: Int
start2 = forall a. MutArray a -> Int
arrStart MutArray a
arr2
        len1 :: Int
len1 = forall a. MutArray a -> Int
arrEnd MutArray a
arr1 forall a. Num a => a -> a -> a
- Int
start1
        len2 :: Int
len2 = forall a. MutArray a -> Int
arrEnd MutArray a
arr2 forall a. Num a => a -> a -> a
- Int
start2
    let newLen :: Int
newLen = Int
len1 forall a. Num a => a -> a -> a
+ Int
len2
    MutByteArray
newArrContents <-
        if MutByteArray -> Bool
Unboxed.isPinned (forall a. MutArray a -> MutByteArray
arrContents MutArray a
arr1)
        then Int -> IO MutByteArray
Unboxed.pinnedNew Int
newLen
        else Int -> IO MutByteArray
Unboxed.new Int
newLen
    let len :: Int
len = Int
len1 forall a. Num a => a -> a -> a
+ Int
len2
    forall (m :: * -> *).
MonadIO m =>
MutByteArray -> Int -> MutByteArray -> Int -> Int -> m ()
putSliceUnsafe (forall a. MutArray a -> MutByteArray
arrContents MutArray a
arr1) Int
start1 MutByteArray
newArrContents Int
0 Int
len1
    forall (m :: * -> *).
MonadIO m =>
MutByteArray -> Int -> MutByteArray -> Int -> Int -> m ()
putSliceUnsafe (forall a. MutArray a -> MutByteArray
arrContents MutArray a
arr2) Int
start2 MutByteArray
newArrContents Int
len1 Int
len2
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
newArrContents Int
0 Int
len Int
len

-- | Really really unsafe, appends the second array into the first array. If
-- the first array does not have enough space it may cause silent data
-- corruption or if you are lucky a segfault.
{-# INLINE spliceUnsafe #-}
spliceUnsafe :: MonadIO m =>
    MutArray a -> MutArray a -> m (MutArray a)
spliceUnsafe :: forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> MutArray a -> m (MutArray a)
spliceUnsafe MutArray a
dst MutArray a
src =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
         let startSrc :: Int
startSrc = forall a. MutArray a -> Int
arrStart MutArray a
src
             srcLen :: Int
srcLen = forall a. MutArray a -> Int
arrEnd MutArray a
src forall a. Num a => a -> a -> a
- Int
startSrc
             endDst :: Int
endDst = forall a. MutArray a -> Int
arrEnd MutArray a
dst
         assertM(Int
endDst forall a. Num a => a -> a -> a
+ Int
srcLen forall a. Ord a => a -> a -> Bool
<= forall a. MutArray a -> Int
arrBound MutArray a
dst)
         forall (m :: * -> *).
MonadIO m =>
MutByteArray -> Int -> MutByteArray -> Int -> Int -> m ()
putSliceUnsafe
             (forall a. MutArray a -> MutByteArray
arrContents MutArray a
src) Int
startSrc (forall a. MutArray a -> MutByteArray
arrContents MutArray a
dst) Int
endDst Int
srcLen
         forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray a
dst {arrEnd :: Int
arrEnd = Int
endDst forall a. Num a => a -> a -> a
+ Int
srcLen}

-- | @spliceWith sizer dst src@ mutates @dst@ to append @src@. If there is no
-- reserved space available in @dst@ it is reallocated to a size determined by
-- the @sizer dstBytes srcBytes@ function, where @dstBytes@ is the size of the
-- first array and @srcBytes@ is the size of the second array, in bytes.
--
-- Note that the returned array may be a mutated version of first array.
--
-- /Pre-release/
{-# INLINE spliceWith #-}
spliceWith :: forall m a. (MonadIO m, Unbox a) =>
    (Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
spliceWith :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
spliceWith Int -> Int -> Int
sizer dst :: MutArray a
dst@(MutArray MutByteArray
_ Int
start Int
end Int
bound) MutArray a
src = do
{-
    let f = writeAppendWith (`sizer` byteLength src) (return dst)
     in D.fold f (toStreamD src)
-}
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
end forall a. Ord a => a -> a -> Bool
<= Int
bound) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let srcBytes :: Int
srcBytes = forall a. MutArray a -> Int
arrEnd MutArray a
src forall a. Num a => a -> a -> a
- forall a. MutArray a -> Int
arrStart MutArray a
src

    MutArray a
dst1 <-
        if Int
end forall a. Num a => a -> a -> a
+ Int
srcBytes forall a. Ord a => a -> a -> Bool
>= Int
bound
        then do
            let dstBytes :: Int
dstBytes = Int
end forall a. Num a => a -> a -> a
- Int
start
                newSizeInBytes :: Int
newSizeInBytes = Int -> Int -> Int
sizer Int
dstBytes Int
srcBytes
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
newSizeInBytes forall a. Ord a => a -> a -> Bool
< Int
dstBytes forall a. Num a => a -> a -> a
+ Int
srcBytes)
                forall a b. (a -> b) -> a -> b
$ forall a. (?callStack::CallStack) => [Char] -> a
error
                    forall a b. (a -> b) -> a -> b
$ [Char]
"splice: newSize is less than the total size "
                    forall a. [a] -> [a] -> [a]
++ [Char]
"of arrays being appended. Please check the "
                    forall a. [a] -> [a] -> [a]
++ [Char]
"sizer function passed."
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
newSizeInBytes MutArray a
dst
        else forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
dst
    forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> MutArray a -> m (MutArray a)
spliceUnsafe MutArray a
dst1 MutArray a
src

-- | The first array is mutated to append the second array. If there is no
-- reserved space available in the first array a new allocation of exact
-- required size is done.
--
-- Note that the returned array may be a mutated version of first array.
--
-- >>> splice = MutArray.spliceWith (+)
--
-- If the original array is pinned the spliced array is also pinned.
--
-- /Pre-release/
{-# INLINE splice #-}
splice :: (MonadIO m, Unbox a) => MutArray a -> MutArray a -> m (MutArray a)
splice :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> MutArray a -> m (MutArray a)
splice = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
spliceWith forall a. Num a => a -> a -> a
(+)

-- | Like 'append' but the growth of the array is exponential. Whenever a new
-- allocation is required the previous array size is at least doubled.
--
-- This is useful to reduce allocations when folding many arrays together.
--
-- Note that the returned array may be a mutated version of first array.
--
-- >>> spliceExp = MutArray.spliceWith (\l1 l2 -> max (l1 * 2) (l1 + l2))
--
-- /Pre-release/
{-# INLINE spliceExp #-}
spliceExp :: (MonadIO m, Unbox a) => MutArray a -> MutArray a -> m (MutArray a)
spliceExp :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> MutArray a -> m (MutArray a)
spliceExp = forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
spliceWith (\Int
l1 Int
l2 -> forall a. Ord a => a -> a -> a
max (Int
l1 forall a. Num a => a -> a -> a
* Int
2) (Int
l1 forall a. Num a => a -> a -> a
+ Int
l2))

-------------------------------------------------------------------------------
-- Splitting
-------------------------------------------------------------------------------

-- | Drops the separator byte
{-# INLINE breakOn #-}
breakOn :: MonadIO m
    => Word8 -> MutArray Word8 -> m (MutArray Word8, Maybe (MutArray Word8))
breakOn :: forall (m :: * -> *).
MonadIO m =>
Word8
-> MutArray Word8 -> m (MutArray Word8, Maybe (MutArray Word8))
breakOn Word8
sep arr :: MutArray Word8
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a b.
MonadIO m =>
MutArray a -> (Ptr a -> m b) -> m b
asPtrUnsafe MutArray Word8
arr forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    -- XXX Instead of using asPtrUnsafe (pinning memory) we can pass unlifted
    -- Addr# to memchr and it should be safe (from ghc 8.4).
    -- XXX We do not need memchr here, we can use a Haskell equivalent.
    Ptr Word8
loc <- Ptr Word8 -> Word8 -> CSize -> IO (Ptr Word8)
c_memchr Ptr Word8
p Word8
sep (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. MutArray a -> Int
byteLength MutArray Word8
arr)
    let sepIndex :: Int
sepIndex = Ptr Word8
loc forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
        if Ptr Word8
loc forall a. Eq a => a -> a -> Bool
== forall a. Ptr a
nullPtr
        then (MutArray Word8
arr, forall a. Maybe a
Nothing)
        else
            ( MutArray
                { arrContents :: MutByteArray
arrContents = MutByteArray
arrContents
                , arrStart :: Int
arrStart = Int
arrStart
                , arrEnd :: Int
arrEnd = Int
arrStart forall a. Num a => a -> a -> a
+ Int
sepIndex -- exclude the separator
                , arrBound :: Int
arrBound = Int
arrStart forall a. Num a => a -> a -> a
+ Int
sepIndex
                }
            , forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ MutArray
                    { arrContents :: MutByteArray
arrContents = MutByteArray
arrContents
                    , arrStart :: Int
arrStart = Int
arrStart forall a. Num a => a -> a -> a
+ (Int
sepIndex forall a. Num a => a -> a -> a
+ Int
1)
                    , arrEnd :: Int
arrEnd = Int
arrEnd
                    , arrBound :: Int
arrBound = Int
arrBound
                    }
            )

-- | Create two slices of an array without copying the original array. The
-- specified index @i@ is the first index of the second slice.
--
splitAt :: forall a. Unbox a => Int -> MutArray a -> (MutArray a, MutArray a)
splitAt :: forall a. Unbox a => Int -> MutArray a -> (MutArray a, MutArray a)
splitAt Int
i arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} =
    let maxIndex :: Int
maxIndex = forall a. Unbox a => MutArray a -> Int
length MutArray a
arr forall a. Num a => a -> a -> a
- Int
1
    in  if Int
i forall a. Ord a => a -> a -> Bool
< Int
0
        then forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"sliceAt: negative array index"
        else if Int
i forall a. Ord a => a -> a -> Bool
> Int
maxIndex
             then forall a. (?callStack::CallStack) => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"sliceAt: specified array index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
i
                        forall a. [a] -> [a] -> [a]
++ [Char]
" is beyond the maximum index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
maxIndex
             else let off :: Int
off = Int
i forall a. Num a => a -> a -> a
* SIZE_OF(a)
                      p :: Int
p = Int
arrStart forall a. Num a => a -> a -> a
+ Int
off
                in ( MutArray
                  { arrContents :: MutByteArray
arrContents = MutByteArray
arrContents
                  , arrStart :: Int
arrStart = Int
arrStart
                  , arrEnd :: Int
arrEnd = Int
p
                  , arrBound :: Int
arrBound = Int
p
                  }
                , MutArray
                  { arrContents :: MutByteArray
arrContents = MutByteArray
arrContents
                  , arrStart :: Int
arrStart = Int
p
                  , arrEnd :: Int
arrEnd = Int
arrEnd
                  , arrBound :: Int
arrBound = Int
arrBound
                  }
                )

-------------------------------------------------------------------------------
-- Casting
-------------------------------------------------------------------------------

-- | Cast an array having elements of type @a@ into an array having elements of
-- type @b@. The array size must be a multiple of the size of type @b@
-- otherwise accessing the last element of the array may result into a crash or
-- a random value.
--
-- /Pre-release/
--
castUnsafe ::
#ifdef DEVBUILD
    Unbox b =>
#endif
    MutArray a -> MutArray b
castUnsafe :: forall a b. MutArray a -> MutArray b
castUnsafe (MutArray MutByteArray
contents Int
start Int
end Int
bound) =
    forall a. MutByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutByteArray
contents Int
start Int
end Int
bound

-- | Cast an @MutArray a@ into an @MutArray Word8@.
--
asBytes :: MutArray a -> MutArray Word8
asBytes :: forall a. MutArray a -> MutArray Word8
asBytes = forall a b. MutArray a -> MutArray b
castUnsafe

-- | Cast an array having elements of type @a@ into an array having elements of
-- type @b@. The length of the array should be a multiple of the size of the
-- target element otherwise 'Nothing' is returned.
--
cast :: forall a b. Unbox b => MutArray a -> Maybe (MutArray b)
cast :: forall a b. Unbox b => MutArray a -> Maybe (MutArray b)
cast MutArray a
arr =
    let len :: Int
len = forall a. MutArray a -> Int
byteLength MutArray a
arr
        r :: Int
r = Int
len forall a. Integral a => a -> a -> a
`mod` SIZE_OF(b)
     in if Int
r forall a. Eq a => a -> a -> Bool
/= Int
0
        then forall a. Maybe a
Nothing
        else forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. MutArray a -> MutArray b
castUnsafe MutArray a
arr

-- XXX We can provide another API for "unsafe" FFI calls passing an unlifted
-- pointer to the FFI call. For unsafe calls we do not need to pin the array.
-- We can pass an unlifted pointer to the FFI routine to avoid GC kicking in
-- before the pointer is wrapped.
--
-- From the GHC manual:
--
-- GHC, since version 8.4, guarantees that garbage collection will never occur
-- during an unsafe call, even in the bytecode interpreter, and further
-- guarantees that unsafe calls will be performed in the calling thread. Making
-- it safe to pass heap-allocated objects to unsafe functions.

-- Should we just name it asPtr, the unsafety is implicit for any pointer
-- operations. And we are safe from Haskell perspective because we will be
-- pinning the memory.

-- | Use a @MutArray a@ as @Ptr a@. This is useful when we want to pass an
-- array as a pointer to some operating system call or to a "safe" FFI call.
--
-- If the array is not pinned it is copied to pinned memory before passing it
-- to the monadic action.
--
-- /Performance Notes:/ Forces a copy if the array is not pinned. It is advised
-- that the programmer keeps this in mind and creates a pinned array
-- opportunistically before this operation occurs, to avoid the cost of a copy
-- if possible.
--
-- /Unsafe/ because of direct pointer operations. The user must ensure that
-- they are writing within the legal bounds of the array.
--
-- /Pre-release/
--
{-# INLINE asPtrUnsafe #-}
asPtrUnsafe :: MonadIO m => MutArray a -> (Ptr a -> m b) -> m b
asPtrUnsafe :: forall (m :: * -> *) a b.
MonadIO m =>
MutArray a -> (Ptr a -> m b) -> m b
asPtrUnsafe MutArray a
arr Ptr a -> m b
f =
    forall (m :: * -> *) a b.
MonadIO m =>
MutByteArray -> (Ptr a -> m b) -> m b
Unboxed.asPtrUnsafe
        (forall a. MutArray a -> MutByteArray
arrContents MutArray a
arr) (\Ptr Any
ptr -> Ptr a -> m b
f (Ptr Any
ptr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` forall a. MutArray a -> Int
arrStart MutArray a
arr))

-------------------------------------------------------------------------------
-- Equality
-------------------------------------------------------------------------------

-- | Compare the length of the arrays. If the length is equal, compare the
-- lexicographical ordering of two underlying byte arrays otherwise return the
-- result of length comparison.
--
-- /Pre-release/
{-# INLINE cmp #-}
cmp :: MonadIO m => MutArray a -> MutArray a -> m Ordering
cmp :: forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> MutArray a -> m Ordering
cmp MutArray a
arr1 MutArray a
arr2 =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        forall a b. (a -> b) -> a -> b
$ do
            let marr1 :: MutableByteArray# RealWorld
marr1 = MutByteArray -> MutableByteArray# RealWorld
getMutableByteArray# (forall a. MutArray a -> MutByteArray
arrContents MutArray a
arr1)
                marr2 :: MutableByteArray# RealWorld
marr2 = MutByteArray -> MutableByteArray# RealWorld
getMutableByteArray# (forall a. MutArray a -> MutByteArray
arrContents MutArray a
arr2)
                !(I# Int#
st1#) = forall a. MutArray a -> Int
arrStart MutArray a
arr1
                !(I# Int#
st2#) = forall a. MutArray a -> Int
arrStart MutArray a
arr2
                !(I# Int#
len#) = forall a. MutArray a -> Int
byteLength MutArray a
arr1
            case forall a. Ord a => a -> a -> Ordering
compare (forall a. MutArray a -> Int
byteLength MutArray a
arr1) (forall a. MutArray a -> Int
byteLength MutArray a
arr2) of
                Ordering
EQ -> do
                    Int
r <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
                             let res :: Int
res =
                                     Int# -> Int
I#
                                         (ByteArray# -> Int# -> ByteArray# -> Int# -> Int# -> Int#
compareByteArrays#
                                              (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
marr1)
                                              Int#
st1#
                                              (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
marr2)
                                              Int#
st2#
                                              Int#
len#)
                              in (# State# RealWorld
s#, Int
res #)
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> a -> Ordering
compare Int
r Int
0
                Ordering
x -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
x

-------------------------------------------------------------------------------
-- NFData
-------------------------------------------------------------------------------

-- | Strip elements which match with predicate from both ends.
--
-- /Pre-release/
{-# INLINE strip #-}
strip :: forall a m. (Unbox a, MonadIO m) =>
    (a -> Bool) -> MutArray a -> m (MutArray a)
strip :: forall a (m :: * -> *).
(Unbox a, MonadIO m) =>
(a -> Bool) -> MutArray a -> m (MutArray a)
strip a -> Bool
eq arr :: MutArray a
arr@MutArray{Int
MutByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutByteArray
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    Int
st <- Int -> IO Int
getStart Int
arrStart
    Int
end <- Int -> Int -> IO Int
getLast Int
arrEnd Int
st
    forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr {arrStart :: Int
arrStart = Int
st, arrEnd :: Int
arrEnd = Int
end, arrBound :: Int
arrBound = Int
end}

    where

    {-
    -- XXX This should have the same perf but it does not, investigate.
    getStart = do
        r <- liftIO $ D.head $ D.findIndices (not . eq) $ toStreamD arr
        pure $
            case r of
                Nothing -> arrEnd
                Just i -> PTR_INDEX(arrStart,i,a)
    -}

    getStart :: Int -> IO Int
getStart Int
cur = do
        if Int
cur forall a. Ord a => a -> a -> Bool
< Int
arrEnd
        then do
            a
r <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
cur MutByteArray
arrContents
            if a -> Bool
eq a
r
            then Int -> IO Int
getStart (INDEX_NEXT(cur,a))
            else forall (m :: * -> *) a. Monad m => a -> m a
return Int
cur
        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
cur

    getLast :: Int -> Int -> IO Int
getLast Int
cur Int
low = do
        if Int
cur forall a. Ord a => a -> a -> Bool
> Int
low
        then do
            let prev :: Int
prev = INDEX_PREV(cur,a)
            a
r <- forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
prev MutByteArray
arrContents
            if a -> Bool
eq a
r
            then Int -> Int -> IO Int
getLast Int
prev Int
low
            else forall (m :: * -> *) a. Monad m => a -> m a
return Int
cur
        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
cur

-- | Given an array sorted in ascending order except the last element being out
-- of order, use bubble sort to place the last element at the right place such
-- that the array remains sorted in ascending order.
--
-- /Pre-release/
{-# INLINE bubble #-}
bubble :: (MonadIO m, Unbox a) => (a -> a -> Ordering) -> MutArray a -> m ()
bubble :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(a -> a -> Ordering) -> MutArray a -> m ()
bubble a -> a -> Ordering
cmp0 MutArray a
arr =
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
l forall a. Ord a => a -> a -> Bool
> Int
1) forall a b. (a -> b) -> a -> b
$ do
        a
x <- forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m a
getIndexUnsafe (Int
l forall a. Num a => a -> a -> a
- Int
1) MutArray a
arr
        forall {m :: * -> *}. MonadIO m => a -> Int -> m ()
go a
x (Int
l forall a. Num a => a -> a -> a
- Int
2)

        where

        l :: Int
l = forall a. Unbox a => MutArray a -> Int
length MutArray a
arr

        go :: a -> Int -> m ()
go a
x Int
i =
            if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0
            then do
                a
x1 <- forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
arr
                case a
x a -> a -> Ordering
`cmp0` a
x1 of
                    Ordering
LT -> do
                        forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndexUnsafe (Int
i forall a. Num a => a -> a -> a
+ Int
1) MutArray a
arr a
x1
                        a -> Int -> m ()
go a
x (Int
i forall a. Num a => a -> a -> a
- Int
1)
                    Ordering
_ -> forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndexUnsafe (Int
i forall a. Num a => a -> a -> a
+ Int
1) MutArray a
arr a
x
            else forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndexUnsafe (Int
i forall a. Num a => a -> a -> a
+ Int
1) MutArray a
arr a
x