{-# LANGUAGE CPP #-}
{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module      : Streamly.Internal.Data.Array.Mut.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.Array.Mut.Type
    (
    -- * Type
    -- $arrayNotes
      MutArray (..)
    , MutableByteArray
    , touch
    , pin
    , unpin

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

    -- *** Uninitialized Arrays
    , newPinned
    , newPinnedBytes
    , newAlignedPinned
    , new
    , newArrayWith

    -- *** Initialized Arrays
    , withNewArrayUnsafe

    -- *** From streams
    , ArrayUnsafe (..)
    , writeNWithUnsafe
    , writeNWith
    , writeNUnsafe
    , writeN
    , writeNAligned

    , writeWith
    , write

    , writeRevN
    -- , writeRev

    -- ** From containers
    , fromListN
    , fromList
    , fromListRevN
    , fromListRev
    , fromStreamDN
    , fromStreamD

    -- * 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
    , toStreamD
    , toStreamDRev
    , 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
    , arrayStreamKFromStreamD
    , 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

    -- ** Appending arrays
    , spliceCopy
    , spliceWith
    , splice
    , spliceExp
    , spliceUnsafe
    , putSliceUnsafe
    -- , 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.Proxy (Proxy(..))
import Data.Word (Word8)
import Foreign.C.Types (CSize(..), CInt(..))
import Foreign.Ptr (plusPtr, minusPtr, nullPtr)
import Streamly.Internal.Data.Unboxed
    ( MutableByteArray(..)
    , Unbox
    , getMutableByteArray#
    , peekWith
    , pokeWith
    , sizeOf
    , touch
    )
import GHC.Base
    ( IO(..)
    , Int(..)
    , byteArrayContents#
    , 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.StreamD.Type (Stream)
import Streamly.Internal.Data.Stream.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.Producer as Producer
import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
import qualified Streamly.Internal.Data.Unboxed as Unboxed
import qualified Prelude

import Prelude hiding
    (length, foldr, 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 :: a -> Int -> Int
bytesToElemCount a
_ Int
n = Int
n Int -> Int -> Int
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 = IO (Ptr Word8) -> IO ()
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 (Int -> CSize
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 (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ CInt
r CInt -> CInt -> Bool
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
    { MutArray a -> MutableByteArray
arrContents :: {-# UNPACK #-} !MutableByteArray
    , MutArray a -> Int
arrStart :: {-# UNPACK #-} !Int  -- ^ index into arrContents
    , MutArray a -> Int
arrEnd   :: {-# UNPACK #-} !Int    -- ^ index into arrContents
                                       -- Represents the first invalid index of
                                       -- the array.
    , MutArray a -> Int
arrBound :: {-# UNPACK #-} !Int    -- ^ first invalid index of arrContents.
    }

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

{-# INLINE pin #-}
pin :: MutArray a -> IO (MutArray a)
pin :: MutArray a -> IO (MutArray a)
pin arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    MutableByteArray
contents <- MutableByteArray -> IO MutableByteArray
Unboxed.pin MutableByteArray
arrContents
    MutArray a -> IO (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> IO (MutArray a)) -> MutArray a -> IO (MutArray a)
forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrContents :: MutableByteArray
arrContents = MutableByteArray
contents}

{-# INLINE unpin #-}
unpin :: MutArray a -> IO (MutArray a)
unpin :: MutArray a -> IO (MutArray a)
unpin arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    MutableByteArray
contents <- MutableByteArray -> IO MutableByteArray
Unboxed.unpin MutableByteArray
arrContents
    MutArray a -> IO (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> IO (MutArray a)) -> MutArray a -> IO (MutArray a)
forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrContents :: MutableByteArray
arrContents = MutableByteArray
contents}

-------------------------------------------------------------------------------
-- 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_newAlignedPinnedByteArrayzh 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 MutableByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith :: (Int -> Int -> m MutableByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith Int -> Int -> m MutableByteArray
alloc Int
alignSize Int
count = do
    let size :: Int
size = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* SIZE_OF(a)) 0
    MutableByteArray
contents <- Int -> Int -> m MutableByteArray
alloc Int
size Int
alignSize
    MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> m (MutArray a)) -> MutArray a -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ MutArray :: forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray
        { arrContents :: MutableByteArray
arrContents = MutableByteArray
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 :: MutArray a
nil = MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
Unboxed.nil Int
0 Int
0 Int
0


-- | 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 newPinnedBytes #-}
newPinnedBytes :: MonadIO m =>
#ifdef DEVBUILD
    Unbox a =>
#endif
    Int -> m (MutArray a)
newPinnedBytes :: Int -> m (MutArray a)
newPinnedBytes Int
bytes = do
    MutableByteArray
contents <- IO MutableByteArray -> m MutableByteArray
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO MutableByteArray -> m MutableByteArray)
-> IO MutableByteArray -> m MutableByteArray
forall a b. (a -> b) -> a -> b
$ Int -> IO MutableByteArray
Unboxed.newPinnedBytes Int
bytes
    MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> m (MutArray a)) -> MutArray a -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ MutArray :: forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray
        { arrContents :: MutableByteArray
arrContents = MutableByteArray
contents
        , arrStart :: Int
arrStart = Int
0
        , arrEnd :: Int
arrEnd   = Int
0
        , arrBound :: Int
arrBound = Int
bytes
        }

-- | 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 newAlignedPinned #-}
newAlignedPinned :: (MonadIO m, Unbox a) => Int -> Int -> m (MutArray a)
newAlignedPinned :: Int -> Int -> m (MutArray a)
newAlignedPinned =
    (Int -> Int -> m MutableByteArray) -> Int -> Int -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> m MutableByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith (\Int
s Int
a -> IO MutableByteArray -> m MutableByteArray
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO MutableByteArray -> m MutableByteArray)
-> IO MutableByteArray -> m MutableByteArray
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IO MutableByteArray
Unboxed.newAlignedPinnedBytes Int
s Int
a)

-- 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 newPinned #-}
newPinned :: forall m a. (MonadIO m, Unbox a) => Int -> m (MutArray a)
newPinned :: Int -> m (MutArray a)
newPinned =
    (Int -> Int -> m MutableByteArray) -> Int -> Int -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> m MutableByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith
        (\Int
s Int
_ -> IO MutableByteArray -> m MutableByteArray
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO MutableByteArray -> m MutableByteArray)
-> IO MutableByteArray -> m MutableByteArray
forall a b. (a -> b) -> a -> b
$ Int -> IO MutableByteArray
Unboxed.newPinnedBytes Int
s)
        ([Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"newPinned: alignSize is not used")

-- | 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 :: Int -> m (MutArray a)
new =
    (Int -> Int -> m MutableByteArray) -> Int -> Int -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> m MutableByteArray) -> Int -> Int -> m (MutArray a)
newArrayWith
        (\Int
s Int
_ -> IO MutableByteArray -> m MutableByteArray
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO MutableByteArray -> m MutableByteArray)
-> IO MutableByteArray -> m MutableByteArray
forall a b. (a -> b) -> a -> b
$ Int -> IO MutableByteArray
Unboxed.newUnpinnedBytes Int
s)
        ([Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"new: alignment is not used in unpinned arrays.")

-- XXX This should create a full length uninitialzed array so that the pointer
-- can be used.

-- | Allocate a pinned MutArray of the given size and run an IO action passing
-- the array start pointer.
--
-- /Internal/
{-# INLINE withNewArrayUnsafe #-}
withNewArrayUnsafe ::
       (MonadIO m, Unbox a) => Int -> (Ptr a -> m ()) -> m (MutArray a)
withNewArrayUnsafe :: Int -> (Ptr a -> m ()) -> m (MutArray a)
withNewArrayUnsafe Int
count Ptr a -> m ()
f = do
    MutArray a
arr <- Int -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
newPinned Int
count
    MutArray a -> (Ptr a -> m (MutArray a)) -> m (MutArray a)
forall (m :: * -> *) a b.
MonadIO m =>
MutArray a -> (Ptr a -> m b) -> m b
asPtrUnsafe MutArray a
arr
        ((Ptr a -> m (MutArray a)) -> m (MutArray a))
-> (Ptr a -> m (MutArray a)) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> Ptr a -> m ()
f Ptr a
p m () -> m (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr

-------------------------------------------------------------------------------
-- 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 :: Int -> MutArray a -> a -> m ()
putIndexUnsafe Int
i MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} a
x = do
    let index :: Int
index = Int
INDEX_OF(arrStart, i, a)
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index, arrEnd, a)) (return ())
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
index a
x

invalidIndex :: String -> Int -> a
invalidIndex :: [Char] -> Int -> a
invalidIndex [Char]
label Int
i =
    [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$ [Char]
label [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
": invalid array index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
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 :: Int -> MutArray a -> a -> m ()
putIndex Int
i MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} a
x = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)
    then IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
index a
x
    else [Char] -> Int -> m ()
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 :: MutArray a -> Fold m (Int, a) ()
putIndices MutArray a
arr = (() -> (Int, a) -> m ()) -> m () -> Fold m (Int, a) ()
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' () -> (Int, a) -> m ()
forall (m :: * -> *). MonadIO m => () -> (Int, a) -> m ()
step (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())

    where

    step :: () -> (Int, a) -> m ()
step () (Int
i, a
x) = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> MutArray a -> a -> IO ()
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.
--
-- /Pre-release/
modifyIndexUnsafe :: forall m a b. (MonadIO m, Unbox a) =>
    Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndexUnsafe :: Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndexUnsafe Int
i MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} a -> (a, b)
f = IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ do
        let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
        Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_NEXT(index,a) <= arrEnd) (return ())
        a
r <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
index
        let (a
x, b
res) = a -> (a, b)
f a
r
        MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
index a
x
        b -> IO b
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 :: Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndex Int
i MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} a -> (a, b)
f = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)
    then IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ do
        a
r <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
index
        let (a
x, b
res) = a -> (a, b)
f a
r
        MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
index a
x
        b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
res
    else [Char] -> Int -> m b
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 :: MutArray a -> (Int -> a -> a) -> Fold m Int ()
modifyIndices MutArray a
arr Int -> a -> a
f = (() -> Int -> m ()) -> m () -> Fold m Int ()
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' () -> Int -> m ()
forall (m :: * -> *). MonadIO m => () -> Int -> m ()
step m ()
initial

    where

    initial :: m ()
initial = () -> m ()
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 Int -> MutArray a -> (a -> (a, ())) -> m ()
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.
--
-- /Pre-release/
modify :: forall m a. (MonadIO m, Unbox a)
    => MutArray a -> (a -> a) -> m ()
modify :: MutArray a -> (a -> a) -> m ()
modify MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} a -> a
f = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
    Int -> IO ()
go Int
arrStart

    where

    go :: Int -> IO ()
go Int
i =
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (INDEX_VALID(i,arrEnd,a)) $ do
            r <- peekWith arrContents i
            pokeWith arrContents i (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
    -> MutableByteArray
    -> Int
    -> Int
    -> IO ()
swapArrayByteIndices :: Proxy a -> MutableByteArray -> Int -> Int -> IO ()
swapArrayByteIndices Proxy a
_ MutableByteArray
arrContents Int
i1 Int
i2 = do
    a
r1 <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
i1
    a
r2 <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
i2
    MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
i1 (a
r2 :: a)
    MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
i2 (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 :: Int -> Int -> MutArray a -> m ()
unsafeSwapIndices Int
i1 Int
i2 MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
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)
        Proxy a -> MutableByteArray -> Int -> Int -> IO ()
forall a.
Unbox a =>
Proxy a -> MutableByteArray -> Int -> Int -> IO ()
swapArrayByteIndices (Proxy a
forall k (t :: k). Proxy t
Proxy :: Proxy a) MutableByteArray
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 :: Int -> Int -> MutArray a -> m ()
swapIndices Int
i1 Int
i2 MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
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)
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| INDEX_INVALID(t1,arrEnd,a))
            (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> Int -> IO ()
forall a. [Char] -> Int -> a
invalidIndex [Char]
"swapIndices" Int
i1
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| INDEX_INVALID(t2,arrEnd,a))
            (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> Int -> IO ()
forall a. [Char] -> Int -> a
invalidIndex [Char]
"swapIndices" Int
i2
        Proxy a -> MutableByteArray -> Int -> Int -> IO ()
forall a.
Unbox a =>
Proxy a -> MutableByteArray -> Int -> Int -> IO ()
swapArrayByteIndices (Proxy a
forall k (t :: k). Proxy t
Proxy :: Proxy a) MutableByteArray
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 Int -> Int -> Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8) Int -> Int -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
largeObjectThreshold
    then
        Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert
            (Int
blockSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& ((Int
blockSize Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
blockSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0))
            ((Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
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 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Bool
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
z6
#else
    1 + z5
#endif

    where

    z0 :: Int
z0 = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    z1 :: Int
z1 = Int
z0 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
z0 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
1
    z2 :: Int
z2 = Int
z1 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
z1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
2
    z3 :: Int
z3 = Int
z2 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
z2 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
4
    z4 :: Int
z4 = Int
z3 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
z3 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
8
    z5 :: Int
z5 = Int
z4 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
z4 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
16
    z6 :: Int
z6 = Int
z5 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
z5 Int -> Int -> Int
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 :: a -> Int -> Int
allocBytesToBytes a
_ Int
n = Int -> Int -> Int
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 :: a -> Int -> Int
allocBytesToElemCount a
x Int
bytes =
    let n :: Int
n = a -> Int -> Int
forall a. Unbox a => a -> Int -> Int
bytesToElemCount a
x (a -> Int -> Int
forall a. Unbox a => a -> Int -> Int
allocBytesToBytes a
x Int
bytes)
     in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
size Int -> Int -> Int
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 :: Int -> Int -> MutArray a -> IO (MutArray a)
reallocExplicit Int
elemSize Int
newCapacityInBytes MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    assertM(Int
arrEnd Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrBound)

    -- Allocate new array
    let newCapMaxInBytes :: Int
newCapMaxInBytes = Int -> Int
roundUpLargeArray Int
newCapacityInBytes
    MutableByteArray
contents <- Int -> IO MutableByteArray
Unboxed.newPinnedBytes Int
newCapMaxInBytes
    let !(MutableByteArray MutableByteArray# RealWorld
mbarrFrom#) = MutableByteArray
arrContents
        !(MutableByteArray MutableByteArray# RealWorld
mbarrTo#) = MutableByteArray
contents

    -- Copy old data
    let oldStart :: Int
oldStart = Int
arrStart
        !(I# Int#
oldStartInBytes#) = Int
oldStart
        oldSizeInBytes :: Int
oldSizeInBytes = Int
arrEnd Int -> Int -> Int
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#) = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
oldSizeInBytes Int
newCapInBytes
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
oldSizeInBytes Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
elemSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
newLenInBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
newLenInBytes Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
elemSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# -> (# MutableByteArray# RealWorld
-> Int#
-> MutableByteArray# RealWorld
-> Int#
-> Int#
-> State# RealWorld
-> State# RealWorld
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#, () #)

    MutArray a -> IO (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> IO (MutArray a)) -> MutArray a -> IO (MutArray a)
forall a b. (a -> b) -> a -> b
$ MutArray :: forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray
        { arrStart :: Int
arrStart = Int
0
        , arrContents :: MutableByteArray
arrContents = MutableByteArray
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).
--
{-# INLINABLE realloc #-}
realloc :: forall m a. (MonadIO m, Unbox a) => Int -> MutArray a -> m (MutArray a)
realloc :: Int -> MutArray a -> m (MutArray a)
realloc Int
bytes MutArray a
arr = IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> MutArray a -> IO (MutArray a)
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 :: [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 = MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
arr
        newCapBytes :: Int
newCapBytes = Int -> Int
capSizer Int
oldSizeBytes
        newSizeBytes :: Int
newSizeBytes = Int
oldSizeBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
minIncrBytes
        safeCapBytes :: Int
safeCapBytes = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
newCapBytes Int
newSizeBytes
    assertM(Int
safeCapBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
newSizeBytes Bool -> Bool -> Bool
|| [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error (Int -> [Char]
forall a. Show a => a -> [Char]
badSize Int
newSizeBytes))

    Int -> MutArray a -> m (MutArray a)
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 =
        [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
            [ [Char]
label
            , [Char]
": new array size (in bytes) is less than required size "
            , a -> [Char]
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 :: Int -> MutArray a -> m (MutArray a)
resize Int
nElems arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    let req :: Int
req = SIZE_OF(a) * nElems
        len :: Int
len = Int
arrEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
arrStart
    if Int
req Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
    then MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr
    else Int -> MutArray a -> m (MutArray a)
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 :: Int -> MutArray a -> m (MutArray a)
resizeExp Int
nElems arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    let req :: Int
req = Int -> Int
roundUpLargeArray (SIZE_OF(a) * nElems)
        req1 :: Int
req1 =
            if Int
req Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
largeObjectThreshold
            then Int -> Int
roundUpToPower2 Int
req
            else Int
req
        len :: Int
len = Int
arrEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
arrStart
    if Int
req1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
    then MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr
    else Int -> MutArray a -> m (MutArray a)
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 :: MutArray a -> m (MutArray a)
rightSize arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
arrEnd Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrBound) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let start :: Int
start = Int
arrStart
        len :: Int
len = Int
arrEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start
        capacity :: Int
capacity = Int
arrBound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start
        target :: Int
target = Int -> Int
roundUpLargeArray Int
len
        waste :: Int
waste = Int
arrBound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
arrEnd
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
target Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
len Int -> Int -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
capacity Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
waste
    then Int -> MutArray a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
target MutArray a
arr
    else MutArray a -> m (MutArray a)
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 :: Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd Int
newEnd arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} a
x = IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ do
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
newEnd Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrBound) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
arrEnd a
x
    MutArray a -> IO (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> IO (MutArray a)) -> MutArray a -> IO (MutArray a)
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 :: MutArray a -> a -> m (MutArray a)
snocUnsafe arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = Int -> MutArray a -> a -> m (MutArray a)
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 :: MutArray a -> a -> m (Maybe (MutArray a))
snocMay arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} a
x = IO (Maybe (MutArray a)) -> m (Maybe (MutArray a))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (MutArray a)) -> m (Maybe (MutArray a)))
-> IO (Maybe (MutArray a)) -> m (Maybe (MutArray a))
forall a b. (a -> b) -> a -> b
$ do
    let newEnd :: Int
newEnd = INDEX_NEXT(arrEnd,a)
    if Int
newEnd Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrBound
    then MutArray a -> Maybe (MutArray a)
forall a. a -> Maybe a
Just (MutArray a -> Maybe (MutArray a))
-> IO (MutArray a) -> IO (Maybe (MutArray a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> MutArray a -> a -> IO (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd Int
newEnd MutArray a
arr a
x
    else Maybe (MutArray a) -> IO (Maybe (MutArray a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (MutArray a)
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 :: (Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWithRealloc Int -> Int
sizer MutArray a
arr a
x = do
    MutArray a
arr1 <- IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ [Char] -> (Int -> Int) -> Int -> MutArray a -> IO (MutArray a)
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
    MutArray a -> a -> m (MutArray a)
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 :: (Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith Int -> Int
allocSize MutArray a
arr a
x = IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ do
    let newEnd :: Int
newEnd = INDEX_NEXT(arrEnd arr,a)
    if Int
newEnd Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= MutArray a -> Int
forall a. MutArray a -> Int
arrBound MutArray a
arr
    then Int -> MutArray a -> a -> IO (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m (MutArray a)
snocNewEnd Int
newEnd MutArray a
arr a
x
    else (Int -> Int) -> MutArray a -> a -> IO (MutArray a)
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 :: MutArray a -> a -> m (MutArray a)
snocLinear = (Int -> Int) -> MutArray a -> a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int -> Int
forall a. Unbox a => a -> Int -> Int
allocBytesToBytes (a
forall a. HasCallStack => 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 :: MutArray a -> a -> m (MutArray a)
snoc = (Int -> Int) -> MutArray a -> a -> m (MutArray a)
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2
        else Int -> Int
roundUpToPower2 Int
oldSize Int -> Int -> Int
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 :: Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)) (return ())
    IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
index

-- | /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 a
getIndex :: Int -> MutArray a -> m a
getIndex Int
i MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    let index :: Int
index = Int
INDEX_OF(arrStart,i,a)
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(index,arrEnd,a)
    then IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
index
    else [Char] -> Int -> m a
forall a. [Char] -> Int -> a
invalidIndex [Char]
"getIndex" Int
i

-- | /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 :: Int -> MutArray a -> m a
getIndexRev Int
i MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = do
    let index :: Int
index = RINDEX_OF(Proxy a -> Int
forall a. Unbox a => Proxy a -> Int
arrEnd,i,a)
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
arrStart
    then IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
index
    else [Char] -> Int -> m a
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 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) = (GetIndicesState MutableByteArray Int Int s
 -> m (Step (GetIndicesState MutableByteArray Int Int s) a))
-> (MutArray a -> m (GetIndicesState MutableByteArray Int Int s))
-> Unfold m (MutArray a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold GetIndicesState MutableByteArray Int Int s
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
forall a.
Unbox a =>
GetIndicesState MutableByteArray Int Int s
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
step MutArray a -> m (GetIndicesState MutableByteArray Int Int s)
forall (m :: * -> *) a.
Monad m =>
MutArray a -> m (GetIndicesState MutableByteArray Int Int s)
inject

    where

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

    {-# INLINE_LATE step #-}
    step :: GetIndicesState MutableByteArray Int Int s
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
step (GetIndicesState MutableByteArray
contents Int
start Int
end s
st) = do
        Step s Int
r <- State StreamK m Int -> s -> m (Step s Int)
stepi State StreamK m Int
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s Int
r of
            D.Yield Int
i s
s -> do
                a
x <- IO a -> m a
forall b. IO b -> m b
liftio (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Int -> MutArray a -> IO a
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m a
getIndex Int
i (MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
contents Int
start Int
end Int
forall a. HasCallStack => a
undefined)
                Step (GetIndicesState MutableByteArray Int Int s) a
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GetIndicesState MutableByteArray Int Int s) a
 -> m (Step (GetIndicesState MutableByteArray Int Int s) a))
-> Step (GetIndicesState MutableByteArray Int Int s) a
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
forall a b. (a -> b) -> a -> b
$ a
-> GetIndicesState MutableByteArray Int Int s
-> Step (GetIndicesState MutableByteArray Int Int s) a
forall s a. a -> s -> Step s a
D.Yield a
x (MutableByteArray
-> Int -> Int -> s -> GetIndicesState MutableByteArray Int Int s
forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState MutableByteArray
contents Int
start Int
end s
s)
            D.Skip s
s -> Step (GetIndicesState MutableByteArray Int Int s) a
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GetIndicesState MutableByteArray Int Int s) a
 -> m (Step (GetIndicesState MutableByteArray Int Int s) a))
-> Step (GetIndicesState MutableByteArray Int Int s) a
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
forall a b. (a -> b) -> a -> b
$ GetIndicesState MutableByteArray Int Int s
-> Step (GetIndicesState MutableByteArray Int Int s) a
forall s a. s -> Step s a
D.Skip (MutableByteArray
-> Int -> Int -> s -> GetIndicesState MutableByteArray Int Int s
forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState MutableByteArray
contents Int
start Int
end s
s)
            Step s Int
D.Stop -> Step (GetIndicesState MutableByteArray Int Int s) a
-> m (Step (GetIndicesState MutableByteArray Int Int s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (GetIndicesState MutableByteArray Int Int s) a
forall s a. Step s a
D.Stop

{-# INLINE getIndices #-}
getIndices :: (MonadIO m, Unbox a) => Stream m Int -> Unfold m (MutArray a) a
getIndices :: Stream m Int -> Unfold m (MutArray a) a
getIndices = (forall b. IO b -> m b) -> Stream m Int -> Unfold m (MutArray a) a
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
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 :: Int -> Int -> MutArray a -> MutArray a
getSliceUnsafe Int
index Int
len (MutArray MutableByteArray
contents Int
start Int
e Int
_) =
    let fp1 :: Int
fp1 = INDEX_OF(start,index,a)
        end :: Int
end = Int
fp1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
* SIZE_OF(a))
     in Bool -> MutArray a -> MutArray a
forall a. HasCallStack => Bool -> a -> a
assert
            (Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
end Int -> Int -> Bool
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.
            (MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
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 :: Int -> Int -> MutArray a -> MutArray a
getSlice Int
index Int
len (MutArray MutableByteArray
contents Int
start Int
e Int
_) =
    let fp1 :: Int
fp1 = INDEX_OF(start,index,a)
        end :: Int
end = Int
fp1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
* SIZE_OF(a))
     in if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
end Int -> Int -> Bool
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 MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
contents Int
fp1 Int
end Int
end
        else [Char] -> MutArray a
forall a. HasCallStack => [Char] -> a
error
                ([Char] -> MutArray a) -> [Char] -> MutArray a
forall a b. (a -> b) -> a -> b
$ [Char]
"getSlice: invalid slice, index "
                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
index [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" length " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
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 :: MutArray a -> m ()
reverse MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
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
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
h) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            Proxy a -> MutableByteArray -> Int -> Int -> IO ()
forall a.
Unbox a =>
Proxy a -> MutableByteArray -> Int -> Int -> IO ()
swapArrayByteIndices (Proxy a
forall k (t :: k). Proxy t
Proxy :: Proxy a) MutableByteArray
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 :: MutArray a -> m Bool
permute = MutArray a -> m Bool
forall a. HasCallStack => 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 :: (a -> Bool) -> MutArray a -> m (MutArray a, MutArray a)
partitionBy a -> Bool
f arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = IO (MutArray a, MutArray a) -> m (MutArray a, MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a, MutArray a) -> m (MutArray a, MutArray a))
-> IO (MutArray a, MutArray a) -> m (MutArray a, MutArray a)
forall a b. (a -> b) -> a -> b
$ do
    if Int
arrStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
arrEnd
    then (MutArray a, MutArray a) -> IO (MutArray a, MutArray a)
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 = MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
arrContents Int
arrStart Int
ptr Int
ptr
            pr :: MutArray a
pr = MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
arrContents Int
ptr Int
arrEnd Int
arrEnd
        (MutArray a, MutArray a) -> IO (MutArray a, MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a
forall a. MutArray a
pl, MutArray a
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 <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
high
        if a -> Bool
f a
h
        then
            -- Correctly classified, continue the loop
            let high1 :: Int
high1 = INDEX_PREV(high,a)
             in if Int
low Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
high1
                then Maybe (Int, a) -> IO (Maybe (Int, a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Int, a)
forall a. Maybe a
Nothing
                else Int -> Int -> IO (Maybe (Int, a))
moveHigh Int
low Int
high1
        else Maybe (Int, a) -> IO (Maybe (Int, a))
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int, a) -> Maybe (Int, a)
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 <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
low
        if a -> Bool
f a
l
        then
            -- low is wrongly classified
            if Int
low Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
high
            then Int -> IO Int
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 -> Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
low
                    Just (Int
high1, a
h) -> do -- low < high1
                        MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
low a
h
                        MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
arrContents Int
high1 a
l
                        let low1 :: Int
low1 = INDEX_NEXT(low,a)
                            high2 :: Int
high2 = INDEX_PREV(high1,a)
                        if Int
low1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
high2
                        then Int -> Int -> IO Int
go Int
low1 Int
high2
                        else Int -> IO Int
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
high
            then Int -> IO Int
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 :: (a -> a -> m Bool) -> MutArray a -> MutArray a -> m ()
shuffleBy = (a -> a -> m Bool) -> MutArray a -> MutArray a -> m ()
forall a. HasCallStack => 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 :: Int
-> (MutArray a -> m (MutArray a, MutArray a)) -> MutArray a -> m ()
divideBy = Int
-> (MutArray a -> m (MutArray a, MutArray a)) -> MutArray a -> m ()
forall a. HasCallStack => 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 :: Int -> (MutArray a -> MutArray a -> m ()) -> MutArray a -> m ()
mergeBy = Int -> (MutArray a -> MutArray a -> m ()) -> MutArray a -> m ()
forall a. HasCallStack => a
undefined

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

-- | /O(1)/ Get the byte length of the array.
--
{-# INLINE byteLength #-}
byteLength :: MutArray a -> Int
byteLength :: MutArray a -> Int
byteLength MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} =
    let len :: Int
len = Int
arrEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
arrStart
    in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
len Int -> Int -> Bool
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 :: MutArray a -> Int
length MutArray a
arr =
    let elemSize :: Int
elemSize = SIZE_OF(a)
        blen :: Int
blen = MutArray a -> Int
forall a. MutArray a -> Int
byteLength MutArray a
arr
     in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
blen Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
elemSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (Int
blen Int -> Int -> Int
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 :: MutArray a -> Int
byteCapacity MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} =
    let len :: Int
len = Int
arrBound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
arrStart
    in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
len Int -> Int -> Bool
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 :: MutArray a -> Int
bytesFree MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} =
    let n :: Int
n = Int
arrBound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
arrEnd
    in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
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

-- | @chunksOf n stream@ groups the input stream into a stream of
-- arrays of size n.
--
-- @chunksOf n = StreamD.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 :: Int -> Stream m a -> Stream m (MutArray a)
chunksOf Int
n (D.Stream State StreamK m a -> s -> m (Step s a)
step s
state) =
    (State StreamK m (MutArray a)
 -> GroupState s MutableByteArray Int Int Int
 -> m (Step
         (GroupState s MutableByteArray Int Int Int) (MutArray a)))
-> GroupState s MutableByteArray Int Int Int
-> Stream m (MutArray a)
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State StreamK m (MutArray a)
-> GroupState s MutableByteArray Int Int Int
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall (m :: * -> *) a a.
State StreamK m a
-> GroupState s MutableByteArray Int Int Int
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
step' (s -> GroupState s MutableByteArray Int Int Int
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 MutableByteArray Int Int Int
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
step' State StreamK m a
_ (GroupStart s
st) = do
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
            -- XXX we can pass the module string from the higher level API
            [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"Streamly.Internal.Data.MutArray.Mut.Type.chunksOf: "
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"the size of arrays [" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"] must be a natural number"
        (MutArray MutableByteArray
contents Int
start Int
end Int
bound :: MutArray a) <- IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
newPinned Int
n
        Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
 -> m (Step
         (GroupState s MutableByteArray Int Int Int) (MutArray a)))
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall a b. (a -> b) -> a -> b
$ GroupState s MutableByteArray Int Int Int
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
forall s a. s -> Step s a
D.Skip (s
-> MutableByteArray
-> Int
-> Int
-> Int
-> GroupState s MutableByteArray Int Int Int
forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
st MutableByteArray
contents Int
start Int
end Int
bound)

    step' State StreamK m a
gst (GroupBuffer s
st MutableByteArray
contents Int
start Int
end Int
bound) = do
        Step s a
r <- State StreamK m a -> s -> m (Step s a)
step (State StreamK m a -> State StreamK m a
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
                IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
contents Int
end a
x
                let end1 :: Int
end1 = INDEX_NEXT(end,a)
                Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
 -> m (Step
         (GroupState s MutableByteArray Int Int Int) (MutArray a)))
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall a b. (a -> b) -> a -> b
$
                    if Int
end1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bound
                    then GroupState s MutableByteArray Int Int Int
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
forall s a. s -> Step s a
D.Skip
                            (MutableByteArray
-> Int
-> Int
-> Int
-> GroupState s MutableByteArray Int Int Int
-> GroupState s MutableByteArray Int Int Int
forall s contents start end bound.
contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
-> GroupState s contents start end bound
GroupYield
                                MutableByteArray
contents Int
start Int
end1 Int
bound (s -> GroupState s MutableByteArray Int Int Int
forall s contents start end bound.
s -> GroupState s contents start end bound
GroupStart s
s))
                    else GroupState s MutableByteArray Int Int Int
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
forall s a. s -> Step s a
D.Skip (s
-> MutableByteArray
-> Int
-> Int
-> Int
-> GroupState s MutableByteArray Int Int Int
forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
s MutableByteArray
contents Int
start Int
end1 Int
bound)
            D.Skip s
s ->
                Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
 -> m (Step
         (GroupState s MutableByteArray Int Int Int) (MutArray a)))
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall a b. (a -> b) -> a -> b
$ GroupState s MutableByteArray Int Int Int
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
forall s a. s -> Step s a
D.Skip (s
-> MutableByteArray
-> Int
-> Int
-> Int
-> GroupState s MutableByteArray Int Int Int
forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
s MutableByteArray
contents Int
start Int
end Int
bound)
            Step s a
D.Stop ->
                Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall (m :: * -> *) a. Monad m => a -> m a
return
                    (Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
 -> m (Step
         (GroupState s MutableByteArray Int Int Int) (MutArray a)))
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
-> m (Step
        (GroupState s MutableByteArray Int Int Int) (MutArray a))
forall a b. (a -> b) -> a -> b
$ GroupState s MutableByteArray Int Int Int
-> Step (GroupState s MutableByteArray Int Int Int) (MutArray a)
forall s a. s -> Step s a
D.Skip (MutableByteArray
-> Int
-> Int
-> Int
-> GroupState s MutableByteArray Int Int Int
-> GroupState s MutableByteArray Int Int Int
forall s contents start end bound.
contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
-> GroupState s contents start end bound
GroupYield MutableByteArray
contents Int
start Int
end Int
bound GroupState s MutableByteArray Int Int Int
forall s contents start end bound.
GroupState s contents start end bound
GroupFinish)

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

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

-- XXX buffer to a list instead?
-- | Buffer the stream into arrays in memory.
{-# INLINE arrayStreamKFromStreamD #-}
arrayStreamKFromStreamD :: forall m a. (MonadIO m, Unbox a) =>
    D.Stream m a -> m (StreamK m (MutArray a))
arrayStreamKFromStreamD :: Stream m a -> m (StreamK m (MutArray a))
arrayStreamKFromStreamD =
    let n :: Int
n = a -> Int -> Int
forall a. Unbox a => a -> Int -> Int
allocBytesToElemCount (a
forall a. HasCallStack => a
undefined :: a) Int
defaultChunkSize
     in (MutArray a -> StreamK m (MutArray a) -> StreamK m (MutArray a))
-> StreamK m (MutArray a)
-> Stream m (MutArray a)
-> m (StreamK m (MutArray a))
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr MutArray a -> StreamK m (MutArray a) -> StreamK m (MutArray a)
forall a (m :: * -> *). a -> StreamK m a -> StreamK m a
K.cons StreamK m (MutArray a)
forall (m :: * -> *) a. StreamK m a
K.nil (Stream m (MutArray a) -> m (StreamK m (MutArray a)))
-> (Stream m a -> Stream m (MutArray a))
-> Stream m a
-> m (StreamK m (MutArray a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Stream m a -> Stream m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> Stream m (MutArray a)
chunksOf 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 :: Stream m (MutArray a) -> Stream m a
flattenArrays (D.Stream State StreamK m (MutArray a) -> s -> m (Step s (MutArray a))
step s
state) = (State StreamK m a
 -> FlattenState s MutableByteArray Any
 -> m (Step (FlattenState s MutableByteArray Any) a))
-> FlattenState s MutableByteArray Any -> Stream m a
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State StreamK m a
-> FlattenState s MutableByteArray Any
-> m (Step (FlattenState s MutableByteArray Any) a)
forall a (m :: * -> *) a a a.
Unbox a =>
State StreamK m a
-> FlattenState s MutableByteArray a
-> m (Step (FlattenState s MutableByteArray a) a)
step' (s -> FlattenState s MutableByteArray Any
forall s contents a. s -> FlattenState s contents a
OuterLoop s
state)

    where

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

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

    step' State StreamK m a
_ (InnerLoop s
st MutableByteArray
contents Int
p Int
end) = do
        a
x <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
contents Int
p
        Step (FlattenState s MutableByteArray a) a
-> m (Step (FlattenState s MutableByteArray a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s MutableByteArray a) a
 -> m (Step (FlattenState s MutableByteArray a) a))
-> Step (FlattenState s MutableByteArray a) a
-> m (Step (FlattenState s MutableByteArray a) a)
forall a b. (a -> b) -> a -> b
$ a
-> FlattenState s MutableByteArray a
-> Step (FlattenState s MutableByteArray a) a
forall s a. a -> s -> Step s a
D.Yield a
x (s
-> MutableByteArray
-> Int
-> Int
-> FlattenState s MutableByteArray a
forall s contents a.
s -> contents -> Int -> Int -> FlattenState s contents a
InnerLoop s
st MutableByteArray
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 :: Stream m (MutArray a) -> Stream m a
flattenArraysRev (D.Stream State StreamK m (MutArray a) -> s -> m (Step s (MutArray a))
step s
state) = (State StreamK m a
 -> FlattenState s MutableByteArray Any
 -> m (Step (FlattenState s MutableByteArray Any) a))
-> FlattenState s MutableByteArray Any -> Stream m a
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State StreamK m a
-> FlattenState s MutableByteArray Any
-> m (Step (FlattenState s MutableByteArray Any) a)
forall a (m :: * -> *) a a a.
Unbox a =>
State StreamK m a
-> FlattenState s MutableByteArray a
-> m (Step (FlattenState s MutableByteArray a) a)
step' (s -> FlattenState s MutableByteArray Any
forall s contents a. s -> FlattenState s contents a
OuterLoop s
state)

    where

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

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

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

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

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

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

fromArrayUnsafe ::
#ifdef DEVBUILD
    Unbox a =>
#endif
    ArrayUnsafe a -> MutArray a
fromArrayUnsafe :: ArrayUnsafe a -> MutArray a
fromArrayUnsafe (ArrayUnsafe MutableByteArray
contents Int
start Int
end) =
         MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
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 b. IO b -> m b) -> Producer m (MutArray a) a
producerWith forall b. IO b -> m b
liftio = (ArrayUnsafe a -> m (Step (ArrayUnsafe a) a))
-> (MutArray a -> m (ArrayUnsafe a))
-> (ArrayUnsafe a -> m (MutArray a))
-> Producer m (MutArray a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
forall a a a.
Unbox a =>
ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step (ArrayUnsafe a -> m (ArrayUnsafe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayUnsafe a -> m (ArrayUnsafe a))
-> (MutArray a -> ArrayUnsafe a) -> MutArray a -> m (ArrayUnsafe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MutArray a -> ArrayUnsafe a
forall a. MutArray a -> ArrayUnsafe a
toArrayUnsafe) ArrayUnsafe a -> m (MutArray a)
forall a. ArrayUnsafe a -> m (MutArray a)
extract
    where

    {-# INLINE_LATE step #-}
    step :: ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step (ArrayUnsafe MutableByteArray
_ Int
cur Int
end)
        | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Int
cur Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
end) (Int
cur Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end) = Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ArrayUnsafe a) a
forall s a. Step s a
D.Stop
    step (ArrayUnsafe MutableByteArray
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 <- IO a -> m a
forall b. IO b -> m b
liftio (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
contents Int
cur
            Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a))
-> Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a)
forall a b. (a -> b) -> a -> b
$ a -> ArrayUnsafe a -> Step (ArrayUnsafe a) a
forall s a. a -> s -> Step s a
D.Yield a
x (MutableByteArray -> Int -> Int -> ArrayUnsafe a
forall a. MutableByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutableByteArray
contents (INDEX_NEXT(cur,a)) end)

    extract :: ArrayUnsafe a -> m (MutArray a)
extract = MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> m (MutArray a))
-> (ArrayUnsafe a -> MutArray a) -> ArrayUnsafe a -> m (MutArray a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ArrayUnsafe a -> MutArray a
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 :: Producer m (MutArray a) a
producer = (forall b. IO b -> m b) -> Producer m (MutArray a) a
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
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 :: Unfold m (MutArray a) a
reader = Producer m (MutArray a) a -> Unfold m (MutArray a) a
forall (m :: * -> *) a b. Producer m a b -> Unfold m a b
Producer.simplify Producer m (MutArray a) a
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 b. IO b -> m b) -> Unfold m (MutArray a) a
readerRevWith forall b. IO b -> m b
liftio = (ArrayUnsafe Any -> m (Step (ArrayUnsafe Any) a))
-> (MutArray a -> m (ArrayUnsafe Any)) -> Unfold m (MutArray a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ArrayUnsafe Any -> m (Step (ArrayUnsafe Any) a)
forall a a a.
Unbox a =>
ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step MutArray a -> m (ArrayUnsafe Any)
forall (m :: * -> *) a a.
Monad m =>
MutArray a -> m (ArrayUnsafe a)
inject
    where

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

    {-# INLINE_LATE step #-}
    step :: ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step (ArrayUnsafe MutableByteArray
_ Int
start Int
p) | Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start = Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ArrayUnsafe a) a
forall s a. Step s a
D.Stop
    step (ArrayUnsafe MutableByteArray
contents Int
start Int
p) = do
        !a
x <- IO a -> m a
forall b. IO b -> m b
liftio (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
contents Int
p
        Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a))
-> Step (ArrayUnsafe a) a -> m (Step (ArrayUnsafe a) a)
forall a b. (a -> b) -> a -> b
$ a -> ArrayUnsafe a -> Step (ArrayUnsafe a) a
forall s a. a -> s -> Step s a
D.Yield a
x (MutableByteArray -> Int -> Int -> ArrayUnsafe a
forall a. MutableByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutableByteArray
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 :: Unfold m (MutArray a) a
readerRev = (forall b. IO b -> m b) -> Unfold m (MutArray a) a
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
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 <- peekWith 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 :: MutArray a -> m [a]
toList MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = IO [a] -> m [a]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [a] -> m [a]) -> IO [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ Int -> IO [a]
forall a. Unbox a => Int -> IO [a]
go Int
arrStart
    where

    go :: Int -> IO [a]
go Int
p | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrEnd) (Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
arrEnd) = [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
    go Int
p = do
        a
x <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
p
        (:) a
x ([a] -> [a]) -> IO [a] -> IO [a]
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 b. IO b -> m b) -> MutArray a -> Stream m a
toStreamDWith forall b. IO b -> m b
liftio MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = (State StreamK m a -> Int -> m (Step Int a)) -> Int -> Stream m a
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State StreamK m a -> Int -> m (Step Int a)
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 | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrEnd) (Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
arrEnd) = Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step Int a
forall s a. Step s a
D.Stop
    step p
_ Int
p = IO (Step Int a) -> m (Step Int a)
forall b. IO b -> m b
liftio (IO (Step Int a) -> m (Step Int a))
-> IO (Step Int a) -> m (Step Int a)
forall a b. (a -> b) -> a -> b
$ do
        a
r <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
p
        Step Int a -> IO (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step Int a -> IO (Step Int a)) -> Step Int a -> IO (Step Int a)
forall a b. (a -> b) -> a -> b
$ a -> Int -> Step Int a
forall s a. a -> s -> Step s a
D.Yield a
r (INDEX_NEXT(p,a))

-- | Use the 'reader' unfold instead.
--
-- @toStreamD = D.unfold reader@
--
-- We can try this if the unfold has any performance issues.
{-# INLINE_NORMAL toStreamD #-}
toStreamD :: forall m a. (MonadIO m, Unbox a) => MutArray a -> D.Stream m a
toStreamD :: MutArray a -> Stream m a
toStreamD = (forall b. IO b -> m b) -> MutArray a -> Stream m a
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
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 b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKWith forall b. IO b -> m b
liftio MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = Int -> StreamK m a
forall a. Unbox a => Int -> StreamK m a
go Int
arrStart

    where

    go :: Int -> StreamK m a
go Int
p | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrEnd) (Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
arrEnd) = StreamK m a
forall (m :: * -> *) a. StreamK m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
p
        in IO a -> m a
forall b. IO b -> m b
liftio IO a
elemM m a -> StreamK m a -> StreamK m a
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 :: MutArray a -> StreamK m a
toStreamK = (forall b. IO b -> m b) -> MutArray a -> StreamK m a
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
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 b. IO b -> m b) -> MutArray a -> Stream m a
toStreamDRevWith forall b. IO b -> m b
liftio MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} =
    let p :: Int
p = INDEX_PREV(arrEnd,a)
    in (State StreamK m a -> Int -> m (Step Int a)) -> Int -> Stream m a
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State StreamK m a -> Int -> m (Step Int a)
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrStart = Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step Int a
forall s a. Step s a
D.Stop
    step p
_ Int
p = IO (Step Int a) -> m (Step Int a)
forall b. IO b -> m b
liftio (IO (Step Int a) -> m (Step Int a))
-> IO (Step Int a) -> m (Step Int a)
forall a b. (a -> b) -> a -> b
$ do
        a
r <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
p
        Step Int a -> IO (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step Int a -> IO (Step Int a)) -> Step Int a -> IO (Step Int a)
forall a b. (a -> b) -> a -> b
$ a -> Int -> Step Int a
forall s a. a -> s -> Step s a
D.Yield a
r (INDEX_PREV(p,a))

-- | Use the 'readerRev' unfold instead.
--
-- @toStreamDRev = D.unfold readerRev@
--
-- We can try this if the unfold has any perf issues.
{-# INLINE_NORMAL toStreamDRev #-}
toStreamDRev :: forall m a. (MonadIO m, Unbox a) => MutArray a -> D.Stream m a
toStreamDRev :: MutArray a -> Stream m a
toStreamDRev = (forall b. IO b -> m b) -> MutArray a -> Stream m a
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
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 b. IO b -> m b) -> MutArray a -> StreamK m a
toStreamKRevWith forall b. IO b -> m b
liftio MutArray {Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} =
    let p :: Int
p = INDEX_PREV(arrEnd,a)
    in Int -> StreamK m a
forall a. Unbox a => Int -> StreamK m a
go Int
p

    where

    go :: Int -> StreamK m a
go Int
p | Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrStart = StreamK m a
forall (m :: * -> *) a. StreamK m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
p
        in IO a -> m a
forall b. IO b -> m b
liftio IO a
elemM m a -> StreamK m a -> StreamK m a
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 :: MutArray a -> StreamK m a
toStreamKRev = (forall b. IO b -> m b) -> MutArray a -> StreamK m a
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
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' :: (b -> a -> b) -> b -> MutArray a -> m b
foldl' b -> a -> b
f b
z MutArray a
arr = (b -> a -> b) -> b -> Stream m a -> m b
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
D.foldl' b -> a -> b
f b
z (Stream m a -> m b) -> Stream m a -> m b
forall a b. (a -> b) -> a -> b
$ MutArray a -> Stream m a
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> Stream m a
toStreamD 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 :: (a -> b -> b) -> b -> MutArray a -> m b
foldr a -> b -> b
f b
z MutArray a
arr = (a -> b -> b) -> b -> Stream m a -> m b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr a -> b -> b
f b
z (Stream m a -> m b) -> Stream m a -> m b
forall a b. (a -> b) -> a -> b
$ MutArray a -> Stream m a
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> Stream m a
toStreamD 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 alloc@ 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 :: Int -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendNUnsafe Int
n m (MutArray a)
action =
    (ArrayUnsafe a -> MutArray a)
-> Fold m a (ArrayUnsafe a) -> Fold m a (MutArray a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ArrayUnsafe a -> MutArray a
forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe (Fold m a (ArrayUnsafe a) -> Fold m a (MutArray a))
-> Fold m a (ArrayUnsafe a) -> Fold m a (MutArray a)
forall a b. (a -> b) -> a -> b
$ (ArrayUnsafe a -> a -> m (ArrayUnsafe a))
-> m (ArrayUnsafe a) -> Fold m a (ArrayUnsafe a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' ArrayUnsafe a -> a -> m (ArrayUnsafe a)
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
        Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
        arr :: MutArray a
arr@(MutArray MutableByteArray
_ Int
_ Int
end Int
bound) <- m (MutArray a)
action
        let free :: Int
free = Int
bound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
end
            needed :: Int
needed = Int
n Int -> Int -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
needed
            then ([Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a))
-> [Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
forall a. a -> a
noinline [Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
[Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
reallocWith [Char]
"writeAppendNUnsafeWith" (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
needed) Int
needed MutArray a
arr
            else MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr
        ArrayUnsafe a -> m (ArrayUnsafe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayUnsafe a -> m (ArrayUnsafe a))
-> ArrayUnsafe a -> m (ArrayUnsafe a)
forall a b. (a -> b) -> a -> b
$ MutArray a -> ArrayUnsafe a
forall a. MutArray a -> ArrayUnsafe a
toArrayUnsafe MutArray a
arr1

    step :: ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step (ArrayUnsafe MutableByteArray
contents Int
start Int
end) a
x = do
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
contents Int
end a
x
        ArrayUnsafe a -> m (ArrayUnsafe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayUnsafe a -> m (ArrayUnsafe a))
-> ArrayUnsafe a -> m (ArrayUnsafe a)
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> Int -> ArrayUnsafe a
forall a. MutableByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutableByteArray
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 :: Int -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendN Int
n m (MutArray a)
initial = Int -> Fold m a (MutArray a) -> Fold m a (MutArray a)
forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n (Int -> m (MutArray a) -> Fold m a (MutArray a)
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 :: (Int -> Int) -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendWith Int -> Int
sizer = (MutArray a -> a -> m (MutArray a))
-> m (MutArray a) -> Fold m a (MutArray a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' ((Int -> Int) -> MutArray a -> a -> m (MutArray a)
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 :: m (MutArray a) -> Fold m a (MutArray a)
writeAppend = (Int -> Int) -> m (MutArray a) -> Fold m a (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int) -> m (MutArray a) -> Fold m a (MutArray a)
writeAppendWith (Int -> Int -> Int
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 :: (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWithUnsafe Int -> m (MutArray a)
alloc Int
n = ArrayUnsafe a -> MutArray a
forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe (ArrayUnsafe a -> MutArray a)
-> Fold m a (ArrayUnsafe a) -> Fold m a (MutArray a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ArrayUnsafe a -> a -> m (ArrayUnsafe a))
-> m (ArrayUnsafe a) -> Fold m a (ArrayUnsafe a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' ArrayUnsafe a -> a -> m (ArrayUnsafe a)
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 = MutArray a -> ArrayUnsafe a
forall a. MutArray a -> ArrayUnsafe a
toArrayUnsafe (MutArray a -> ArrayUnsafe a)
-> m (MutArray a) -> m (ArrayUnsafe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (MutArray a)
alloc (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
n Int
0)

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

-- | 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.newPinned
--
{-# INLINE_NORMAL writeNUnsafe #-}
writeNUnsafe :: forall m a. (MonadIO m, Unbox a)
    => Int -> Fold m a (MutArray a)
writeNUnsafe :: Int -> Fold m a (MutArray a)
writeNUnsafe = (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWithUnsafe Int -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
newPinned

-- | @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 :: (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWith Int -> m (MutArray a)
alloc Int
n = Int -> Fold m a (MutArray a) -> Fold m a (MutArray a)
forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n ((Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
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)

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

-- | 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 :: (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWithUnsafe Int -> m (MutArray a)
alloc Int
n = ArrayUnsafe a -> MutArray a
forall a. ArrayUnsafe a -> MutArray a
fromArrayUnsafe (ArrayUnsafe a -> MutArray a)
-> Fold m a (ArrayUnsafe a) -> Fold m a (MutArray a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ArrayUnsafe a -> a -> m (ArrayUnsafe a))
-> m (ArrayUnsafe a) -> Fold m a (ArrayUnsafe a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' ArrayUnsafe a -> a -> m (ArrayUnsafe a)
forall (m :: * -> *) a a a.
(MonadIO m, Unbox a) =>
ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step m (ArrayUnsafe a)
forall a. m (ArrayUnsafe a)
initial

    where

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

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

    step :: ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step (ArrayUnsafe MutableByteArray
contents Int
start Int
end) a
x = do
        let ptr :: Int
ptr = INDEX_PREV(start,a)
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> a -> IO ()
forall a. Unbox a => MutableByteArray -> Int -> a -> IO ()
pokeWith MutableByteArray
contents Int
ptr a
x
        ArrayUnsafe a -> m (ArrayUnsafe a)
forall (m :: * -> *) a. Monad m => a -> m a
return
          (ArrayUnsafe a -> m (ArrayUnsafe a))
-> ArrayUnsafe a -> m (ArrayUnsafe a)
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> Int -> Int -> ArrayUnsafe a
forall a. MutableByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutableByteArray
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 :: (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWith Int -> m (MutArray a)
alloc Int
n = Int -> Fold m a (MutArray a) -> Fold m a (MutArray a)
forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n ((Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
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 :: Int -> Fold m a (MutArray a)
writeRevN = (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeRevNWith Int -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
newPinned

-- | @writeNAligned align n@ folds a maximum of @n@ elements from the input
-- stream to a 'MutArray' aligned to the given size.
--
-- >>> writeNAligned align = MutArray.writeNWith (MutArray.newAlignedPinned align)
-- >>> writeNAligned align n = MutArray.writeAppendN n (MutArray.newAlignedPinned align n)
--
-- /Pre-release/
--
{-# INLINE_NORMAL writeNAligned #-}
writeNAligned :: forall m a. (MonadIO m, Unbox a)
    => Int -> Int -> Fold m a (MutArray a)
writeNAligned :: Int -> Int -> Fold m a (MutArray a)
writeNAligned Int
align = (Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> m (MutArray a)) -> Int -> Fold m a (MutArray a)
writeNWith (Int -> Int -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Int -> m (MutArray a)
newAlignedPinned 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 :: Int -> Fold m a (StreamK n (MutArray a))
writeChunks Int
n = Fold m a (MutArray a)
-> Fold m (MutArray a) (StreamK n (MutArray a))
-> Fold m a (StreamK n (MutArray a))
forall (m :: * -> *) a b c.
Monad m =>
Fold m a b -> Fold m b c -> Fold m a c
FL.many (Int -> Fold m a (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeN Int
n) Fold m (MutArray a) (StreamK n (MutArray a))
forall (m :: * -> *) a (n :: * -> *).
Monad m =>
Fold m a (StreamK n a)
FL.toStreamK

-- 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.newPinned 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) (newPinned n)
writeWith :: Int -> Fold m a (MutArray a)
writeWith Int
elemCount =
    (MutArray a -> m (MutArray a))
-> Fold m a (MutArray a) -> Fold m a (MutArray a)
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 (Fold m a (MutArray a) -> Fold m a (MutArray a))
-> Fold m a (MutArray a) -> Fold m a (MutArray a)
forall a b. (a -> b) -> a -> b
$ (MutArray a -> a -> m (MutArray a))
-> m (MutArray a) -> Fold m a (MutArray a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' MutArray a -> a -> m (MutArray a)
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
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
elemCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"writeWith: elemCount is negative"
        IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
newPinned Int
elemCount

    step :: MutArray a -> a -> m (MutArray a)
step arr :: MutArray a
arr@(MutArray MutableByteArray
_ 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 = MutArray a -> a -> m (MutArray a)
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 = IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> (MutArray a -> IO (MutArray a)) -> MutArray a -> m (MutArray a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MutArray a -> IO (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutArray a -> m (MutArray a)
rightSize

-- | 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 :: Fold m a (MutArray a)
write = Int -> Fold m a (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeWith (a -> Int -> Int
forall a. Unbox a => a -> Int -> Int
allocBytesToElemCount (a
forall a. HasCallStack => a
undefined :: a) Int
arrayChunkBytes)

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

-- | 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 :: Int -> Stream m a -> m (MutArray a)
fromStreamDN Int
limit Stream m a
str = do
    (MutArray a
arr :: MutArray a) <- IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
newPinned Int
limit
    Int
end <- (Int -> a -> m Int) -> m Int -> Stream m a -> m Int
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> m b
D.foldlM' (MutableByteArray -> Int -> a -> m Int
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
MutableByteArray -> Int -> a -> m Int
fwrite (MutArray a -> MutableByteArray
forall a. MutArray a -> MutableByteArray
arrContents MutArray a
arr)) (Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
arr) (Stream m a -> m Int) -> Stream m a -> m Int
forall a b. (a -> b) -> a -> b
$ Int -> Stream m a -> Stream m a
forall (m :: * -> *) a.
Applicative m =>
Int -> Stream m a -> Stream m a
D.take Int
limit Stream m a
str
    MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> m (MutArray a)) -> MutArray a -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrEnd :: Int
arrEnd = Int
end}

    where

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

-- | 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 :: Int -> [a] -> m (MutArray a)
fromListN Int
n [a]
xs = Int -> Stream m a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> m (MutArray a)
fromStreamDN Int
n (Stream m a -> m (MutArray a)) -> Stream m a -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ [a] -> Stream m a
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 :: Int -> [a] -> m (MutArray a)
fromListRevN Int
n [a]
xs = Fold m a (MutArray a) -> Stream m a -> m (MutArray a)
forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold (Int -> Fold m a (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (MutArray a)
writeRevN Int
n) (Stream m a -> m (MutArray a)) -> Stream m a -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ [a] -> Stream m a
forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

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

{-# INLINE arrayStreamKLength #-}
arrayStreamKLength :: (Monad m, Unbox a) => StreamK m (MutArray a) -> m Int
arrayStreamKLength :: StreamK m (MutArray a) -> m Int
arrayStreamKLength StreamK m (MutArray a)
as = (Int -> Int -> Int) -> Int -> StreamK m Int -> m Int
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> StreamK m a -> m b
K.foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ((MutArray a -> Int) -> StreamK m (MutArray a) -> StreamK m Int
forall a b (m :: * -> *). (a -> b) -> StreamK m a -> StreamK m b
K.map MutArray a -> Int
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 :: StreamK m (MutArray a) -> m (MutArray a)
fromArrayStreamK StreamK m (MutArray a)
as = do
    Int
len <- StreamK m (MutArray a) -> m Int
forall (m :: * -> *) a.
(Monad m, Unbox a) =>
StreamK m (MutArray a) -> m Int
arrayStreamKLength StreamK m (MutArray a)
as
    Int -> Stream m a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> m (MutArray a)
fromStreamDN Int
len (Stream m a -> m (MutArray a)) -> Stream m a -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ Unfold m (MutArray a) a -> Stream m (MutArray a) -> Stream m a
forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> Stream m a -> Stream m b
D.unfoldMany Unfold m (MutArray a) a
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Unfold m (MutArray a) a
reader (Stream m (MutArray a) -> Stream m a)
-> Stream m (MutArray a) -> Stream m a
forall a b. (a -> b) -> a -> b
$ StreamK m (MutArray a) -> Stream m (MutArray a)
forall (m :: * -> *) a. Applicative m => StreamK m a -> Stream m a
D.fromStreamK StreamK m (MutArray a)
as

-- 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 :: Stream m a -> m (MutArray a)
fromStreamD Stream m a
m = Stream m a -> m (StreamK m (MutArray a))
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m a -> m (StreamK m (MutArray a))
arrayStreamKFromStreamD Stream m a
m m (StreamK m (MutArray a))
-> (StreamK m (MutArray a) -> m (MutArray a)) -> m (MutArray a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= StreamK m (MutArray a) -> m (MutArray a)
forall a (m :: * -> *).
(Unbox a, MonadIO m) =>
StreamK m (MutArray a) -> m (MutArray a)
fromArrayStreamK

-- | 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 :: [a] -> m (MutArray a)
fromList [a]
xs = Stream m a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m a -> m (MutArray a)
fromStreamD (Stream m a -> m (MutArray a)) -> Stream m a -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ [a] -> Stream m a
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 :: [a] -> m (MutArray a)
fromListRev [a]
xs = Int -> [a] -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> [a] -> m (MutArray a)
fromListRevN ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
Prelude.length [a]
xs) [a]
xs

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

-- | Put a sub range of a source array into a subrange of a destination array.
-- This is not safe as it does not check the bounds.
{-# INLINE putSliceUnsafe #-}
putSliceUnsafe :: MonadIO m => MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe :: MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
src Int
srcStartBytes MutArray a
dst Int
dstStartBytes Int
lenBytes = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    assertM(Int
lenBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= MutArray a -> Int
forall a. MutArray a -> Int
arrBound MutArray a
dst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
dstStartBytes)
    assertM(Int
lenBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
src Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
srcStartBytes)
    let !(I# Int#
srcStartBytes#) = Int
srcStartBytes
        !(I# Int#
dstStartBytes#) = Int
dstStartBytes
        !(I# Int#
lenBytes#) = Int
lenBytes
    let arrS# :: MutableByteArray# RealWorld
arrS# = MutableByteArray -> MutableByteArray# RealWorld
getMutableByteArray# (MutArray a -> MutableByteArray
forall a. MutArray a -> MutableByteArray
arrContents MutArray a
src)
        arrD# :: MutableByteArray# RealWorld
arrD# = MutableByteArray -> MutableByteArray# RealWorld
getMutableByteArray# (MutArray a -> MutableByteArray
forall a. MutArray a -> MutableByteArray
arrContents MutArray a
dst)
    (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# -> (# MutableByteArray# RealWorld
-> Int#
-> MutableByteArray# RealWorld
-> Int#
-> Int#
-> State# RealWorld
-> State# RealWorld
forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray#
                    MutableByteArray# RealWorld
arrS# Int#
srcStartBytes# MutableByteArray# RealWorld
arrD# Int#
dstStartBytes# Int#
lenBytes# State# RealWorld
s#
                , () #)

-- | Copy two arrays into a newly allocated array.
{-# INLINE spliceCopy #-}
spliceCopy :: forall m a. MonadIO m =>
#ifdef DEVBUILD
    Unbox a =>
#endif
    MutArray a -> MutArray a -> m (MutArray a)
spliceCopy :: MutArray a -> MutArray a -> m (MutArray a)
spliceCopy MutArray a
arr1 MutArray a
arr2 = IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ do
    let start1 :: Int
start1 = MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
arr1
        start2 :: Int
start2 = MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
arr2
        len1 :: Int
len1 = MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
arr1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start1
        len2 :: Int
len2 = MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
arr2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start2
    MutableByteArray
newArrContents <- IO MutableByteArray -> IO MutableByteArray
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO MutableByteArray -> IO MutableByteArray)
-> IO MutableByteArray -> IO MutableByteArray
forall a b. (a -> b) -> a -> b
$ Int -> IO MutableByteArray
Unboxed.newPinnedBytes (Int
len1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len2)
    let len :: Int
len = Int
len1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len2
        newArr :: MutArray a
newArr = MutableByteArray -> Int -> Int -> Int -> MutArray a
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
newArrContents Int
0 Int
len Int
len
    MutArray a -> Int -> MutArray a -> Int -> Int -> IO ()
forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
arr1 Int
start1 MutArray a
forall a. MutArray a
newArr Int
0 Int
len1
    MutArray a -> Int -> MutArray a -> Int -> Int -> IO ()
forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
arr2 Int
start2 MutArray a
forall a. MutArray a
newArr Int
len1 Int
len2
    MutArray a -> IO (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
forall a. MutArray a
newArr

-- | 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 :: MutArray a -> MutArray a -> m (MutArray a)
spliceUnsafe MutArray a
dst MutArray a
src =
    IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ do
         let startSrc :: Int
startSrc = MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
src
             srcLen :: Int
srcLen = MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
src Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
startSrc
             endDst :: Int
endDst = MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
dst
         assertM(Int
endDst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
srcLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= MutArray a -> Int
forall a. MutArray a -> Int
arrBound MutArray a
dst)
         MutArray a -> Int -> MutArray a -> Int -> Int -> IO ()
forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
src Int
startSrc MutArray a
dst Int
endDst Int
srcLen
         MutArray a -> IO (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArray a -> IO (MutArray a)) -> MutArray a -> IO (MutArray a)
forall a b. (a -> b) -> a -> b
$ MutArray a
dst {arrEnd :: Int
arrEnd = Int
endDst Int -> Int -> Int
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 :: (Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
spliceWith Int -> Int -> Int
sizer dst :: MutArray a
dst@(MutArray MutableByteArray
_ Int
start Int
end Int
bound) MutArray a
src = do
{-
    let f = writeAppendWith (`sizer` byteLength src) (return dst)
     in D.fold f (toStreamD src)
-}
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
end Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
bound) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let srcBytes :: Int
srcBytes = MutArray a -> Int
forall a. MutArray a -> Int
arrEnd MutArray a
src Int -> Int -> Int
forall a. Num a => a -> a -> a
- MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
src

    MutArray a
dst1 <-
        if Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
srcBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bound
        then do
            let dstBytes :: Int
dstBytes = Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start
                newSizeInBytes :: Int
newSizeInBytes = Int -> Int -> Int
sizer Int
dstBytes Int
srcBytes
            Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
newSizeInBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
dstBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
srcBytes)
                (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error
                    ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"splice: newSize is less than the total size "
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"of arrays being appended. Please check the "
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"sizer function passed."
            IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
forall a b. (a -> b) -> a -> b
$ Int -> MutArray a -> IO (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m (MutArray a)
realloc Int
newSizeInBytes MutArray a
dst
        else MutArray a -> m (MutArray a)
forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
dst
    MutArray a -> MutArray a -> m (MutArray a)
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 (+)
--
-- /Pre-release/
{-# INLINE splice #-}
splice :: (MonadIO m, Unbox a) => MutArray a -> MutArray a -> m (MutArray a)
splice :: MutArray a -> MutArray a -> m (MutArray a)
splice = (Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
spliceWith Int -> Int -> Int
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 :: MutArray a -> MutArray a -> m (MutArray a)
spliceExp = (Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
(Int -> Int -> Int) -> MutArray a -> MutArray a -> m (MutArray a)
spliceWith (\Int
l1 Int
l2 -> Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
l1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) (Int
l1 Int -> Int -> Int
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 :: Word8
-> MutArray Word8 -> m (MutArray Word8, Maybe (MutArray Word8))
breakOn Word8
sep arr :: MutArray Word8
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = MutArray Word8
-> (Ptr Word8 -> m (MutArray Word8, Maybe (MutArray Word8)))
-> m (MutArray Word8, Maybe (MutArray Word8))
forall (m :: * -> *) a b.
MonadIO m =>
MutArray a -> (Ptr a -> m b) -> m b
asPtrUnsafe MutArray Word8
arr ((Ptr Word8 -> m (MutArray Word8, Maybe (MutArray Word8)))
 -> m (MutArray Word8, Maybe (MutArray Word8)))
-> (Ptr Word8 -> m (MutArray Word8, Maybe (MutArray Word8)))
-> m (MutArray Word8, Maybe (MutArray Word8))
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> IO (MutArray Word8, Maybe (MutArray Word8))
-> m (MutArray Word8, Maybe (MutArray Word8))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray Word8, Maybe (MutArray Word8))
 -> m (MutArray Word8, Maybe (MutArray Word8)))
-> IO (MutArray Word8, Maybe (MutArray Word8))
-> m (MutArray Word8, Maybe (MutArray Word8))
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 (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CSize) -> Int -> CSize
forall a b. (a -> b) -> a -> b
$ MutArray Word8 -> Int
forall a. MutArray a -> Int
byteLength MutArray Word8
arr)
    let sepIndex :: Int
sepIndex = Ptr Word8
loc Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p
    (MutArray Word8, Maybe (MutArray Word8))
-> IO (MutArray Word8, Maybe (MutArray Word8))
forall (m :: * -> *) a. Monad m => a -> m a
return ((MutArray Word8, Maybe (MutArray Word8))
 -> IO (MutArray Word8, Maybe (MutArray Word8)))
-> (MutArray Word8, Maybe (MutArray Word8))
-> IO (MutArray Word8, Maybe (MutArray Word8))
forall a b. (a -> b) -> a -> b
$
        if Ptr Word8
loc Ptr Word8 -> Ptr Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word8
forall a. Ptr a
nullPtr
        then (MutArray Word8
arr, Maybe (MutArray Word8)
forall a. Maybe a
Nothing)
        else
            ( MutArray :: forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray
                { arrContents :: MutableByteArray
arrContents = MutableByteArray
arrContents
                , arrStart :: Int
arrStart = Int
arrStart
                , arrEnd :: Int
arrEnd = Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sepIndex -- exclude the separator
                , arrBound :: Int
arrBound = Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sepIndex
                }
            , MutArray Word8 -> Maybe (MutArray Word8)
forall a. a -> Maybe a
Just (MutArray Word8 -> Maybe (MutArray Word8))
-> MutArray Word8 -> Maybe (MutArray Word8)
forall a b. (a -> b) -> a -> b
$ MutArray :: forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray
                    { arrContents :: MutableByteArray
arrContents = MutableByteArray
arrContents
                    , arrStart :: Int
arrStart = Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
sepIndex Int -> Int -> Int
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 :: Int -> MutArray a -> (MutArray a, MutArray a)
splitAt Int
i arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} =
    let maxIndex :: Int
maxIndex = MutArray a -> Int
forall a. Unbox a => MutArray a -> Int
length MutArray a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    in  if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        then [Char] -> (MutArray a, MutArray a)
forall a. HasCallStack => [Char] -> a
error [Char]
"sliceAt: negative array index"
        else if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxIndex
             then [Char] -> (MutArray a, MutArray a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (MutArray a, MutArray a))
-> [Char] -> (MutArray a, MutArray a)
forall a b. (a -> b) -> a -> b
$ [Char]
"sliceAt: specified array index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i
                        [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is beyond the maximum index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
maxIndex
             else let off :: Int
off = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* SIZE_OF(a)
                      p :: Int
p = Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off
                in ( MutArray :: forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray
                  { arrContents :: MutableByteArray
arrContents = MutableByteArray
arrContents
                  , arrStart :: Int
arrStart = Int
arrStart
                  , arrEnd :: Int
arrEnd = Int
p
                  , arrBound :: Int
arrBound = Int
p
                  }
                , MutArray :: forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray
                  { arrContents :: MutableByteArray
arrContents = MutableByteArray
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 :: MutArray a -> MutArray b
castUnsafe (MutArray MutableByteArray
contents Int
start Int
end Int
bound) =
    MutableByteArray -> Int -> Int -> Int -> MutArray b
forall a. MutableByteArray -> Int -> Int -> Int -> MutArray a
MutArray MutableByteArray
contents Int
start Int
end Int
bound

-- | Cast an @MutArray a@ into an @MutArray Word8@.
--
asBytes :: MutArray a -> MutArray Word8
asBytes :: MutArray a -> MutArray Word8
asBytes = MutArray a -> MutArray Word8
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 :: MutArray a -> Maybe (MutArray b)
cast MutArray a
arr =
    let len :: Int
len = MutArray a -> Int
forall a. MutArray a -> Int
byteLength MutArray a
arr
        r :: Int
r = Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` SIZE_OF(b)
     in if Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
        then Maybe (MutArray b)
forall a. Maybe a
Nothing
        else MutArray b -> Maybe (MutArray b)
forall a. a -> Maybe a
Just (MutArray b -> Maybe (MutArray b))
-> MutArray b -> Maybe (MutArray b)
forall a b. (a -> b) -> a -> b
$ MutArray a -> MutArray 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.

-- Unsafe because of direct pointer operations. The user must ensure that they
-- are writing within the legal bounds of the array. 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 an @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/
--
-- /Pre-release/
--
asPtrUnsafe :: MonadIO m => MutArray a -> (Ptr a -> m b) -> m b
asPtrUnsafe :: MutArray a -> (Ptr a -> m b) -> m b
asPtrUnsafe MutArray a
arr Ptr a -> m b
f = do
  let contents :: MutableByteArray
contents = MutArray a -> MutableByteArray
forall a. MutArray a -> MutableByteArray
arrContents MutArray a
arr
      !ptr :: Ptr a
ptr = Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr (ByteArray# -> Addr#
byteArrayContents#
                     (MutableByteArray# RealWorld -> ByteArray#
unsafeCoerce# (MutableByteArray -> MutableByteArray# RealWorld
getMutableByteArray# MutableByteArray
contents)))
  -- XXX Check if the array is pinned, if not, copy it to a pinned array
  -- XXX We should probably pass to the IO action the byte length of the array
  -- as well so that bounds can be checked.
  b
r <- Ptr a -> m b
f (Ptr Any
forall a. Ptr a
ptr Ptr Any -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
arr)
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray -> IO ()
touch MutableByteArray
contents
  b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
r

-------------------------------------------------------------------------------
-- 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 :: MutArray a -> MutArray a -> m Ordering
cmp MutArray a
arr1 MutArray a
arr2 =
    IO Ordering -> m Ordering
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        (IO Ordering -> m Ordering) -> IO Ordering -> m Ordering
forall a b. (a -> b) -> a -> b
$ do
            let marr1 :: MutableByteArray# RealWorld
marr1 = MutableByteArray -> MutableByteArray# RealWorld
getMutableByteArray# (MutArray a -> MutableByteArray
forall a. MutArray a -> MutableByteArray
arrContents MutArray a
arr1)
                marr2 :: MutableByteArray# RealWorld
marr2 = MutableByteArray -> MutableByteArray# RealWorld
getMutableByteArray# (MutArray a -> MutableByteArray
forall a. MutArray a -> MutableByteArray
arrContents MutArray a
arr2)
                !(I# Int#
st1#) = MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
arr1
                !(I# Int#
st2#) = MutArray a -> Int
forall a. MutArray a -> Int
arrStart MutArray a
arr2
                !(I# Int#
len#) = MutArray a -> Int
forall a. MutArray a -> Int
byteLength MutArray a
arr1
            case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (MutArray a -> Int
forall a. MutArray a -> Int
byteLength MutArray a
arr1) (MutArray a -> Int
forall a. MutArray a -> Int
byteLength MutArray a
arr2) of
                Ordering
EQ -> do
                    Int
r <- IO Int -> IO Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> IO Int) -> IO Int -> IO Int
forall a b. (a -> b) -> a -> b
$ (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int)
-> (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
                             let res :: Int
res =
                                     Int# -> Int
I#
                                         (ByteArray# -> Int# -> ByteArray# -> Int# -> Int# -> Int#
compareByteArrays#
                                              (MutableByteArray# RealWorld -> ByteArray#
unsafeCoerce# MutableByteArray# RealWorld
marr1)
                                              Int#
st1#
                                              (MutableByteArray# RealWorld -> ByteArray#
unsafeCoerce# MutableByteArray# RealWorld
marr2)
                                              Int#
st2#
                                              Int#
len#)
                              in (# State# RealWorld
s#, Int
res #)
                    Ordering -> IO Ordering
forall (m :: * -> *) a. Monad m => a -> m a
return (Ordering -> IO Ordering) -> Ordering -> IO Ordering
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
r Int
0
                Ordering
x -> Ordering -> IO Ordering
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 :: (a -> Bool) -> MutArray a -> m (MutArray a)
strip a -> Bool
eq arr :: MutArray a
arr@MutArray{Int
MutableByteArray
arrBound :: Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
arrBound :: forall a. MutArray a -> Int
arrEnd :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents :: forall a. MutArray a -> MutableByteArray
..} = IO (MutArray a) -> m (MutArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MutArray a) -> m (MutArray a))
-> IO (MutArray a) -> m (MutArray a)
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
    MutArray a -> IO (MutArray a)
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrEnd
        then do
            a
r <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
cur
            if a -> Bool
eq a
r
            then Int -> IO Int
getStart (INDEX_NEXT(cur,a))
            else Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
cur
        else Int -> IO Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
low
        then do
            let prev :: Int
prev = INDEX_PREV(cur,a)
            a
r <- MutableByteArray -> Int -> IO a
forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
arrContents Int
prev
            if a -> Bool
eq a
r
            then Int -> Int -> IO Int
getLast Int
prev Int
low
            else Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
cur
        else Int -> IO Int
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 :: (a -> a -> Ordering) -> MutArray a -> m ()
bubble a -> a -> Ordering
cmp0 MutArray a
arr =
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
x <- Int -> MutArray a -> m a
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> m a
getIndexUnsafe (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) MutArray a
arr
        a -> Int -> m ()
forall (m :: * -> *). MonadIO m => a -> Int -> m ()
go a
x (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)

        where

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

        go :: a -> Int -> m ()
go a
x Int
i =
            if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            then do
                a
x1 <- Int -> MutArray a -> m a
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
                        Int -> MutArray a -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndexUnsafe (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MutArray a
arr a
x1
                        a -> Int -> m ()
go a
x (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                    Ordering
_ -> Int -> MutArray a -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndexUnsafe (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MutArray a
arr a
x
            else Int -> MutArray a -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> MutArray a -> a -> m ()
putIndexUnsafe (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MutArray a
arr a
x