{-# 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
--
module Streamly.Internal.Data.Array.Mut.Type
(
    -- * Type
    -- $arrayNotes
      Array (..)

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

    -- *** Uninitialized Arrays
    , newArray
    -- , newArrayWith

    -- *** From streams
    , writeNUnsafe
    , writeN

    -- , writeWith
    -- , write

    -- , writeRevN
    -- , writeRev

    -- ** From containers
    -- , fromListN
    -- , fromList
    -- , 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

    -- * 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
    -- , appendNUnsafe
    -- , appendN
    -- , appendWith
    -- , append

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

    -- ** To streams
    , read
    -- , readRev

    -- ** To containers
    , toStreamD
    -- , toStreamDRev
    , toStreamK
    -- , toStreamKRev
    , toList

    -- experimental
    , producer

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

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

    -- * 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
    -- , arraysOf
    -- , 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
    -- , putSlice
    -- , appendSlice
    -- , appendSliceFrom
    )
where

#include "inline.hs"

import Control.Exception (assert)
import Control.Monad.IO.Class (MonadIO(..))
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 qualified Streamly.Internal.Data.Fold.Type as FL
import qualified Streamly.Internal.Data.Producer as Producer
import qualified Streamly.Internal.Data.Stream.StreamD as D
import qualified Streamly.Internal.Data.Stream.StreamK as K

import Prelude hiding (read)

-- $setup
-- >>> :m
-- >>> import qualified Streamly.Internal.Data.Array.Mut.Type as Array
-- >>> import qualified Streamly.Internal.Data.Stream.IsStream as Stream
-- >>> import qualified Streamly.Internal.Data.Stream.StreamD as StreamD
-- >>> import qualified Streamly.Internal.Data.Fold as Fold
-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold

-------------------------------------------------------------------------------
-- Array Data Type
-------------------------------------------------------------------------------

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

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

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

        , Array 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 :: a
bottomElement =
    [Char] -> a
forall a. HasCallStack => [Char] -> a
error
        ([Char] -> a) -> [Char] -> a
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.Array.Mut.Type.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.
-- | @newArray count@ allocates an empty array that can hold 'count' items.
--
-- /Pre-release/
{-# INLINE newArray #-}
newArray :: forall m a. MonadIO m => Int -> m (Array a)
newArray :: Int -> m (Array a)
newArray n :: Int
n@(I# Int#
n#) =
    IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ (State# RealWorld -> (# State# RealWorld, Array a #))
-> IO (Array a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        ((State# RealWorld -> (# State# RealWorld, Array a #))
 -> IO (Array a))
-> (State# RealWorld -> (# State# RealWorld, Array a #))
-> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              case Int#
-> a
-> State# RealWorld
-> (# State# RealWorld, MutableArray# RealWorld a #)
forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
newArray# Int#
n# a
forall a. a
bottomElement State# RealWorld
s# of
                  (# State# RealWorld
s1#, MutableArray# RealWorld a
arr# #) ->
                      let ma :: Array a
ma = MutableArray# RealWorld a -> Int -> Int -> Int -> Array a
forall a. MutableArray# RealWorld a -> Int -> Int -> Int -> Array a
Array MutableArray# RealWorld a
arr# Int
0 Int
0 Int
n
                       in (# State# RealWorld
s1#, Array a
ma #)

-------------------------------------------------------------------------------
-- 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 => Array a -> Int -> a -> m ()
putIndexUnsafe :: Array a -> Int -> a -> m ()
putIndexUnsafe Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} Int
i a
x =
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ (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# ->
              case Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
arrStart of
                  I# Int#
n# ->
                      let s1# :: State# RealWorld
s1# = MutableArray# RealWorld a
-> Int# -> a -> State# RealWorld -> State# RealWorld
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#, () #)

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 arr ix val = Array.modifyIndex arr ix (const (val, ()))
--
-- /Pre-release/
{-# INLINE putIndex #-}
putIndex :: MonadIO m => Array a -> Int -> a -> m ()
putIndex :: Array a -> Int -> a -> m ()
putIndex arr :: Array a
arr@Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} Int
i a
x =
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrLen
    then Array a -> Int -> a -> m ()
forall (m :: * -> *) a. MonadIO m => Array a -> Int -> a -> m ()
putIndexUnsafe Array a
arr Int
i a
x
    else [Char] -> Int -> m ()
forall a. [Char] -> Int -> a
invalidIndex [Char]
"putIndex" Int
i

-- | 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 => Array a -> Int -> (a -> (a, b)) -> m b
modifyIndexUnsafe :: Array a -> Int -> (a -> (a, b)) -> m b
modifyIndexUnsafe Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} Int
i a -> (a, b)
f = do
    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
$ (State# RealWorld -> (# State# RealWorld, b #)) -> IO b
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        ((State# RealWorld -> (# State# RealWorld, b #)) -> IO b)
-> (State# RealWorld -> (# State# RealWorld, b #)) -> IO b
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              case Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
arrStart of
                  I# Int#
n# ->
                      case MutableArray# RealWorld a
-> Int# -> State# RealWorld -> (# State# RealWorld, a #)
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# = MutableArray# RealWorld a
-> Int# -> a -> State# RealWorld -> State# RealWorld
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 => Array a -> Int -> (a -> (a, b)) -> m b
modifyIndex :: Array a -> Int -> (a -> (a, b)) -> m b
modifyIndex arr :: Array a
arr@Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} Int
i a -> (a, b)
f = do
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrLen
    then Array a -> Int -> (a -> (a, b)) -> m b
forall (m :: * -> *) a b.
MonadIO m =>
Array a -> Int -> (a -> (a, b)) -> m b
modifyIndexUnsafe Array a
arr Int
i a -> (a, b)
f
    else [Char] -> Int -> m b
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 -> Array a -> m (Array a)
realloc :: Int -> Array a -> m (Array a)
realloc Int
n Array a
arr = do
    Array a
arr1 <- Int -> m (Array a)
forall (m :: * -> *) a. MonadIO m => Int -> m (Array a)
newArray Int
n
    let !newLen :: Int
newLen@(I# Int#
newLen#) = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
n (Array a -> Int
forall a. Array a -> Int
arrLen Array a
arr)
        !(I# Int#
arrS#) = Array a -> Int
forall a. Array a -> Int
arrStart Array a
arr
        !(I# Int#
arr1S#) = Array a -> Int
forall a. Array a -> Int
arrStart Array a
arr1
        arrC# :: MutableArray# RealWorld a
arrC# = Array a -> MutableArray# RealWorld a
forall a. Array a -> MutableArray# RealWorld a
arrContents# Array a
arr
        arr1C# :: MutableArray# RealWorld a
arr1C# = Array a -> MutableArray# RealWorld a
forall a. Array a -> MutableArray# RealWorld a
arrContents# Array a
arr1
    IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ (State# RealWorld -> (# State# RealWorld, Array a #))
-> IO (Array a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        ((State# RealWorld -> (# State# RealWorld, Array a #))
 -> IO (Array a))
-> (State# RealWorld -> (# State# RealWorld, Array a #))
-> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              let s1# :: State# RealWorld
s1# = MutableArray# RealWorld a
-> Int#
-> MutableArray# RealWorld a
-> Int#
-> Int#
-> State# RealWorld
-> State# RealWorld
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#, Array a
arr1 {arrLen :: Int
arrLen = Int
newLen, arrTrueLen :: Int
arrTrueLen = Int
n} #)

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

    where

    badSize :: [Char]
badSize = [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
        [ [Char]
label
        , [Char]
": new array size is less than required size "
        , Int -> [Char]
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 => Array a -> a -> m (Array a)
snocUnsafe :: Array a -> a -> m (Array a)
snocUnsafe arr :: Array a
arr@Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} a
a = do
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
arrLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrTrueLen) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Array a -> Int -> a -> m ()
forall (m :: * -> *) a. MonadIO m => Array a -> Int -> a -> m ()
putIndexUnsafe Array a
arr Int
arrLen a
a
    Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> m (Array a)) -> Array a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Array a
arr {arrLen :: Int
arrLen = Int
arrLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1}

-- NOINLINE to move it out of the way and not pollute the instruction cache.
{-# NOINLINE snocWithRealloc #-}
snocWithRealloc :: MonadIO m => (Int -> Int) -> Array a -> a -> m (Array a)
snocWithRealloc :: (Int -> Int) -> Array a -> a -> m (Array a)
snocWithRealloc Int -> Int
sizer Array a
arr a
x = do
    Array a
arr1 <- [Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
forall (m :: * -> *) a.
MonadIO m =>
[Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
reallocWith [Char]
"snocWithRealloc" Int -> Int
sizer Int
1 Array a
arr
    Array a -> a -> m (Array a)
forall (m :: * -> *) a. MonadIO m => Array a -> a -> m (Array a)
snocUnsafe Array 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) -> Array a -> a -> m (Array a)
snocWith :: (Int -> Int) -> Array a -> a -> m (Array a)
snocWith Int -> Int
sizer arr :: Array a
arr@Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} a
x = do
    if Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
arrLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrTrueLen
    then Array a -> a -> m (Array a)
forall (m :: * -> *) a. MonadIO m => Array a -> a -> m (Array a)
snocUnsafe Array a
arr a
x
    else (Int -> Int) -> Array a -> a -> m (Array a)
forall (m :: * -> *) a.
MonadIO m =>
(Int -> Int) -> Array a -> a -> m (Array a)
snocWithRealloc Int -> Int
sizer Array 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 = Array.snocWith (* 2)
--
-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
--
-- /Pre-release/
{-# INLINE snoc #-}
snoc :: MonadIO m => Array a -> a -> m (Array a)
snoc :: Array a -> a -> m (Array a)
snoc = (Int -> Int) -> Array a -> a -> m (Array a)
forall (m :: * -> *) a.
MonadIO m =>
(Int -> Int) -> Array a -> a -> m (Array a)
snocWith (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2)

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

-- | 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 => Array a -> Int -> m a
getIndexUnsafe :: Array a -> Int -> m a
getIndexUnsafe Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} Int
n =
    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
$ (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
        ((State# RealWorld -> (# State# RealWorld, a #)) -> IO a)
-> (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s# ->
              let !(I# Int#
i#) = Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n
               in MutableArray# RealWorld a
-> Int# -> State# RealWorld -> (# State# RealWorld, a #)
forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# MutableArray# RealWorld a
arrContents# Int#
i# State# RealWorld
s#

-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
--
{-# INLINE getIndex #-}
getIndex :: MonadIO m => Array a -> Int -> m a
getIndex :: Array a -> Int -> m a
getIndex arr :: Array a
arr@Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} Int
i =
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arrLen
    then Array a -> Int -> m a
forall (m :: * -> *) a. MonadIO m => Array a -> Int -> m a
getIndexUnsafe Array a
arr Int
i
    else [Char] -> Int -> m a
forall a. [Char] -> Int -> a
invalidIndex [Char]
"getIndex" Int
i

-------------------------------------------------------------------------------
-- 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
    -> Array a
    -> Array a
getSliceUnsafe :: Int -> Int -> Array a -> Array a
getSliceUnsafe Int
index Int
len arr :: Array a
arr@Array {Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} =
    Bool -> Array a -> Array 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
index Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrLen)
        (Array a -> Array a) -> Array a -> Array a
forall a b. (a -> b) -> a -> b
$ Array a
arr {arrStart :: Int
arrStart = Int
arrStart Int -> Int -> Int
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
    -> Array a
    -> Array a
getSlice :: Int -> Int -> Array a -> Array a
getSlice Int
index Int
len arr :: Array a
arr@Array{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} =
    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
index Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arrLen
    then Array a
arr {arrStart :: Int
arrStart = Int
arrStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
index, arrLen :: Int
arrLen = Int
len}
    else [Char] -> Array a
forall a. HasCallStack => [Char] -> a
error
             ([Char] -> Array a) -> [Char] -> Array 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

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

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

-- | Use the 'read' unfold instead.
--
-- @toStreamD = D.unfold read@
--
-- We can try this if the unfold has any performance issues.
{-# INLINE_NORMAL toStreamD #-}
toStreamD :: MonadIO m => Array a -> D.Stream m a
toStreamD :: Array a -> Stream m a
toStreamD arr :: Array a
arr@Array{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} =
    (Int -> m a) -> Stream m Int -> Stream m a
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
D.mapM (Array a -> Int -> m a
forall (m :: * -> *) a. MonadIO m => Array a -> Int -> m a
getIndexUnsafe Array a
arr) (Stream m Int -> Stream m a) -> Stream m Int -> Stream m a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Stream m Int
forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> Stream m a
D.enumerateFromToIntegral Int
0 (Int
arrLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE toStreamK #-}
toStreamK :: MonadIO m => Array a -> K.Stream m a
toStreamK :: Array a -> Stream m a
toStreamK arr :: Array a
arr@Array{Int
MutableArray# RealWorld a
arrTrueLen :: Int
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrTrueLen :: forall a. Array a -> Int
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} = (Int -> m (Maybe (a, Int))) -> Int -> Stream m a
forall (m :: * -> *) b a.
Monad m =>
(b -> m (Maybe (a, b))) -> b -> Stream m a
K.unfoldrM Int -> m (Maybe (a, Int))
forall (m :: * -> *). MonadIO m => Int -> m (Maybe (a, Int))
step Int
0

    where

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


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

-- | 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 (Array a)
writeNUnsafe :: Int -> Fold m a (Array a)
writeNUnsafe Int
n = (Array a -> a -> m (Step (Array a) (Array a)))
-> m (Step (Array a) (Array a))
-> (Array a -> m (Array a))
-> Fold m a (Array a)
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold Array a -> a -> m (Step (Array a) (Array a))
forall (f :: * -> *) a b.
MonadIO f =>
Array a -> a -> f (Step (Array a) b)
step m (Step (Array a) (Array a))
forall a b. m (Step (Array a) b)
initial Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return

    where

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

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

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

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

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

    where

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

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

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

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