{-# LANGUAGE CPP #-}
{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module      : Streamly.Internal.Data.MutArray.Generic
-- Copyright   : (c) 2020 Composewell Technologies
-- License     : BSD3-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.MutArray.Generic
(
    -- * Type
    -- $arrayNotes
      MutArray (..)

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

    -- *** Uninitialized Arrays
    , new
    -- , newArrayWith

    -- *** From streams
    , writeNUnsafe
    , writeN
    , writeWith
    , write
    , fromStreamN
    , fromStream
    , fromPureStream

    -- , writeRevN
    -- , writeRev

    -- ** From containers
    , fromListN
    , fromList

    -- * 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

    -- * 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.

    -- ** Reallocation
    , realloc
    , uninit

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

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

    -- ** Truncation
    -- These are not the same as slicing the array at the beginning, they may
    -- reduce the length as well as the capacity of the array.
    -- , truncateWith
    -- , truncate
    -- , truncateExp

    -- * Eliminating and Reading

    -- ** Unfolds
    , reader
    -- , readerRev
    , producerWith -- experimental
    , producer -- experimental

    -- ** To containers
    , read
    , readRev
    , toStreamK
    -- , toStreamKRev
    , toList

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

    -- * Size
    , length

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

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

    -- * 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
    , putSliceUnsafe
    -- , appendSlice
    -- , appendSliceFrom

    , clone
    )
where

#include "inline.hs"
#include "assert.hs"

import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Functor.Identity (Identity(..))
import GHC.Base
    ( MutableArray#
    , RealWorld
    , copyMutableArray#
    , newArray#
    , readArray#
    , writeArray#
    )
import GHC.IO (IO(..))
import GHC.Int (Int(..))
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Producer.Type (Producer (..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.Data.Stream.Type (Stream)
import Streamly.Internal.Data.SVar.Type (adaptState)

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

import Prelude hiding (read, length)

#include "DocTestDataMutArrayGeneric.hs"

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

data MutArray a =
    MutArray
        { forall a. MutArray a -> MutableArray# RealWorld a
arrContents# :: MutableArray# RealWorld a
          -- ^ The internal contents of the array representing the entire array.

        , forall a. MutArray a -> Int
arrStart :: {-# UNPACK #-}!Int
          -- ^ The starting index of this slice.

        , forall a. MutArray a -> Int
arrLen :: {-# UNPACK #-}!Int
          -- ^ The length of this slice.

        , forall a. MutArray a -> Int
arrTrueLen :: {-# UNPACK #-}!Int
          -- ^ This is the true length of the array. Coincidentally, this also
          -- represents the first index beyond the maximum acceptable index of
          -- the array. This is specific to the array contents itself and not
          -- dependent on the slice. This value should not change and is shared
          -- across all the slices.
        }

{-# INLINE bottomElement #-}
bottomElement :: a
bottomElement :: forall a. a
bottomElement =
    forall a. HasCallStack => [Char] -> a
error
        forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
unwords
              [ [Char]
funcName
              , [Char]
"This is the bottom element of the array."
              , [Char]
"This is a place holder and should never be reached!"
              ]

    where

    funcName :: [Char]
funcName = [Char]
"Streamly.Internal.Data.MutArray.Generic.bottomElement:"

-- XXX Would be nice if GHC can provide something like newUninitializedArray# so
-- that we do not have to write undefined or error in the whole array.

-- | @new count@ allocates a zero length array that can be extended to hold
-- up to 'count' items without reallocating.
--
-- /Pre-release/
{-# INLINE new #-}
new :: MonadIO m => Int -> m (MutArray a)
new :: forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new n :: Int
n@(I# Int#
n#) =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              case forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
newArray# Int#
n# forall a. a
bottomElement State# RealWorld
s# of
                  (# State# RealWorld
s1#, MutableArray# RealWorld a
arr# #) ->
                      let ma :: MutArray a
ma = forall a.
MutableArray# RealWorld a -> Int -> Int -> Int -> MutArray a
MutArray MutableArray# RealWorld a
arr# Int
0 Int
0 Int
n
                       in (# State# RealWorld
s1#, MutArray a
ma #)

-- XXX This could be pure?

-- |
-- Definition:
--
-- >>> nil = MutArray.new 0
{-# INLINE nil #-}
nil :: MonadIO m => m (MutArray a)
nil :: forall (m :: * -> *) a. MonadIO m => m (MutArray a)
nil = forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new Int
0

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

-- | Write the given element to the given index of the 'MutableArray#'. Does not
-- check if the index is out of bounds of the array.
--
-- /Pre-release/
{-# INLINE putIndexUnsafeWith #-}
putIndexUnsafeWith :: MonadIO m => Int -> MutableArray# RealWorld a -> a -> m ()
putIndexUnsafeWith :: forall (m :: * -> *) a.
MonadIO m =>
Int -> MutableArray# RealWorld a -> a -> m ()
putIndexUnsafeWith Int
n MutableArray# RealWorld a
_arrContents# a
x =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              case Int
n of
                  I# Int#
n# ->
                      let s1# :: State# RealWorld
s1# = forall d a. MutableArray# d a -> Int# -> a -> State# d -> State# d
writeArray# MutableArray# RealWorld a
_arrContents# Int#
n# a
x State# RealWorld
s#
                       in (# State# RealWorld
s1#, () #)

-- | 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 => Int -> MutArray a -> a -> m ()
putIndexUnsafe :: forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> a -> m ()
putIndexUnsafe Int
i MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} a
x =
    forall a. HasCallStack => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
< Int
arrLen)
    forall (m :: * -> *) a.
MonadIO m =>
Int -> MutableArray# RealWorld a -> a -> m ()
putIndexUnsafeWith (Int
i forall a. Num a => a -> a -> a
+ Int
arrStart) MutableArray# RealWorld a
arrContents# a
x

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

-- | /O(1)/ Write the given element at the given index in the array.
-- Performs in-place mutation of the array.
--
-- >>> putIndex ix arr val = MutArray.modifyIndex ix arr (const (val, ()))
--
-- /Pre-release/
{-# INLINE putIndex #-}
putIndex :: MonadIO m => Int -> MutArray a -> a -> m ()
putIndex :: forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> a -> m ()
putIndex Int
i arr :: MutArray a
arr@MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} a
x =
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
< Int
arrLen
    then forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> a -> m ()
putIndexUnsafe Int
i MutArray a
arr a
x
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"putIndex" Int
i

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

    where

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

-- | Modify a given index of an array using a modifier function without checking
-- the bounds.
--
-- Unsafe because it does not check the bounds of the array.
--
-- /Pre-release/
modifyIndexUnsafe :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndexUnsafe :: forall (m :: * -> *) a b.
MonadIO m =>
Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndexUnsafe Int
i MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} a -> (a, b)
f = do
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              case Int
i forall a. Num a => a -> a -> a
+ Int
arrStart of
                  I# Int#
n# ->
                      case forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# MutableArray# RealWorld a
arrContents# Int#
n# State# RealWorld
s# of
                          (# State# RealWorld
s1#, a
a #) ->
                              let (a
a1, b
b) = a -> (a, b)
f a
a
                                  s2# :: State# RealWorld
s2# = forall d a. MutableArray# d a -> Int# -> a -> State# d -> State# d
writeArray# MutableArray# RealWorld a
arrContents# Int#
n# a
a1 State# RealWorld
s1#
                               in (# State# RealWorld
s2#, b
b #)

-- | Modify a given index of an array using a modifier function.
--
-- /Pre-release/
modifyIndex :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndex :: forall (m :: * -> *) a b.
MonadIO m =>
Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndex Int
i arr :: MutArray a
arr@MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} a -> (a, b)
f = do
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
< Int
arrLen
    then forall (m :: * -> *) a b.
MonadIO m =>
Int -> MutArray a -> (a -> (a, b)) -> m b
modifyIndexUnsafe Int
i MutArray a
arr a -> (a, b)
f
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"modifyIndex" Int
i

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

-- | Reallocates the array according to the new size. This is a safe function
-- that always creates a new array and copies the old array into the new one.
-- If the reallocated size is less than the original array it results in a
-- truncated version of the original array.
--
realloc :: MonadIO m => Int -> MutArray a -> m (MutArray a)
realloc :: forall (m :: * -> *) a.
MonadIO m =>
Int -> MutArray a -> m (MutArray a)
realloc Int
n MutArray a
arr = do
    MutArray a
arr1 <- forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new Int
n
    let !newLen :: Int
newLen@(I# Int#
newLen#) = forall a. Ord a => a -> a -> a
min Int
n (forall a. MutArray a -> Int
arrLen MutArray a
arr)
        !(I# Int#
arrS#) = forall a. MutArray a -> Int
arrStart MutArray a
arr
        !(I# Int#
arr1S#) = forall a. MutArray a -> Int
arrStart MutArray a
arr1
        arrC# :: MutableArray# RealWorld a
arrC# = forall a. MutArray a -> MutableArray# RealWorld a
arrContents# MutArray a
arr
        arr1C# :: MutableArray# RealWorld a
arr1C# = forall a. MutArray a -> MutableArray# RealWorld a
arrContents# MutArray a
arr1
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              let s1# :: State# RealWorld
s1# = forall d a.
MutableArray# d a
-> Int#
-> MutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableArray# MutableArray# RealWorld a
arrC# Int#
arrS# MutableArray# RealWorld a
arr1C# Int#
arr1S# Int#
newLen# State# RealWorld
s#
               in (# State# RealWorld
s1#, MutArray a
arr1 {arrLen :: Int
arrLen = Int
newLen, arrTrueLen :: Int
arrTrueLen = Int
n} #)

reallocWith ::
       MonadIO m => String -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
reallocWith :: forall (m :: * -> *) a.
MonadIO m =>
[Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
reallocWith [Char]
label Int -> Int
sizer Int
reqSize MutArray a
arr = do
    let oldSize :: Int
oldSize = forall a. MutArray a -> Int
arrLen MutArray a
arr
        newSize :: Int
newSize = Int -> Int
sizer Int
oldSize
        safeSize :: Int
safeSize = forall a. Ord a => a -> a -> a
max Int
newSize (Int
oldSize forall a. Num a => a -> a -> a
+ Int
reqSize)
    forall a. HasCallStack => Bool -> a -> a
assert (Int
newSize forall a. Ord a => a -> a -> Bool
>= Int
oldSize forall a. Num a => a -> a -> a
+ Int
reqSize Bool -> Bool -> Bool
|| forall a. HasCallStack => [Char] -> a
error [Char]
badSize) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall (m :: * -> *) a.
MonadIO m =>
Int -> MutArray a -> m (MutArray a)
realloc Int
safeSize MutArray a
arr

    where

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

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

-- XXX Not sure of the behavior of writeArray# if we specify an index which is
-- out of bounds. This comment should be rewritten based on that.
-- | Really really unsafe, appends the element into the first array, may
-- cause silent data corruption or if you are lucky a segfault if the index
-- is out of bounds.
--
-- /Internal/
{-# INLINE snocUnsafe #-}
snocUnsafe :: MonadIO m => MutArray a -> a -> m (MutArray a)
snocUnsafe :: forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> a -> m (MutArray a)
snocUnsafe arr :: MutArray a
arr@MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} a
a = do
    forall a. HasCallStack => Bool -> a -> a
assert (Int
arrStart forall a. Num a => a -> a -> a
+ Int
arrLen forall a. Ord a => a -> a -> Bool
< Int
arrTrueLen) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let arr1 :: MutArray a
arr1 = MutArray a
arr {arrLen :: Int
arrLen = Int
arrLen forall a. Num a => a -> a -> a
+ Int
1}
    forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> a -> m ()
putIndexUnsafe Int
arrLen MutArray a
arr1 a
a
    forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr1

-- NOINLINE to move it out of the way and not pollute the instruction cache.
{-# NOINLINE snocWithRealloc #-}
snocWithRealloc :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWithRealloc :: forall (m :: * -> *) a.
MonadIO m =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWithRealloc Int -> Int
sizer MutArray a
arr a
x = do
    MutArray a
arr1 <- forall (m :: * -> *) a.
MonadIO m =>
[Char] -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
reallocWith [Char]
"snocWithRealloc" Int -> Int
sizer Int
1 MutArray a
arr
    forall (m :: * -> *) a.
MonadIO m =>
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 oldSize@ function, where @oldSize@ is the
-- original size of the array.
--
-- Note that the returned array may be a mutated version of the original array.
--
-- /Pre-release/
{-# INLINE snocWith #-}
snocWith :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith :: forall (m :: * -> *) a.
MonadIO m =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith Int -> Int
sizer arr :: MutArray a
arr@MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} a
x = do
    if Int
arrStart forall a. Num a => a -> a -> a
+ Int
arrLen forall a. Ord a => a -> a -> Bool
< Int
arrTrueLen
    then forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> a -> m (MutArray a)
snocUnsafe MutArray a
arr a
x
    else forall (m :: * -> *) a.
MonadIO m =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWithRealloc Int -> Int
sizer MutArray a
arr a
x

-- XXX round it to next power of 2.

-- | 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.
--
-- /Pre-release/
{-# INLINE snoc #-}
snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)
snoc :: forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> a -> m (MutArray a)
snoc = forall (m :: * -> *) a.
MonadIO m =>
(Int -> Int) -> MutArray a -> a -> m (MutArray a)
snocWith (forall a. Num a => a -> a -> a
* Int
2)

-- | Make the uninitialized memory in the array available for use extending it
-- by the supplied length beyond the current length of the array. The array may
-- be reallocated.
--
{-# INLINE uninit #-}
uninit :: MonadIO m => MutArray a -> Int -> m (MutArray a)
uninit :: forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> m (MutArray a)
uninit arr :: MutArray a
arr@MutArray{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} Int
len =
    if Int
arrStart forall a. Num a => a -> a -> a
+ Int
arrLen forall a. Num a => a -> a -> a
+ Int
len forall a. Ord a => a -> a -> Bool
<= Int
arrTrueLen
    then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrLen :: Int
arrLen = Int
arrLen forall a. Num a => a -> a -> a
+ Int
len}
    else forall (m :: * -> *) a.
MonadIO m =>
Int -> MutArray a -> m (MutArray a)
realloc (Int
arrLen forall a. Num a => a -> a -> a
+ Int
len) MutArray a
arr

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

-- | Return the element at the specified index without checking the bounds from
-- a @MutableArray# RealWorld@.
--
-- Unsafe because it does not check the bounds of the array.
{-# INLINE getIndexUnsafeWith #-}
getIndexUnsafeWith :: MonadIO m => MutableArray# RealWorld a -> Int -> m a
getIndexUnsafeWith :: forall (m :: * -> *) a.
MonadIO m =>
MutableArray# RealWorld a -> Int -> m a
getIndexUnsafeWith MutableArray# RealWorld a
_arrContents# Int
n =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              let !(I# Int#
i#) = Int
n
               in forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# MutableArray# RealWorld a
_arrContents# Int#
i# State# RealWorld
s#

-- | 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 :: MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe :: forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
n MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} = forall (m :: * -> *) a.
MonadIO m =>
MutableArray# RealWorld a -> Int -> m a
getIndexUnsafeWith MutableArray# RealWorld a
arrContents# (Int
n forall a. Num a => a -> a -> a
+ Int
arrStart)

-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
--
{-# INLINE getIndex #-}
getIndex :: MonadIO m => Int -> MutArray a -> m (Maybe a)
getIndex :: forall (m :: * -> *) a.
MonadIO m =>
Int -> MutArray a -> m (Maybe a)
getIndex Int
i arr :: MutArray a
arr@MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} =
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
< Int
arrLen
    then forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
arr
    else forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

-------------------------------------------------------------------------------
-- 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
    :: Int -- ^ from index
    -> Int -- ^ length of the slice
    -> MutArray a
    -> MutArray a
getSliceUnsafe :: forall a. Int -> Int -> MutArray a -> MutArray a
getSliceUnsafe Int
index Int
len arr :: MutArray a
arr@MutArray {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} =
    forall a. HasCallStack => Bool -> a -> a
assert (Int
index forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
index forall a. Num a => a -> a -> a
+ Int
len forall a. Ord a => a -> a -> Bool
<= Int
arrLen)
        forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrStart :: Int
arrStart = Int
arrStart forall a. Num a => a -> a -> a
+ Int
index, arrLen :: Int
arrLen = Int
len}

-- | /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
    :: Int -- ^ from index
    -> Int -- ^ length of the slice
    -> MutArray a
    -> MutArray a
getSlice :: forall a. Int -> Int -> MutArray a -> MutArray a
getSlice Int
index Int
len arr :: MutArray a
arr@MutArray{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} =
    if Int
index forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
index forall a. Num a => a -> a -> a
+ Int
len forall a. Ord a => a -> a -> Bool
<= Int
arrLen
    then MutArray a
arr {arrStart :: Int
arrStart = Int
arrStart forall a. Num a => a -> a -> a
+ Int
index, arrLen :: Int
arrLen = Int
len}
    else forall a. HasCallStack => [Char] -> a
error
             forall a b. (a -> b) -> a -> b
$ [Char]
"getSlice: invalid slice, index "
             forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
index forall a. [a] -> [a] -> [a]
++ [Char]
" length " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
len

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

-- XXX Maybe faster to create a list explicitly instead of mapM, if list fusion
-- does not work well.

-- | Convert an 'Array' into a list.
--
-- /Pre-release/
{-# INLINE toList #-}
toList :: MonadIO m => MutArray a -> m [a]
toList :: forall (m :: * -> *) a. MonadIO m => MutArray a -> m [a]
toList arr :: MutArray a
arr@MutArray{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
`getIndexUnsafe` MutArray a
arr) [Int
0 .. (Int
arrLen forall a. Num a => a -> a -> a
- Int
1)]

-- | Generates a stream from the elements of a @MutArray@.
--
-- >>> read = Stream.unfold MutArray.reader
--
{-# INLINE_NORMAL read #-}
read :: MonadIO m => MutArray a -> D.Stream m a
read :: forall (m :: * -> *) a. MonadIO m => MutArray a -> Stream m a
read arr :: MutArray a
arr@MutArray{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} =
    forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
D.mapM (forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
`getIndexUnsafe` MutArray a
arr) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> Stream m a
D.enumerateFromToIntegral Int
0 (Int
arrLen forall a. Num a => a -> a -> a
- Int
1)

-- Check equivalence with StreamK.fromStream . toStreamD and remove
{-# INLINE toStreamK #-}
toStreamK :: MonadIO m => MutArray a -> K.StreamK m a
toStreamK :: forall (m :: * -> *) a. MonadIO m => MutArray a -> StreamK m a
toStreamK arr :: MutArray a
arr@MutArray{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} = forall (m :: * -> *) b a.
Monad m =>
(b -> m (Maybe (a, b))) -> b -> StreamK m a
K.unfoldrM forall {m :: * -> *}. MonadIO m => Int -> m (Maybe (a, Int))
step Int
0

    where

    step :: Int -> m (Maybe (a, Int))
step Int
i
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
arrLen = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
        | Bool
otherwise = do
            a
x <- forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
arr
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (a
x, Int
i forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE_NORMAL readRev #-}
readRev :: MonadIO m => MutArray a -> D.Stream m a
readRev :: forall (m :: * -> *) a. MonadIO m => MutArray a -> Stream m a
readRev arr :: MutArray a
arr@MutArray{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. MutArray a -> Int
arrLen :: forall a. MutArray a -> Int
arrStart :: forall a. MutArray a -> Int
arrContents# :: forall a. MutArray a -> MutableArray# RealWorld a
..} =
    forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
D.mapM (forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
`getIndexUnsafe` MutArray a
arr)
        forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
D.enumerateFromThenToIntegral (Int
arrLen forall a. Num a => a -> a -> a
- Int
1) (Int
arrLen forall a. Num a => a -> a -> a
- Int
2) Int
0

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

-- XXX deduplicate this across unboxed array and this module?

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

-- | 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.
--
-- /Pre-release/
{-# INLINE_NORMAL writeNUnsafe #-}
writeNUnsafe :: MonadIO m => Int -> Fold m a (MutArray a)
writeNUnsafe :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (MutArray a)
writeNUnsafe Int
n = forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> (s -> m b) -> Fold m a b
Fold forall {f :: * -> *} {a} {b}.
MonadIO f =>
MutArray a -> a -> f (Step (MutArray a) b)
step forall {a} {b}. m (Step (MutArray a) b)
initial forall (m :: * -> *) a. Monad m => a -> m a
return forall (m :: * -> *) a. Monad m => a -> m a
return

    where

    initial :: m (Step (MutArray a) b)
initial = forall s b. s -> Step s b
FL.Partial forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new (forall a. Ord a => a -> a -> a
max Int
n Int
0)

    step :: MutArray a -> a -> f (Step (MutArray a) b)
step MutArray a
arr a
x = forall s b. s -> Step s b
FL.Partial forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> a -> m (MutArray a)
snocUnsafe MutArray a
arr a
x

-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
-- 'Array'.
--
-- >>> writeN n = Fold.take n (MutArray.writeNUnsafe n)
--
-- /Pre-release/
{-# INLINE_NORMAL writeN #-}
writeN :: MonadIO m => Int -> Fold m a (MutArray a)
writeN :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (MutArray a)
writeN Int
n = forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (MutArray a)
writeNUnsafe Int
n

-- >>> f n = MutArray.writeAppendWith (* 2) (MutArray.pinnedNew n)
-- >>> writeWith n = Fold.rmapM MutArray.rightSize (f n)
-- >>> writeWith n = Fold.rmapM MutArray.fromArrayStreamK (MutArray.writeChunks n)

-- | @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./
--
-- /Pre-release/
{-# INLINE_NORMAL writeWith #-}
writeWith :: MonadIO m => Int -> Fold m a (MutArray a)
-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (pinnedNew n)
writeWith :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (MutArray a)
writeWith Int
elemCount = forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> Fold m a b -> Fold m a c
FL.rmapM forall {a}. a -> m a
extract forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> a -> m (MutArray a)
step forall {a}. m (MutArray a)
initial

    where

    initial :: m (MutArray a)
initial = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
elemCount forall a. Ord a => a -> a -> Bool
< Int
0) forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => [Char] -> a
error [Char]
"writeWith: elemCount is negative"
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new Int
elemCount

    step :: MutArray a -> a -> m (MutArray a)
step arr :: MutArray a
arr@(MutArray MutableArray# RealWorld a
_ Int
start Int
end Int
bound) a
x
        | Int
end forall a. Eq a => a -> a -> Bool
== Int
bound = do
        let oldSize :: Int
oldSize = Int
end forall a. Num a => a -> a -> a
- Int
start
            newSize :: Int
newSize = forall a. Ord a => a -> a -> a
max (Int
oldSize forall a. Num a => a -> a -> a
* Int
2) Int
1
        MutArray a
arr1 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadIO m =>
Int -> MutArray a -> m (MutArray a)
realloc Int
newSize MutArray a
arr
        forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> a -> m (MutArray a)
snocUnsafe MutArray a
arr1 a
x
    step MutArray a
arr a
x = forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> a -> m (MutArray a)
snocUnsafe MutArray a
arr a
x

    -- extract = liftIO . rightSize
    extract :: a -> m a
extract = forall (m :: * -> *) a. Monad m => a -> m a
return

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

-- | Create a 'MutArray' from the first @n@ elements of a stream. The
-- array is allocated to size @n@, if the stream terminates before @n@
-- elements then the array may hold less than @n@ elements.
--
{-# INLINE fromStreamN #-}
fromStreamN :: MonadIO m => Int -> Stream m a -> m (MutArray a)
fromStreamN :: forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (MutArray a)
fromStreamN Int
n = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold (forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (MutArray a)
writeN Int
n)

{-# INLINE fromStream #-}
fromStream :: MonadIO m => Stream m a -> m (MutArray a)
fromStream :: forall (m :: * -> *) a. MonadIO m => Stream m a -> m (MutArray a)
fromStream = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold forall (m :: * -> *) a. MonadIO m => Fold m a (MutArray a)
write

{-# INLINABLE fromListN #-}
fromListN :: MonadIO m => Int -> [a] -> m (MutArray a)
fromListN :: forall (m :: * -> *) a. MonadIO m => Int -> [a] -> m (MutArray a)
fromListN Int
n [a]
xs = forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (MutArray a)
fromStreamN Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

{-# INLINABLE fromList #-}
fromList :: MonadIO m => [a] -> m (MutArray a)
fromList :: forall (m :: * -> *) a. MonadIO m => [a] -> m (MutArray a)
fromList [a]
xs = forall (m :: * -> *) a. MonadIO m => Stream m a -> m (MutArray a)
fromStream forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

{-# INLINABLE fromPureStream #-}
fromPureStream :: MonadIO m => Stream Identity a -> m (MutArray a)
fromPureStream :: forall (m :: * -> *) a.
MonadIO m =>
Stream Identity a -> m (MutArray a)
fromPureStream Stream Identity a
xs =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold forall (m :: * -> *) a. MonadIO m => Fold m a (MutArray a)
write forall a b. (a -> b) -> a -> b
$ forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> Stream m a -> Stream n a
D.morphInner (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Identity a -> a
runIdentity) Stream Identity a
xs

-------------------------------------------------------------------------------
-- Chunking
-------------------------------------------------------------------------------

data GroupState s a start end bound
    = GroupStart s
    | GroupBuffer s (MutableArray# RealWorld a) start end bound
    | GroupYield
          (MutableArray# RealWorld a)
          start
          end
          bound
          (GroupState s a start end bound)
    | GroupFinish

-- | @chunksOf n stream@ groups the input stream into a stream of
-- arrays of size n.
--
-- @chunksOf n = foldMany (MutArray.writeN n)@
--
-- /Pre-release/
{-# INLINE_NORMAL chunksOf #-}
chunksOf :: forall m a. MonadIO m
    => Int -> D.Stream m a -> D.Stream m (MutArray a)
-- XXX the idiomatic implementation leads to large regression in the D.reverse'
-- benchmark. It seems it has difficulty producing optimized code when
-- converting to StreamK. Investigate GHC optimizations.
-- chunksOf n = D.foldMany (writeN n)
chunksOf :: forall (m :: * -> *) a.
MonadIO m =>
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) =
    forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {a}.
State StreamK m a
-> GroupState s a Int Int Int
-> m (Step (GroupState s a Int Int Int) (MutArray a))
step' (forall s a start end bound. s -> GroupState s a start end bound
GroupStart s
state)

    where

    -- start is always 0
    -- end and len are always equal

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

    step' State StreamK m a
gst (GroupBuffer s
st MutableArray# RealWorld a
contents Int
start Int
end Int
bound) = do
        Step s a
r <- State StreamK m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State StreamK m a
gst) s
st
        case Step s a
r of
            D.Yield a
x s
s -> do
                forall (m :: * -> *) a.
MonadIO m =>
Int -> MutableArray# RealWorld a -> a -> m ()
putIndexUnsafeWith Int
end MutableArray# RealWorld a
contents a
x
                let end1 :: Int
end1 = Int
end forall a. Num a => a -> a -> a
+ Int
1
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
                    if Int
end1 forall a. Ord a => a -> a -> Bool
>= Int
bound
                    then forall s a. s -> Step s a
D.Skip
                            (forall s a start end bound.
MutableArray# RealWorld a
-> start
-> end
-> bound
-> GroupState s a start end bound
-> GroupState s a start end bound
GroupYield
                                MutableArray# RealWorld a
contents Int
start Int
end1 Int
bound (forall s a start end bound. s -> GroupState s a start end bound
GroupStart s
s))
                    else forall s a. s -> Step s a
D.Skip (forall s a start end bound.
s
-> MutableArray# RealWorld a
-> start
-> end
-> bound
-> GroupState s a start end bound
GroupBuffer s
s MutableArray# RealWorld a
contents Int
start Int
end1 Int
bound)
            D.Skip s
s ->
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s a start end bound.
s
-> MutableArray# RealWorld a
-> start
-> end
-> bound
-> GroupState s a start end bound
GroupBuffer s
s MutableArray# RealWorld a
contents Int
start Int
end Int
bound)
            Step s a
D.Stop ->
                forall (m :: * -> *) a. Monad m => a -> m a
return
                    forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s a start end bound.
MutableArray# RealWorld a
-> start
-> end
-> bound
-> GroupState s a start end bound
-> GroupState s a start end bound
GroupYield MutableArray# RealWorld a
contents Int
start Int
end Int
bound forall s a start end bound. GroupState s a start end bound
GroupFinish)

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

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

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

-- | Resumable unfold of an array.
--
{-# INLINE_NORMAL producerWith #-}
producerWith :: Monad m => (forall b. IO b -> m b) -> Producer m (MutArray a) a
producerWith :: forall (m :: * -> *) a.
Monad m =>
(forall b. IO b -> m b) -> Producer m (MutArray a) a
producerWith forall b. IO b -> m b
liftio = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer forall {a}. (MutArray a, Int) -> m (Step (MutArray a, Int) a)
step forall {m :: * -> *} {b} {a}. (Monad m, Num b) => a -> m (a, b)
inject forall {m :: * -> *} {a}.
Monad m =>
(MutArray a, Int) -> m (MutArray a)
extract

    where

    {-# INLINE inject #-}
    inject :: a -> m (a, b)
inject a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return (a
arr, b
0)

    {-# INLINE extract #-}
    extract :: (MutArray a, Int) -> m (MutArray a)
extract (MutArray a
arr, Int
i) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray a
arr {arrStart :: Int
arrStart = forall a. MutArray a -> Int
arrStart MutArray a
arr forall a. Num a => a -> a -> a
+ Int
i, arrLen :: Int
arrLen = forall a. MutArray a -> Int
arrLen MutArray a
arr forall a. Num a => a -> a -> a
- Int
i}

    {-# INLINE_LATE step #-}
    step :: (MutArray a, Int) -> m (Step (MutArray a, Int) a)
step (MutArray a
arr, Int
i)
        | forall a. HasCallStack => Bool -> a -> a
assert (forall a. MutArray a -> Int
arrLen MutArray a
arr forall a. Ord a => a -> a -> Bool
>= Int
0) (Int
i forall a. Eq a => a -> a -> Bool
== forall a. MutArray a -> Int
arrLen MutArray a
arr) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (MutArray a
arr, Int
i) = do
        a
x <- forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
arr
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (MutArray a
arr, Int
i forall a. Num a => a -> a -> a
+ Int
1)

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

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

--------------------------------------------------------------------------------
-- Appending arrays
--------------------------------------------------------------------------------

-- | 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 :: forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
src Int
srcStart MutArray a
dst Int
dstStart Int
len = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    assertM(Int
len forall a. Ord a => a -> a -> Bool
<= forall a. MutArray a -> Int
arrLen MutArray a
dst)
    assertM(Int
len forall a. Ord a => a -> a -> Bool
<= forall a. MutArray a -> Int
arrLen MutArray a
src)
    let !(I# Int#
srcStart#) = Int
srcStart forall a. Num a => a -> a -> a
+ forall a. MutArray a -> Int
arrStart MutArray a
src
        !(I# Int#
dstStart#) = Int
dstStart forall a. Num a => a -> a -> a
+ forall a. MutArray a -> Int
arrStart MutArray a
dst
        !(I# Int#
len#) = Int
len
    let arrS# :: MutableArray# RealWorld a
arrS# = forall a. MutArray a -> MutableArray# RealWorld a
arrContents# MutArray a
src
        arrD# :: MutableArray# RealWorld a
arrD# = forall a. MutArray a -> MutableArray# RealWorld a
arrContents# MutArray a
dst
    forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# -> (# forall d a.
MutableArray# d a
-> Int#
-> MutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableArray#
                    MutableArray# RealWorld a
arrS# Int#
srcStart# MutableArray# RealWorld a
arrD# Int#
dstStart# Int#
len# State# RealWorld
s#
                , () #)

{-# INLINE clone #-}
clone :: MonadIO m => MutArray a -> m (MutArray a)
clone :: forall (m :: * -> *) a. MonadIO m => MutArray a -> m (MutArray a)
clone MutArray a
src = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let len :: Int
len = forall a. MutArray a -> Int
arrLen MutArray a
src
    MutArray a
dst <- forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new Int
len
    forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
src Int
0 MutArray a
dst Int
0 Int
len
    forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
dst

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

{-# INLINE length #-}
length :: MutArray a -> Int
length :: forall a. MutArray a -> Int
length = forall a. MutArray a -> Int
arrLen

-------------------------------------------------------------------------------
-- 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, Ord a) => MutArray a -> MutArray a -> m Ordering
cmp :: forall (m :: * -> *) a.
(MonadIO m, Ord a) =>
MutArray a -> MutArray a -> m Ordering
cmp MutArray a
a1 MutArray a
a2 =
    case forall a. Ord a => a -> a -> Ordering
compare Int
lenA1 Int
lenA2 of
        Ordering
EQ -> forall {m :: * -> *}. MonadIO m => Int -> m Ordering
loop (Int
lenA1 forall a. Num a => a -> a -> a
- Int
1)
        Ordering
x -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
x

    where

    lenA1 :: Int
lenA1 = forall a. MutArray a -> Int
length MutArray a
a1
    lenA2 :: Int
lenA2 = forall a. MutArray a -> Int
length MutArray a
a2

    loop :: Int -> m Ordering
loop Int
i
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
EQ
        | Bool
otherwise = do
            a
v1 <- forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
a1
            a
v2 <- forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
a2
            case forall a. Ord a => a -> a -> Ordering
compare a
v1 a
v2 of
                Ordering
EQ -> Int -> m Ordering
loop (Int
i forall a. Num a => a -> a -> a
- Int
1)
                Ordering
x -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
x

{-# INLINE eq #-}
eq :: (MonadIO m, Eq a) => MutArray a -> MutArray a -> m Bool
eq :: forall (m :: * -> *) a.
(MonadIO m, Eq a) =>
MutArray a -> MutArray a -> m Bool
eq MutArray a
a1 MutArray a
a2 =
    if Int
lenA1 forall a. Eq a => a -> a -> Bool
== Int
lenA2
    then forall {m :: * -> *}. MonadIO m => Int -> m Bool
loop (Int
lenA1 forall a. Num a => a -> a -> a
- Int
1)
    else forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

    where

    lenA1 :: Int
lenA1 = forall a. MutArray a -> Int
length MutArray a
a1
    lenA2 :: Int
lenA2 = forall a. MutArray a -> Int
length MutArray a
a2

    loop :: Int -> m Bool
loop Int
i
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
        | Bool
otherwise = do
            a
v1 <- forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
a1
            a
v2 <- forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
i MutArray a
a2
            if a
v1 forall a. Eq a => a -> a -> Bool
== a
v2
            then Int -> m Bool
loop (Int
i forall a. Num a => a -> a -> a
- Int
1)
            else forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

{-# INLINE strip #-}
strip :: MonadIO m => (a -> Bool) -> MutArray a -> m (MutArray a)
strip :: forall (m :: * -> *) a.
MonadIO m =>
(a -> Bool) -> MutArray a -> m (MutArray a)
strip a -> Bool
p MutArray a
arr = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let lastIndex :: Int
lastIndex = forall a. MutArray a -> Int
length MutArray a
arr forall a. Num a => a -> a -> a
- Int
1
    Int
indexR <- forall {m :: * -> *}. MonadIO m => Int -> m Int
getIndexR Int
lastIndex -- last predicate failing index
    if Int
indexR forall a. Ord a => a -> a -> Bool
< Int
0
    then forall (m :: * -> *) a. MonadIO m => m (MutArray a)
nil
    else do
        Int
indexL <- forall {m :: * -> *}. MonadIO m => Int -> m Int
getIndexL Int
0 -- first predicate failing index
        if Int
indexL forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
indexR forall a. Eq a => a -> a -> Bool
== Int
lastIndex
        then forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr
        else
           let newLen :: Int
newLen = Int
indexR forall a. Num a => a -> a -> a
- Int
indexL forall a. Num a => a -> a -> a
+ Int
1
            in forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> MutArray a -> MutArray a
getSliceUnsafe Int
indexL Int
newLen MutArray a
arr

    where

    getIndexR :: Int -> m Int
getIndexR Int
idx
        | Int
idx forall a. Ord a => a -> a -> Bool
< Int
0 = forall (m :: * -> *) a. Monad m => a -> m a
return Int
idx
        | Bool
otherwise = do
            a
r <- forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
idx MutArray a
arr
            if a -> Bool
p a
r
            then Int -> m Int
getIndexR (Int
idx forall a. Num a => a -> a -> a
- Int
1)
            else forall (m :: * -> *) a. Monad m => a -> m a
return Int
idx

    getIndexL :: Int -> m Int
getIndexL Int
idx = do
        a
r <- forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
getIndexUnsafe Int
idx MutArray a
arr
        if a -> Bool
p a
r
        then Int -> m Int
getIndexL (Int
idx forall a. Num a => a -> a -> a
+ Int
1)
        else forall (m :: * -> *) a. Monad m => a -> m a
return Int
idx