{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module      : Streamly.Internal.Data.Array.Foreign.Mut.Type
-- Copyright   : (c) 2020 Composewell Technologies
-- License     : BSD3-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Mutable arrays and file system files are quite similar, they can grow and
-- their content is mutable. Therefore, both have similar APIs as well. We
-- strive to keep the API consistent for both. Ideally, you should be able to
-- replace one with another with little changes to the code.

module Streamly.Internal.Data.Array.Foreign.Mut.Type
    (
    -- * Type
    -- $arrayNotes
      Array (..)
    , ArrayContents
    , arrayToFptrContents
    , fptrToArrayContents
    , nilArrayContents
    , touch

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

    -- *** Uninitialized Arrays
    , newArray
    , newArrayAligned
    , newArrayAlignedUnmanaged
    , newArrayWith

    -- *** Initialized Arrays
    , withNewArrayUnsafe

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

    , writeWith
    , write

    -- , writeRevN
    -- , writeRev

    -- ** From containers
    , fromForeignPtrUnsafe
    , 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
    , unsafeSwapIndices

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

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

    -- ** Appending streams
    , appendNUnsafe
    , appendN
    , appendWith
    , append

    -- * Eliminating and Reading

    -- ** To streams
    , ReadUState(..)
    , read
    , readRev

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

    -- experimental
    , producer

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

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

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

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

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

    -- * Folding
    , foldl'
    , foldr
    , cmp

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

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

    -- ** Construct from streams
    , 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

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

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

#ifdef USE_C_MALLOC
#define USE_FOREIGN_PTR
#endif

import Control.Exception (assert)
import Control.DeepSeq (NFData(..))
import Control.Monad (when, void)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Bits (shiftR, (.|.), (.&.))
#if __GLASGOW_HASKELL__ < 808
import Data.Semigroup (Semigroup(..))
#endif
import Data.Word (Word8)
import Foreign.C.Types (CSize(..), CInt(..))
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
#ifndef USE_FOREIGN_PTR
import Foreign.Marshal.Alloc (mallocBytes)
#endif
import Foreign.Ptr (plusPtr, minusPtr, castPtr, nullPtr)
import Foreign.Storable (Storable(..))
import GHC.Base
    ( touch#, IO(..), byteArrayContents#
    , Int(..), newAlignedPinnedByteArray#
    )
#ifndef USE_FOREIGN_PTR
import GHC.Base (RealWorld, MutableByteArray#)
#endif
#if __GLASGOW_HASKELL__ < 802
#define noinline
#else
import GHC.Base (noinline)
#endif
import GHC.Exts (unsafeCoerce#)
import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..))
#ifdef USE_C_MALLOC
import GHC.ForeignPtr (mallocForeignPtrAlignedBytes)
#endif
import GHC.Ptr (Ptr(..))

import Streamly.Internal.BaseCompat
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Producer.Type (Producer (..))
import Streamly.Internal.Data.Stream.Serial (SerialT(..))
import Streamly.Internal.Data.SVar.Type (adaptState, defState)
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.System.IO (arrayPayloadSize, defaultChunkSize)
import System.IO.Unsafe (unsafePerformIO)

#ifdef DEVBUILD
import qualified Data.Foldable as F
#endif
import qualified Streamly.Internal.Data.Fold.Type as FL
import qualified Streamly.Internal.Data.Producer as Producer
import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
#ifdef USE_FOREIGN_PTR
import qualified Streamly.Internal.Foreign.Malloc as Malloc
#endif

import Prelude hiding
    (length, foldr, read, unlines, splitAt, reverse, truncate)

-- $setup
-- >>> :m
-- >>> import qualified Streamly.Internal.Data.Array.Foreign.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

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

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

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

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

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

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

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

-------------------------------------------------------------------------------
-- Array Contents
-------------------------------------------------------------------------------

-- We support using ForeignPtrContents or MutableByteArray.

#ifdef USE_FOREIGN_PTR
newtype ArrayContents = ArrayContents ForeignPtrContents
#define UNPACKIF
#else
-- XXX can use UnliftedNewtypes
data ArrayContents = ArrayContents !(MutableByteArray# RealWorld)
#define UNPACKIF {-# UNPACK #-}
#endif

{-# INLINE touch #-}
touch :: ArrayContents -> IO ()
touch :: ArrayContents -> IO ()
touch (ArrayContents MutableByteArray# RealWorld
contents) =
    forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case touch# :: forall a. a -> State# RealWorld -> State# RealWorld
touch# MutableByteArray# RealWorld
contents State# RealWorld
s of State# RealWorld
s' -> (# State# RealWorld
s', () #)

fptrToArrayContents :: ForeignPtrContents -> ArrayContents
arrayToFptrContents :: ArrayContents -> ForeignPtrContents
#ifdef USE_FOREIGN_PTR
fptrToArrayContents = ArrayContents
arrayToFptrContents (ArrayContents contents) = contents
#else
fptrToArrayContents :: ForeignPtrContents -> ArrayContents
fptrToArrayContents (PlainPtr MutableByteArray# RealWorld
mbarr) = MutableByteArray# RealWorld -> ArrayContents
ArrayContents MutableByteArray# RealWorld
mbarr
fptrToArrayContents ForeignPtrContents
_ = forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Unsupported foreign ptr"
arrayToFptrContents :: ArrayContents -> ForeignPtrContents
arrayToFptrContents (ArrayContents MutableByteArray# RealWorld
contents) = MutableByteArray# RealWorld -> ForeignPtrContents
PlainPtr MutableByteArray# RealWorld
contents
#endif

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

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

-- XXX Rename the fields to better names.

-- | An unboxed, pinned mutable array. An array is created with a given length
-- and capacity. Length is the number of valid elements in the array.  Capacity
-- is the maximum number of elements that the array can be expanded to without
-- having to reallocate the memory.
--
-- The elements in the array can be mutated in-place without changing the
-- reference (constructor). However, the length of the array cannot be mutated
-- in-place.  A new array reference is generated when the length changes.  When
-- the length is increased (upto the maximum reserved capacity of the array),
-- the array is not reallocated and the new reference uses the same underlying
-- memory as the old one.
--
-- Several routines in this module allow the programmer to control the capacity
-- of the array. The programmer can control the trade-off between memory usage
-- and performance impact due to reallocations when growing or shrinking the
-- array.
--
data Array a =
#ifdef DEVBUILD
    Storable a =>
#endif
    Array
    { forall a. Array a -> ArrayContents
arrContents :: UNPACKIF !ArrayContents
    , forall a. Array a -> Ptr a
arrStart :: {-# UNPACK #-} !(Ptr a)      -- ^ first address
    , forall a. Array a -> Ptr a
aEnd   :: {-# UNPACK #-} !(Ptr a)        -- ^ first unused address
    , forall a. Array a -> Ptr a
aBound :: {-# UNPACK #-} !(Ptr a)        -- ^ first address beyond allocated memory
    }

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

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

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

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

newAlignedArrayContents :: Int -> Int -> IO (ArrayContents, Ptr a)
#ifdef USE_C_MALLOC
newAlignedArrayContents size align = do
    (ForeignPtr addr contents) <- mallocForeignPtrAlignedBytes size align
    return (ArrayContents contents, Ptr addr)
#else
newAlignedArrayContents :: forall a. Int -> Int -> IO (ArrayContents, Ptr a)
newAlignedArrayContents Int
size Int
_align | Int
size forall a. Ord a => a -> a -> Bool
< Int
0 =
  forall a. [Char] -> a
errorWithoutStackTrace [Char]
"newAlignedArrayContents: size must be >= 0"
newAlignedArrayContents (I# Int#
size) (I# Int#
align) = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
    case forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# Int#
size Int#
align State# RealWorld
s of
        (# State# RealWorld
s', MutableByteArray# RealWorld
mbarr# #) ->
           let p :: Ptr a
p = forall a. Addr# -> Ptr a
Ptr (ByteArray# -> Addr#
byteArrayContents# (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
mbarr#))
#ifdef USE_FOREIGN_PTR
               c = ArrayContents (PlainPtr mbarr#)
#else
               c :: ArrayContents
c = MutableByteArray# RealWorld -> ArrayContents
ArrayContents MutableByteArray# RealWorld
mbarr#
#endif
            in (# State# RealWorld
s', (ArrayContents
c, forall {a}. Ptr a
p) #)
#endif

{-# NOINLINE nilArrayContents #-}
nilArrayContents :: ArrayContents
nilArrayContents :: ArrayContents
nilArrayContents =
    forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> IO (ArrayContents, Ptr a)
newAlignedArrayContents Int
0 Int
0

nil ::
#ifdef DEVBUILD
    Storable a =>
#endif
    Array a
nil :: forall a. Array a
nil = forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
nilArrayContents forall {a}. Ptr a
nullPtr forall {a}. Ptr a
nullPtr forall {a}. Ptr a
nullPtr

-- | @fromForeignPtrUnsafe foreignPtr end bound@ creates an 'Array' that starts
-- at the memory pointed by the @foreignPtr@, @end@ is the first unused
-- address, and @bound@ is the first address beyond the allocated memory.
--
-- Unsafe: Make sure that foreignPtr <= end <= bound and (end - start) is an
-- integral multiple of the element size. Only PlainPtr type ForeignPtr is
-- supported.
--
-- /Pre-release/
--
{-# INLINE fromForeignPtrUnsafe #-}
fromForeignPtrUnsafe ::
#ifdef DEVBUILD
    Storable a =>
#endif
    ForeignPtr a -> Ptr a -> Ptr a -> Array a
fromForeignPtrUnsafe :: forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
fromForeignPtrUnsafe (ForeignPtr Addr#
start ForeignPtrContents
_) Ptr a
_ Ptr a
_
    | forall a. Addr# -> Ptr a
Ptr Addr#
start forall a. Eq a => a -> a -> Bool
== forall {a}. Ptr a
nullPtr = forall a. Array a
nil
fromForeignPtrUnsafe fp :: ForeignPtr a
fp@(ForeignPtr Addr#
start ForeignPtrContents
contents) Ptr a
end Ptr a
bound =
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
fp forall a. Ord a => a -> a -> Bool
<= Ptr a
end Bool -> Bool -> Bool
&& Ptr a
end forall a. Ord a => a -> a -> Bool
<= Ptr a
bound)
           (forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array (ForeignPtrContents -> ArrayContents
fptrToArrayContents ForeignPtrContents
contents) (forall a. Addr# -> Ptr a
Ptr Addr#
start) Ptr a
end Ptr a
bound)

-- | Like 'newArrayWith' but using an allocator that allocates unmanaged pinned
-- memory. The memory will never be freed by GHC.  This could be useful in
-- allocate-once global data structures. Use carefully as incorrect use can
-- lead to memory leak.
--
-- /Internal/
{-# INLINE newArrayAlignedUnmanaged #-}
newArrayAlignedUnmanaged :: forall m a. (MonadIO m, Storable a) =>
    Int -> Int -> m (Array a)
#ifdef USE_FOREIGN_PTR
newArrayAlignedUnmanaged = do
    newArrayWith mallocForeignPtrAlignedUnmanagedBytes

    where

    mallocForeignPtrAlignedUnmanagedBytes size align = do
        ForeignPtr addr contents <-
            liftIO $ Malloc.mallocForeignPtrAlignedUnmanagedBytes size align
        return (ArrayContents contents, Ptr addr)
#else
newArrayAlignedUnmanaged :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAlignedUnmanaged Int
_align Int
count = do
    let size :: Int
size = forall a. Ord a => a -> a -> a
max (Int
count forall a. Num a => a -> a -> a
* SIZE_OF(a)) 0
    Ptr a
p <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Int -> IO (Ptr a)
mallocBytes Int
size
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Array
        { arrContents :: ArrayContents
arrContents = ArrayContents
nilArrayContents
        , arrStart :: Ptr a
arrStart = Ptr a
p
        , aEnd :: Ptr a
aEnd = Ptr a
p
        , aBound :: Ptr a
aBound = Ptr a
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size
        }
#endif

-- | Like 'newArrayWith' but using an allocator that aligns the memory to the
-- alignment dictated by the 'Storable' instance of the type.
--
-- /Internal/
{-# INLINE newArrayAligned #-}
newArrayAligned :: (MonadIO m, Storable a) => Int -> Int -> m (Array a)
newArrayAligned :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAligned = forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> Int -> m (ArrayContents, Ptr a))
-> Int -> Int -> m (Array a)
newArrayWith (\Int
s Int
a -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> IO (ArrayContents, Ptr a)
newAlignedArrayContents Int
s Int
a)

-- XXX can unaligned allocation be more efficient when alignment is not needed?
--
-- | Allocates an empty array that can hold 'count' items.  The memory of the
-- array is uninitialized and the allocation is aligned as per the 'Storable'
-- instance of the type.
--
-- /Pre-release/
{-# INLINE newArray #-}
newArray :: forall m a. (MonadIO m, Storable a) => Int -> m (Array a)
newArray :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray = forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAligned (forall a. Storable a => a -> Int
alignment (forall a. (?callStack::CallStack) => a
undefined :: a))

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

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

-- | Write the given element to the given index of the array. Does not check if
-- the index is out of bounds of the array.
--
-- /Pre-release/
{-# INLINE putIndexUnsafe #-}
putIndexUnsafe :: forall m a. (MonadIO m, Storable a)
    => Int -> a -> Array a -> m ()
putIndexUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> a -> Array a -> m ()
putIndexUnsafe Int
i a
x arr :: Array a
arr@(Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..}) =
    forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let elemPtr :: Ptr b
elemPtr = PTR_INDEX(ptr,i,a)
        forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& PTR_VALID(elemPtr,aEndInt
,a)) (return ())
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
poke forall {a}. Ptr a
elemPtr a
x

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

{-# INLINE putIndexPtr #-}
putIndexPtr :: forall m a. (MonadIO m, Storable a) =>
    Ptr a -> Ptr a -> Int -> a -> m ()
putIndexPtr :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> a -> m ()
putIndexPtr Ptr a
start Ptr a
end Int
i a
x = do
    let elemPtr :: Ptr b
elemPtr = PTR_INDEX(start,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& PTR_VALID(elemPtr,forall a. Ord a => a -> a -> a
end,Int
a)
    then forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
poke forall {a}. Ptr a
elemPtr a
x
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"putIndexPtr" 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 ix (const (val, ())) arr
-- >>> f = Array.putIndices
-- >>> putIndex ix val arr = Stream.fold (f arr) (Stream.fromPure (ix, val))
--
-- /Pre-release/
{-# INLINE putIndex #-}
putIndex :: (MonadIO m, Storable a) => Int -> a -> Array a -> m ()
putIndex :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> a -> Array a -> m ()
putIndex Int
i a
x Array a
arr =
    forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr
        forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> a -> m ()
putIndexPtr Ptr a
p (forall a. Array a -> Ptr a
aEnd Array a
arr) Int
i a
x

-- | Write an input stream of (index, value) pairs to an array. Throws an
-- error if any index is out of bounds.
--
-- /Pre-release/
{-# INLINE putIndices #-}
putIndices :: forall m a. (MonadIO m, Storable a)
    => Array a -> Fold m (Int, a) ()
putIndices :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Fold m (Int, a) ()
putIndices Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall s a (m :: * -> *) b.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
FL.mkFoldM forall {f :: * -> *} {b}.
MonadIO f =>
() -> (Int, a) -> f (Step () b)
step forall {b}. m (Step () b)
initial forall {m :: * -> *}. MonadIO m => () -> m ()
extract

    where

    initial :: m (Step () b)
initial = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
FL.Partial ()

    step :: () -> (Int, a) -> f (Step () b)
step () (Int
i, 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 => IO a -> m a
liftIO (forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> a -> m ()
putIndexPtr Ptr a
arrStart Ptr a
aEnd Int
i a
x)

    extract :: () -> m ()
extract () = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
arrContents

-- | Modify a given index of an array using a modifier function.
--
-- /Pre-release/
modifyIndexUnsafe :: forall m a b. (MonadIO m, Storable a) =>
    Int -> (a -> (a, b)) -> Array a -> m b
modifyIndexUnsafe :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Int -> (a -> (a, b)) -> Array a -> m b
modifyIndexUnsafe Int
i a -> (a, b)
f arr :: Array a
arr@(Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..}) = do
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let elemPtr :: Ptr b
elemPtr = PTR_INDEX(ptr,i,a)
        forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& PTR_NEXT(elemPtr,a) <= aEnd) (return ())
        a
r <- forall a. Storable a => Ptr a -> IO a
peek forall {a}. Ptr a
elemPtr
        let (a
x, b
res) = a -> (a, b)
f a
r
        forall a. Storable a => Ptr a -> a -> IO ()
poke forall {a}. Ptr a
elemPtr a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return b
res

{-# INLINE modifyIndexPtr #-}
modifyIndexPtr :: forall m a b. (MonadIO m, Storable a) =>
    Int -> (a -> (a, b)) -> Ptr a -> Ptr a -> m b
modifyIndexPtr :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Int -> (a -> (a, b)) -> Ptr a -> Ptr a -> m b
modifyIndexPtr Int
i a -> (a, b)
f Ptr a
start Ptr a
end = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let elemPtr :: Ptr b
elemPtr = PTR_INDEX(start,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& PTR_VALID(elemPtr,forall a. Ord a => a -> a -> a
end,Int
a)
    then do
        a
r <- forall a. Storable a => Ptr a -> IO a
peek forall {a}. Ptr a
elemPtr
        let (a
x, b
res) = a -> (a, b)
f a
r
        forall a. Storable a => Ptr a -> a -> IO ()
poke forall {a}. Ptr a
elemPtr a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return b
res
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"modifyIndex" Int
i

-- | Modify a given index of an array using a modifier function.
--
-- /Pre-release/
modifyIndex :: forall m a b. (MonadIO m, Storable a) =>
    Int -> (a -> (a, b)) -> Array a -> m b
modifyIndex :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Int -> (a -> (a, b)) -> Array a -> m b
modifyIndex Int
i a -> (a, b)
f arr :: Array a
arr@(Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..}) = do
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Int -> (a -> (a, b)) -> Ptr a -> Ptr a -> m b
modifyIndexPtr Int
i a -> (a, b)
f Ptr a
ptr Ptr a
aEnd

-- | Modify the array indices generated by the supplied stream.
--
-- /Pre-release/
{-# INLINE modifyIndices #-}
modifyIndices :: forall m a. (MonadIO m, Storable a)
    => (a -> a) -> Array a -> Fold m Int ()
modifyIndices :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(a -> a) -> Array a -> Fold m Int ()
modifyIndices a -> a
f Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold forall {f :: * -> *} {b}. MonadIO f => () -> Int -> f (Step () b)
step forall {b}. m (Step () b)
initial forall {m :: * -> *}. MonadIO m => () -> m ()
extract

    where

    initial :: m (Step () b)
initial = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
FL.Partial ()

    step :: () -> Int -> f (Step () b)
step () Int
i =
        let f1 :: a -> (a, ())
f1 a
x = (a -> a
f a
x, ())
         in 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 => IO a -> m a
liftIO (forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Int -> (a -> (a, b)) -> Ptr a -> Ptr a -> m b
modifyIndexPtr Int
i a -> (a, ())
f1 Ptr a
arrStart Ptr a
aEnd)

    extract :: () -> m ()
extract () = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
arrContents

-- | Modify each element of an array using the supplied modifier function.
--
-- /Pre-release/
modify :: forall m a. (MonadIO m, Storable a)
    => (a -> a) -> Array a -> m ()
modify :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(a -> a) -> Array a -> m ()
modify a -> a
f arr :: Array a
arr@Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr Ptr a -> IO ()
go

    where

    go :: Ptr a -> IO ()
go Ptr a
ptr =
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (PTR_VALID(ptr,aEndInt
,a)) $ do
            r <- peek ptr
            poke ptr (f r)
            go (PTR_NEXT(ptr,a))

{-# INLINE swapPtrs #-}
swapPtrs :: Storable a => Ptr a -> Ptr a -> IO ()
swapPtrs :: forall a. Storable a => Ptr a -> Ptr a -> IO ()
swapPtrs Ptr a
ptr1 Ptr a
ptr2 = do
    a
r1 <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
ptr1
    a
r2 <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
ptr2
    forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
ptr1 a
r2
    forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
ptr2 a
r1

-- | Swap the elements at two indices without validating the indices.
--
-- /Unsafe/: This could result in memory corruption if indices are not valid.
--
-- /Pre-release/
{-# INLINE unsafeSwapIndices #-}
unsafeSwapIndices :: forall m a. (MonadIO m, Storable a)
    => Int -> Int -> Array a -> m ()
unsafeSwapIndices :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Array a -> m ()
unsafeSwapIndices Int
i1 Int
i2 Array a
arr = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let ptr1 :: Ptr b
ptr1 = PTR_INDEX(ptr,i1,a)
            ptr2 :: Ptr b
ptr2 = PTR_INDEX(ptr,i2,a)
        forall a. Storable a => Ptr a -> Ptr a -> IO ()
swapPtrs forall {a}. Ptr a
ptr1 (forall {a}. Ptr a
ptr2 :: Ptr a)

-- | Swap the elements at two indices.
--
-- /Pre-release/
swapIndices :: forall m a. (MonadIO m, Storable a)
    => Int -> Int -> Array a -> m ()
swapIndices :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Array a -> m ()
swapIndices Int
i1 Int
i2 Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let ptr1 :: Ptr b
ptr1 = PTR_INDEX(arrStart,i1,a)
            ptr2 :: Ptr b
ptr2 = PTR_INDEX(arrStart,i2,a)
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i1 forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| PTR_INVALID(ptr1,aEnd,a))
            forall a b. (a -> b) -> a -> b
$ forall a. [Char] -> Int -> a
invalidIndex [Char]
"swapIndices" Int
i1
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i2 forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| PTR_INVALID(ptr2,aEnd,a))
            forall a b. (a -> b) -> a -> b
$ forall a. [Char] -> Int -> a
invalidIndex [Char]
"swapIndices" Int
i2
        forall a. Storable a => Ptr a -> Ptr a -> IO ()
swapPtrs forall {a}. Ptr a
ptr1 (forall {a}. Ptr a
ptr2 :: Ptr a)

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

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

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

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

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

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

    where

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

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

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

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

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

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

-- XXX See if resizing can be implemented by reading the old array as a stream
-- and then using writeN to the new array.
--
{-# NOINLINE reallocAligned #-}
reallocAligned :: Int -> Int -> Int -> Array a -> IO (Array a)
reallocAligned :: forall a. Int -> Int -> Int -> Array a -> IO (Array a)
reallocAligned Int
elemSize Int
alignSize Int
newCapacity Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = do
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
aEnd forall a. Ord a => a -> a -> Bool
<= Ptr a
aBound) (forall (m :: * -> *) a. Monad m => a -> m a
return ())

    -- Allocate new array
    let newCapMax :: Int
newCapMax = Int -> Int
roundUpLargeArray Int
newCapacity
    (ArrayContents
contents, Ptr a
pNew) <- forall a. Int -> Int -> IO (ArrayContents, Ptr a)
newAlignedArrayContents Int
newCapMax Int
alignSize

    -- Copy old data
    let oldStart :: Ptr a
oldStart = Ptr a
arrStart
        oldSize :: Int
oldSize = Ptr a
aEnd forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
oldStart
        newCap :: Int
newCap = Int -> Int -> Int
roundDownTo Int
elemSize Int
newCapMax
        newLen :: Int
newLen = forall a. Ord a => a -> a -> a
min Int
oldSize Int
newCap
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
oldSize forall a. Integral a => a -> a -> a
`mod` Int
elemSize forall a. Eq a => a -> a -> Bool
== Int
0) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
newLen forall a. Ord a => a -> a -> Bool
>= Int
0) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
newLen forall a. Integral a => a -> a -> a
`mod` Int
elemSize forall a. Eq a => a -> a -> Bool
== Int
0) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (forall a b. Ptr a -> Ptr b
castPtr Ptr a
pNew) (forall a b. Ptr a -> Ptr b
castPtr Ptr a
oldStart) Int
newLen
    ArrayContents -> IO ()
touch ArrayContents
arrContents

    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Array
        { arrStart :: Ptr a
arrStart = Ptr a
pNew
        , arrContents :: ArrayContents
arrContents = ArrayContents
contents
        , aEnd :: Ptr a
aEnd   = Ptr a
pNew forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
newLen
        , aBound :: Ptr a
aBound = Ptr a
pNew forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
newCap
        }

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

-- | @reallocWith label capSizer minIncrement array@. The label is used
-- in error messages and the capSizer is used to determine the capacity of the
-- new array in bytes given the current byte length of the array.
reallocWith :: forall m a. (MonadIO m , Storable a) =>
       String
    -> (Int -> Int)
    -> Int
    -> Array a
    -> m (Array a)
reallocWith :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
[Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
reallocWith [Char]
label Int -> Int
capSizer Int
minIncr Array a
arr = do
    let oldSize :: Int
oldSize = forall a. Array a -> Ptr a
aEnd Array a
arr forall a b. Ptr a -> Ptr b -> Int
`minusPtr` forall a. Array a -> Ptr a
arrStart Array a
arr
        newCap :: Int
newCap = Int -> Int
capSizer Int
oldSize
        newSize :: Int
newSize = Int
oldSize forall a. Num a => a -> a -> a
+ Int
minIncr
        safeCap :: Int
safeCap = forall a. Ord a => a -> a -> a
max Int
newCap Int
newSize
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
newCap forall a. Ord a => a -> a -> Bool
>= Int
newSize Bool -> Bool -> Bool
|| forall a. (?callStack::CallStack) => [Char] -> a
error (forall a. Show a => a -> [Char]
badSize Int
newSize)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m (Array a)
realloc Int
safeCap Array a
arr

    where

    badSize :: a -> [Char]
badSize a
newSize = 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 a
newSize
        , [Char]
". Please check the sizing function passed."
        ]

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

-- | Like 'resize' but if the capacity is more than 'largeObjectThreshold' then
-- it is rounded up to the closest power of 2.
--
-- /Pre-release/
{-# INLINE resizeExp #-}
resizeExp :: forall m a. (MonadIO m, Storable a) =>
    Int -> Array a -> m (Array a)
resizeExp :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m (Array a)
resizeExp Int
n arr :: Array a
arr@Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = do
    let req :: Int
req = Int -> Int
roundUpLargeArray (SIZE_OF(a) * n)
        req1 :: Int
req1 =
            if Int
req forall a. Ord a => a -> a -> Bool
> Int
largeObjectThreshold
            then Int -> Int
roundUpToPower2 Int
req
            else Int
req
        len :: Int
len = Ptr a
aEnd forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
arrStart
    if Int
req1 forall a. Ord a => a -> a -> Bool
< Int
len
    then forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr
    else forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m (Array a)
realloc Int
req1 Array a
arr

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

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

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

-- | Snoc using a 'Ptr'. Low level reusable function.
--
-- /Internal/
{-# INLINE snocNewEnd #-}
snocNewEnd :: (MonadIO m, Storable a) => Ptr a -> Array a -> a -> m (Array a)
snocNewEnd :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Array a -> a -> m (Array a)
snocNewEnd Ptr a
newEnd arr :: Array a
arr@Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} a
x = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
newEnd forall a. Ord a => a -> a -> Bool
<= Ptr a
aBound) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
aEnd a
x
    ArrayContents -> IO ()
touch ArrayContents
arrContents
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Array a
arr {aEnd :: Ptr a
aEnd = Ptr a
newEnd}

-- | Really really unsafe, appends the element into the first array, may
-- cause silent data corruption or if you are lucky a segfault if the first
-- array does not have enough space to append the element.
--
-- /Internal/
{-# INLINE snocUnsafe #-}
snocUnsafe :: forall m a. (MonadIO m, Storable a) =>
    Array a -> a -> m (Array a)
snocUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> a -> m (Array a)
snocUnsafe arr :: Array a
arr@Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Array a -> a -> m (Array a)
snocNewEnd (PTR_NEXT(aEnd,a)) arr

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

-- NOINLINE to move it out of the way and not pollute the instruction cache.
{-# NOINLINE snocWithRealloc #-}
snocWithRealloc :: forall m a. (MonadIO m, Storable a) =>
       (Int -> Int)
    -> Array a
    -> a
    -> m (Array a)
snocWithRealloc :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> Int) -> Array a -> a -> m (Array a)
snocWithRealloc Int -> Int
sizer Array a
arr a
x = do
    Array a
arr1 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
[Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
reallocWith [Char]
"snocWith" Int -> Int
sizer (SIZE_OF(a)) arr
    forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
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 oldSizeBytes@ function, where
-- @oldSizeBytes@ is the original size of the array in bytes.
--
-- If the new array size is more than 'largeObjectThreshold' we automatically
-- round it up to 'blockSize'.
--
-- Note that the returned array may be a mutated version of the original array.
--
-- /Pre-release/
{-# INLINE snocWith #-}
snocWith :: forall m a. (MonadIO m, Storable a) =>
       (Int -> Int)
    -> Array a
    -> a
    -> m (Array a)
snocWith :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> Int) -> Array a -> a -> m (Array a)
snocWith Int -> Int
allocSize Array a
arr a
x = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let newEnd :: Ptr b
newEnd = PTR_NEXT(aEnd arr,a)
    if forall {a}. Ptr a
newEnd forall a. Ord a => a -> a -> Bool
<= forall a. Array a -> Ptr a
aBound Array a
arr
    then forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Array a -> a -> m (Array a)
snocNewEnd forall {a}. Ptr a
newEnd Array a
arr a
x
    else forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> Int) -> Array a -> a -> m (Array a)
snocWithRealloc Int -> Int
allocSize Array a
arr a
x

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

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

    where

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

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

-- | Return the element at the specified index without checking the bounds.
--
-- Unsafe because it does not check the bounds of the array.
{-# INLINE_NORMAL getIndexUnsafe #-}
getIndexUnsafe :: forall m a. (MonadIO m, Storable a) => Int -> Array a -> m a
getIndexUnsafe :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m a
getIndexUnsafe Int
i arr :: Array a
arr@(Array {Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..}) =
    forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let elemPtr :: Ptr b
elemPtr = PTR_INDEX(ptr,i,a)
        forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& PTR_VALID(elemPtr,aEndInt
,a)) (return ())
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek forall {a}. Ptr a
elemPtr

{-# INLINE getIndexPtr #-}
getIndexPtr :: forall m a. (MonadIO m, Storable a) =>
    Ptr a -> Ptr a -> Int -> m a
getIndexPtr :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtr Ptr a
start Ptr a
end Int
i = do
    let elemPtr :: Ptr b
elemPtr = PTR_INDEX(start,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& PTR_VALID(elemPtr,forall a. Ord a => a -> a -> a
end,Int
a)
    then forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek forall {a}. Ptr a
elemPtr
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"getIndexPtr" Int
i

-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
--
{-# INLINE getIndex #-}
getIndex :: (MonadIO m, Storable a) => Int -> Array a -> m a
getIndex :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m a
getIndex Int
i Array a
arr =
    forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr
        forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtr Ptr a
p (forall a. Array a -> Ptr a
aEnd Array a
arr) Int
i

{-# INLINE getIndexPtrRev #-}
getIndexPtrRev :: forall m a. (MonadIO m, Storable a) =>
    Ptr a -> Ptr a -> Int -> m a
getIndexPtrRev :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtrRev Ptr a
start Ptr a
end Int
i = do
    let elemPtr :: Ptr b
elemPtr = PTR_RINDEX(end,i,a)
    if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& forall {a}. Ptr a
elemPtr forall a. Ord a => a -> a -> Bool
>= Ptr a
start
    then forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek forall {a}. Ptr a
elemPtr
    else forall a. [Char] -> Int -> a
invalidIndex [Char]
"getIndexPtrRev" Int
i

-- | /O(1)/ Lookup the element at the given index from the end of the array.
-- Index starts from 0.
--
-- Slightly faster than computing the forward index and using getIndex.
--
{-# INLINE getIndexRev #-}
getIndexRev :: (MonadIO m, Storable a) => Int -> Array a -> m a
getIndexRev :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m a
getIndexRev Int
i Array a
arr =
    forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array a
arr
        forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtrRev Ptr a
p (forall a. Array a -> Ptr a
aEnd Array a
arr) Int
i

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

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

    where

    inject :: Array a -> m (GetIndicesState ArrayContents (Ptr a) (Ptr a) s)
inject (Array ArrayContents
contents Ptr a
start (Ptr Addr#
end) Ptr a
_) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState ArrayContents
contents Ptr a
start (forall a. Addr# -> Ptr a
Ptr Addr#
end) s
sti

    {-# INLINE_LATE step #-}
    step :: GetIndicesState ArrayContents (Ptr a) (Ptr a) s
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
step (GetIndicesState ArrayContents
contents Ptr a
start Ptr a
end s
st) = do
        Step s Int
r <- State Stream m Int -> s -> m (Step s Int)
stepi forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s Int
r of
            D.Yield Int
i s
s -> do
                a
x <- forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtr Ptr a
start Ptr a
end Int
i
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState ArrayContents
contents Ptr a
start Ptr a
end s
s)
            D.Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall contents start end st.
contents
-> start -> end -> st -> GetIndicesState contents start end st
GetIndicesState ArrayContents
contents Ptr a
start Ptr a
end s
s)
            Step s Int
D.Stop -> do
                forall b. IO b -> m b
liftio forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
                forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop

{-# INLINE getIndices #-}
getIndices :: (MonadIO m, Storable a) => SerialT m Int -> Unfold m (Array a) a
getIndices :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
SerialT m Int -> Unfold m (Array a) a
getIndices (SerialT Stream m Int
stream) = forall (m :: * -> *) a.
(Monad m, Storable a) =>
(forall b. IO b -> m b) -> Stream m Int -> Unfold m (Array a) a
getIndicesD forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => Stream m a -> Stream m a
D.fromStreamK Stream m Int
stream

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

-- XXX We can also get immutable slices.

-- | /O(1)/ Slice an array in constant time.
--
-- Unsafe: The bounds of the slice are not checked.
--
-- /Unsafe/
--
-- /Pre-release/
{-# INLINE getSliceUnsafe #-}
getSliceUnsafe :: forall a. Storable a
    => Int -- ^ from index
    -> Int -- ^ length of the slice
    -> Array a
    -> Array a
getSliceUnsafe :: forall a. Storable a => Int -> Int -> Array a -> Array a
getSliceUnsafe Int
index Int
len (Array ArrayContents
contents Ptr a
start Ptr a
e Ptr a
_) =
    let fp1 :: Ptr b
fp1 = PTR_INDEX(start,index,a)
        end :: Ptr b
end = forall {a}. Ptr a
fp1 forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
len forall a. Num a => a -> a -> a
* SIZE_OF(a))
     in forall a. (?callStack::CallStack) => Bool -> a -> a
assert
            (Int
index forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& forall {a}. Ptr a
end forall a. Ord a => a -> a -> Bool
<= Ptr a
e)
            (forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents forall {a}. Ptr a
fp1 forall {a}. Ptr a
end forall {a}. Ptr a
end)

-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
-- extends out of the array bounds.
--
-- /Pre-release/
{-# INLINE getSlice #-}
getSlice :: forall a. Storable a =>
       Int -- ^ from index
    -> Int -- ^ length of the slice
    -> Array a
    -> Array a
getSlice :: forall a. Storable a => Int -> Int -> Array a -> Array a
getSlice Int
index Int
len (Array ArrayContents
contents Ptr a
start Ptr a
e Ptr a
_) =
    let fp1 :: Ptr b
fp1 = PTR_INDEX(start,index,a)
        end :: Ptr b
end = forall {a}. Ptr a
fp1 forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
len forall a. Num a => a -> a -> a
* SIZE_OF(a))
     in if Int
index forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& forall {a}. Ptr a
end forall a. Ord a => a -> a -> Bool
<= Ptr a
e
        then forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents forall {a}. Ptr a
fp1 forall {a}. Ptr a
end forall {a}. Ptr a
end
        else forall a. (?callStack::CallStack) => [Char] -> a
error
                forall a b. (a -> b) -> a -> b
$ [Char]
"getSlice: invalid slice, index "
                forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
index forall a. [a] -> [a] -> [a]
++ [Char]
" length " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
len

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

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

-- | You may not need to reverse an array because you can consume it in reverse
-- using 'readRev'. To reverse large arrays you can read in reverse and write
-- to another array. However, in-place reverse can be useful to take adavantage
-- of cache locality and when you do not want to allocate additional memory.
--
-- /Pre-release/
{-# INLINE reverse #-}
reverse :: forall m a. (MonadIO m, Storable a) => Array a -> m ()
reverse :: forall (m :: * -> *) a. (MonadIO m, Storable a) => Array a -> m ()
reverse Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let l :: Ptr a
l = Ptr a
arrStart
        h :: Ptr b
h = PTR_PREV(aEnd,a)
     in forall a. Storable a => Ptr a -> Ptr a -> IO ()
swap Ptr a
l forall {a}. Ptr a
h

    where

    swap :: Ptr b -> Ptr b -> IO ()
swap Ptr b
l Ptr b
h = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ptr b
l forall a. Ord a => a -> a -> Bool
< Ptr b
h) forall a b. (a -> b) -> a -> b
$ do
            forall a. Storable a => Ptr a -> Ptr a -> IO ()
swapPtrs Ptr b
l Ptr b
h
            Ptr b -> Ptr b -> IO ()
swap (PTR_NEXT(l,a)) (PTR_PREV(h,a))

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

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

    where

    -- Invariant low < high on entry, and on return as well
    moveHigh :: Ptr b -> Ptr a -> IO (Maybe (Ptr a, a))
moveHigh Ptr b
low Ptr a
high = do
        a
h <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
high
        if a -> Bool
f a
h
        then
            -- Correctly classified, continue the loop
            let high1 :: Ptr b
high1 = PTR_PREV(high,a)
             in if Ptr b
low forall a. Eq a => a -> a -> Bool
== forall {a}. Ptr a
high1
                then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
                else Ptr b -> Ptr a -> IO (Maybe (Ptr a, a))
moveHigh Ptr b
low forall {a}. Ptr a
high1
        else forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just (Ptr a
high, a
h)) -- incorrectly classified

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

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

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

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

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

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

-- | /O(1)/ Get the byte length of the array.
--
-- @since 0.7.0
{-# INLINE byteLength #-}
byteLength :: Array a -> Int
byteLength :: forall a. Array a -> Int
byteLength Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} =
    let len :: Int
len = Ptr a
aEnd forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
arrStart
    in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
len forall a. Ord a => a -> a -> Bool
>= Int
0) Int
len

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

-- | Get the total capacity of an array. An array may have space reserved
-- beyond the current used length of the array.
--
-- /Pre-release/
{-# INLINE byteCapacity #-}
byteCapacity :: Array a -> Int
byteCapacity :: forall a. Array a -> Int
byteCapacity Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} =
    let len :: Int
len = Ptr a
aBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
arrStart
    in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
len forall a. Ord a => a -> a -> Bool
>= Int
0) Int
len

-- | The remaining capacity in the array for appending more elements without
-- reallocation.
--
-- /Pre-release/
{-# INLINE bytesFree #-}
bytesFree :: Array a -> Int
bytesFree :: forall a. Array a -> Int
bytesFree Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} =
    let n :: Int
n = Ptr a
aBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
aEnd
    in forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
>= Int
0) Int
n

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

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

-- | @arraysOf n stream@ groups the input stream into a stream of
-- arrays of size n.
--
-- @arraysOf n = StreamD.foldMany (Array.writeN n)@
--
-- /Pre-release/
{-# INLINE_NORMAL arraysOf #-}
arraysOf :: forall m a. (MonadIO m, Storable a)
    => Int -> D.Stream m a -> D.Stream m (Array 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.
-- arraysOf n = D.foldMany (writeN n)
arraysOf :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> Stream m (Array a)
arraysOf Int
n (D.Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {a}.
State Stream m a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
step' (forall s contents start end bound.
s -> GroupState s contents start end bound
GroupStart s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
step' State Stream m a
_ (GroupStart s
st) = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n forall a. Ord a => a -> a -> Bool
<= Int
0) forall a b. (a -> b) -> a -> b
$
            -- XXX we can pass the module string from the higher level API
            forall a. (?callStack::CallStack) => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Streamly.Internal.Data.Array.Foreign.Mut.Type.arraysOf: "
                    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"
        Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray Int
n
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
st ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound)

    step' State Stream m a
gst (GroupBuffer s
st ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            D.Yield a
x s
s -> do
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ArrayContents -> IO ()
touch ArrayContents
contents
                let end1 :: Ptr b
end1 = PTR_NEXT(end,a)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
                    if forall {a}. Ptr a
end1 forall a. Ord a => a -> a -> Bool
>= Ptr a
bound
                    then forall s a. s -> Step s a
D.Skip
                            (forall s contents start end bound.
contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
-> GroupState s contents start end bound
GroupYield
                                ArrayContents
contents Ptr a
start forall {a}. Ptr a
end1 Ptr a
bound (forall s contents start end bound.
s -> GroupState s contents start end bound
GroupStart s
s))
                    else forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
s ArrayContents
contents Ptr a
start forall {a}. Ptr a
end1 Ptr a
bound)
            D.Skip s
s ->
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
s
-> contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
GroupBuffer s
s ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound)
            Step s a
D.Stop ->
                forall (m :: * -> *) a. Monad m => a -> m a
return
                    forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
D.Skip (forall s contents start end bound.
contents
-> start
-> end
-> bound
-> GroupState s contents start end bound
-> GroupState s contents start end bound
GroupYield ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound forall s contents start end bound.
GroupState s contents start end bound
GroupFinish)

    step' State Stream m a
_ (GroupYield ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound) GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
next

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

-- XXX buffer to a list instead?
-- | Buffer the stream into arrays in memory.
{-# INLINE arrayStreamKFromStreamD #-}
arrayStreamKFromStreamD :: forall m a. (MonadIO m, Storable a) =>
    D.Stream m a -> m (K.Stream m (Array a))
arrayStreamKFromStreamD :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m a -> m (Stream m (Array a))
arrayStreamKFromStreamD =
    let n :: Int
n = forall a. Storable a => a -> Int -> Int
allocBytesToElemCount (forall a. (?callStack::CallStack) => a
undefined :: a) Int
defaultChunkSize
     in forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr forall a (m :: * -> *). a -> Stream m a -> Stream m a
K.cons forall (m :: * -> *) a. Stream m a
K.nil forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> Stream m (Array a)
arraysOf Int
n

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

data FlattenState s contents a =
      OuterLoop s
    | InnerLoop s contents !(Ptr a) !(Ptr a)

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

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> FlattenState s ArrayContents a
-> m (Step (FlattenState s ArrayContents a) a)
step' State Stream m a
gst (OuterLoop s
st) = do
        Step s (Array a)
r <- State Stream m (Array a) -> s -> m (Step s (Array a))
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s (Array a)
r of
            D.Yield Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} s
s ->
                forall s a. s -> Step s a
D.Skip (forall s contents a.
s -> contents -> Ptr a -> Ptr a -> FlattenState s contents a
InnerLoop s
s ArrayContents
arrContents Ptr a
arrStart Ptr a
aEnd)
            D.Skip s
s -> forall s a. s -> Step s a
D.Skip (forall s contents a. s -> FlattenState s contents a
OuterLoop s
s)
            Step s (Array a)
D.Stop -> forall s a. Step s a
D.Stop

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

    step' State Stream m a
_ (InnerLoop s
st ArrayContents
contents Ptr a
p Ptr a
end) = do
        a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
                    a
r <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
                    ArrayContents -> IO ()
touch ArrayContents
contents
                    forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall s contents a.
s -> contents -> Ptr a -> Ptr a -> FlattenState s contents a
InnerLoop s
st ArrayContents
contents (PTR_NEXT(p,a)) end)

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

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> FlattenState s ArrayContents a
-> m (Step (FlattenState s ArrayContents a) a)
step' State Stream m a
gst (OuterLoop s
st) = do
        Step s (Array a)
r <- State Stream m (Array a) -> s -> m (Step s (Array a))
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s (Array a)
r of
            D.Yield Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} s
s ->
                let p :: Ptr b
p = PTR_PREV(aEnd,a)
                 in forall s a. s -> Step s a
D.Skip (forall s contents a.
s -> contents -> Ptr a -> Ptr a -> FlattenState s contents a
InnerLoop s
s ArrayContents
arrContents forall {a}. Ptr a
p Ptr a
arrStart)
            D.Skip s
s -> forall s a. s -> Step s a
D.Skip (forall s contents a. s -> FlattenState s contents a
OuterLoop s
s)
            Step s (Array a)
D.Stop -> forall s a. Step s a
D.Stop

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

    step' State Stream m a
_ (InnerLoop s
st ArrayContents
contents Ptr a
p Ptr a
start) = do
        a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
                    a
r <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
                    ArrayContents -> IO ()
touch ArrayContents
contents
                    forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        let cur :: Ptr b
cur = PTR_PREV(p,a)
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall s contents a.
s -> contents -> Ptr a -> Ptr a -> FlattenState s contents a
InnerLoop s
st ArrayContents
contents forall {a}. Ptr a
cur Ptr a
start)

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

data ReadUState a = ReadUState
    UNPACKIF !ArrayContents  -- contents
    !(Ptr a)           -- end address
    !(Ptr a)           -- current address

toReadUState :: Array a -> ReadUState a
toReadUState :: forall a. Array a -> ReadUState a
toReadUState (Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
_) = forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
end Ptr a
start

-- | Resumable unfold of an array.
--
{-# INLINE_NORMAL producer #-}
producer :: forall m a. (MonadIO m, Storable a) => Producer m (Array a) a
producer :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Producer m (Array a) a
producer = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer forall {m :: * -> *} {a}.
(MonadIO m, Storable a) =>
ReadUState a -> m (Step (ReadUState a) a)
step (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Array a -> ReadUState a
toReadUState) forall {m :: * -> *} {a}. Monad m => ReadUState a -> m (Array a)
extract
    where

    {-# INLINE_LATE step #-}
    step :: ReadUState a -> m (Step (ReadUState a) a)
step (ReadUState ArrayContents
contents Ptr a
end Ptr a
cur)
        | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
cur forall a. Ord a => a -> a -> Bool
<= Ptr a
end) (Ptr a
cur forall a. Eq a => a -> a -> Bool
== Ptr a
end) = do
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
            forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (ReadUState ArrayContents
contents Ptr a
end Ptr a
cur) = do
            !a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek Ptr a
cur
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
end (PTR_NEXT(cur,a)))

    extract :: ReadUState a -> m (Array a)
extract (ReadUState ArrayContents
contents Ptr a
end Ptr a
cur) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
cur Ptr a
end Ptr a
end

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

-- | Unfold an array into a stream in reverse order.
--
-- /Pre-release/
{-# INLINE_NORMAL readRev #-}
readRev :: forall m a. (MonadIO m, Storable a) => Unfold m (Array a) a
readRev :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Unfold m (Array a) a
readRev = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {m :: * -> *} {a}.
(MonadIO m, Storable a) =>
ReadUState a -> m (Step (ReadUState a) a)
step forall {m :: * -> *} {a}. Monad m => Array a -> m (ReadUState a)
inject
    where

    inject :: Array a -> m (ReadUState a)
inject (Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
_) =
        let p :: Ptr b
p = PTR_PREV(end,a)
         in forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
start forall {a}. Ptr a
p

    {-# INLINE_LATE step #-}
    step :: ReadUState a -> m (Step (ReadUState a) a)
step (ReadUState ArrayContents
contents Ptr a
start Ptr a
p) | Ptr a
p forall a. Ord a => a -> a -> Bool
< Ptr a
start = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (ReadUState ArrayContents
contents Ptr a
start Ptr a
p) = do
        a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
x (forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
start (PTR_PREV(p,a)))

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

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

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

-- XXX Monadic foldr/build fusion?
-- Reference: https://www.researchgate.net/publication/220676509_Monadic_augment_and_generalised_short_cut_fusion
-- | Convert an 'Array' into a list.
--
-- @since 0.7.0
{-# INLINE toList #-}
toList :: forall m a. (MonadIO m, Storable a) => Array a -> m [a]
toList :: forall (m :: * -> *) a. (MonadIO m, Storable a) => Array a -> m [a]
toList Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Ptr a -> IO [a]
go Ptr a
arrStart
    where

    go :: Ptr a -> IO [a]
go Ptr a
p | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
p forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = forall (m :: * -> *) a. Monad m => a -> m a
return []
    go Ptr a
p = do
        a
x <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        ArrayContents -> IO ()
touch ArrayContents
arrContents
        (:) a
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> IO [a]
go (PTR_NEXT(p,a))

-- | Use the 'read' unfold instead.
--
-- @toStreamD = D.unfold read@
--
-- We can try this if the unfold has any performance issues.
{-# INLINE_NORMAL toStreamD #-}
toStreamD :: forall m a. (MonadIO m, Storable a) => Array a -> D.Stream m a
toStreamD :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Stream m a
toStreamD Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {p} {b}.
MonadIO m =>
p -> Ptr a -> m (Step (Ptr b) a)
step Ptr a
arrStart

    where

    {-# INLINE_LATE step #-}
    step :: p -> Ptr a -> m (Step (Ptr b) a)
step p
_ Ptr a
p | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
p forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step p
_ Ptr a
p = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        a
r <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        ArrayContents -> IO ()
touch ArrayContents
arrContents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
r (PTR_NEXT(p,a))

{-# INLINE toStreamK #-}
toStreamK :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a
toStreamK :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Stream m a
toStreamK Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall {m :: * -> *}. MonadIO m => Ptr a -> Stream m a
go Ptr a
arrStart

    where

    go :: Ptr a -> Stream m a
go Ptr a
p | forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
p forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = forall (m :: * -> *) a. Stream m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = do
              a
r <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
              ArrayContents -> IO ()
touch ArrayContents
arrContents
              forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO a
elemM forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
`K.consM` Ptr a -> Stream m a
go (PTR_NEXT(p,a))

-- | Use the 'readRev' unfold instead.
--
-- @toStreamDRev = D.unfold readRev@
--
-- We can try this if the unfold has any perf issues.
{-# INLINE_NORMAL toStreamDRev #-}
toStreamDRev :: forall m a. (MonadIO m, Storable a) => Array a -> D.Stream m a
toStreamDRev :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Stream m a
toStreamDRev Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} =
    let p :: Ptr b
p = PTR_PREV(aEnd,a)
    in forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {p} {b}.
MonadIO m =>
p -> Ptr a -> m (Step (Ptr b) a)
step forall {a}. Ptr a
p

    where

    {-# INLINE_LATE step #-}
    step :: p -> Ptr a -> m (Step (Ptr b) a)
step p
_ Ptr a
p | Ptr a
p forall a. Ord a => a -> a -> Bool
< Ptr a
arrStart = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step p
_ Ptr a
p = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        a
r <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        ArrayContents -> IO ()
touch ArrayContents
arrContents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield a
r (PTR_PREV(p,a))

{-# INLINE toStreamKRev #-}
toStreamKRev :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a
toStreamKRev :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Stream m a
toStreamKRev Array {Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} =
    let p :: Ptr b
p = PTR_PREV(aEnd,a)
    in forall {m :: * -> *}. MonadIO m => Ptr a -> Stream m a
go forall {a}. Ptr a
p

    where

    go :: Ptr a -> Stream m a
go Ptr a
p | Ptr a
p forall a. Ord a => a -> a -> Bool
< Ptr a
arrStart = forall (m :: * -> *) a. Stream m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = do
              a
r <- forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
              ArrayContents -> IO ()
touch ArrayContents
arrContents
              forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO a
elemM forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
`K.consM` Ptr a -> Stream m a
go (PTR_PREV(p,a))

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

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

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

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

data ArrayUnsafe a = ArrayUnsafe
    UNPACKIF !ArrayContents  -- contents
    {-# UNPACK #-} !(Ptr a)  -- start address
    {-# UNPACK #-} !(Ptr a)  -- first unused address

toArrayUnsafe :: Array a -> ArrayUnsafe a
toArrayUnsafe :: forall a. Array a -> ArrayUnsafe a
toArrayUnsafe (Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
_) =
    forall a. ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe a
ArrayUnsafe ArrayContents
contents Ptr a
start Ptr a
end

fromArrayUnsafe ::
#ifdef DEVBUILD
    Storable a =>
#endif
    ArrayUnsafe a -> Array a
fromArrayUnsafe :: forall a. ArrayUnsafe a -> Array a
fromArrayUnsafe (ArrayUnsafe ArrayContents
contents Ptr a
start Ptr a
end) =
         forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
end

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

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

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

    where

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

    step :: ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step (ArrayUnsafe ArrayContents
contents Ptr a
start Ptr a
end) a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ArrayContents -> IO ()
touch ArrayContents
contents
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe a
ArrayUnsafe ArrayContents
contents Ptr a
start (PTR_NEXT(end,a))

-- | Append @n@ elements to an existing array. Any free space left in the array
-- after appending @n@ elements is lost.
--
-- >>> appendN initial n = Fold.take n (Array.appendNUnsafe initial n)
--
-- /Pre-release/
{-# INLINE_NORMAL appendN #-}
appendN :: forall m a. (MonadIO m, Storable a) =>
    m (Array a) -> Int -> Fold m a (Array a)
appendN :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
m (Array a) -> Int -> Fold m a (Array a)
appendN m (Array a)
initial Int
n = forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n (forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
m (Array a) -> Int -> Fold m a (Array a)
appendNUnsafe m (Array a)
initial Int
n)

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

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

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

    where

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

    step :: ArrayUnsafe a -> a -> m (Step (ArrayUnsafe a) b)
step (ArrayUnsafe ArrayContents
contents Ptr a
start Ptr a
end) a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ArrayContents -> IO ()
touch ArrayContents
contents
        forall (m :: * -> *) a. Monad m => a -> m a
return
          forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
FL.Partial
          forall a b. (a -> b) -> a -> b
$ forall a. ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe a
ArrayUnsafe ArrayContents
contents Ptr a
start (PTR_NEXT(end,a))

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

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

-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
-- 'Array'.
--
-- >>> writeN = Array.writeNWith Array.newArray
-- >>> writeN n = Fold.take n (Array.writeNUnsafe n)
-- >>> writeN n = Array.appendN (Array.newArray n) n
--
-- @since 0.7.0
{-# INLINE_NORMAL writeN #-}
writeN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
writeN :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Fold m a (Array a)
writeN = forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWith forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray

-- | @writeNAligned align n@ folds a maximum of @n@ elements from the input
-- stream to an 'Array' aligned to the given size.
--
-- >>> writeNAligned align = Array.writeNWith (Array.newArrayAligned align)
-- >>> writeNAligned align n = Array.appendN (Array.newArrayAligned align n) n
--
-- /Pre-release/
--
{-# INLINE_NORMAL writeNAligned #-}
writeNAligned :: forall m a. (MonadIO m, Storable a)
    => Int -> Int -> Fold m a (Array a)
writeNAligned :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
writeNAligned Int
align = forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWith (forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAligned Int
align)

-- | @writeNAlignedUnmanaged align n@ folds a maximum of @n@ elements from the
-- input stream to an 'Array' whose starting address is aligned to @align@
-- bytes and is allocated using unmanaged memory (never freed).  This could be
-- useful to allocate memory that we need to allocate only once in the lifetime
-- of the program.
--
-- >>> f = Array.newArrayAlignedUnmanaged
-- >>> writeNAlignedUnmanaged a = Array.writeNWith (f a)
-- >>> writeNAlignedUnmanaged a n = Array.appendN (f a n) n
--
-- /Pre-release/
--
{-# INLINE_NORMAL writeNAlignedUnmanaged #-}
writeNAlignedUnmanaged :: forall m a. (MonadIO m, Storable a)
    => Int -> Int -> Fold m a (Array a)
writeNAlignedUnmanaged :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
writeNAlignedUnmanaged Int
align = forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWith (forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAlignedUnmanaged Int
align)

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

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

-- | @writeWith minCount@ folds the whole input to a single array. The array
-- starts at a size big enough to hold minCount elements, the size is doubled
-- every time the array needs to be grown.
--
-- /Caution! Do not use this on infinite streams./
--
-- >>> f n = Array.appendWith (* 2) (Array.newArray n)
-- >>> writeWith n = Fold.rmapM Array.rightSize (f n)
-- >>> writeWith n = Fold.rmapM Array.fromArrayStreamK (Array.writeChunks n)
--
-- /Pre-release/
{-# INLINE_NORMAL writeWith #-}
writeWith :: forall m a. (MonadIO m, Storable a)
    => Int -> Fold m a (Array a)
-- writeWith n = FL.rmapM rightSize $ appendWith (* 2) (newArray n)
writeWith :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Fold m a (Array a)
writeWith Int
elemCount =
    forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> Fold m a b -> Fold m a c
FL.rmapM Array a -> m (Array 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, Storable a) =>
Array a -> a -> m (Array a)
step m (Array a)
initial

    where

    insertElem :: Array a -> a -> m (Array a)
insertElem (Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound) a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
start (PTR_NEXT(end,a)) bound

    initial :: m (Array a)
initial = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
elemCount forall a. Ord a => a -> a -> Bool
< Int
0) forall a b. (a -> b) -> a -> b
$ forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"writeWith: elemCount is negative"
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAligned (forall a. Storable a => a -> Int
alignment (forall a. (?callStack::CallStack) => a
undefined :: a)) Int
elemCount
    step :: Array a -> a -> m (Array a)
step arr :: Array a
arr@(Array ArrayContents
_ Ptr a
start Ptr a
end Ptr a
bound) a
x
        | PTR_NEXT(end,a) > bound = do
        let oldSize = end `minusPtr` start
            newSize = max (oldSize * 2) 1
        arr1 <-
            liftIO
                $ reallocAligned
                    (SIZE_OF(a))
                    (alignment (undefined :: a))
                    newSize
                    arr
        insertElem arr1 x
    step Array a
arr a
x = forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> a -> m (Array a)
insertElem Array a
arr a
x
    extract :: Array a -> m (Array a)
extract = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> m (Array a)
rightSize

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

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

-- | Use the 'writeN' fold instead.
--
-- >>> fromStreamDN n = StreamD.fold (Array.writeN n)
--
{-# INLINE_NORMAL fromStreamDN #-}
fromStreamDN :: forall m a. (MonadIO m, Storable a)
    => Int -> D.Stream m a -> m (Array a)
-- fromStreamDN n = D.fold (writeN n)
fromStreamDN :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> m (Array a)
fromStreamDN Int
limit Stream m a
str = do
    Array a
arr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray Int
limit
    Ptr a
end <- forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> m b
D.foldlM' forall {m :: * -> *} {a} {b}.
(MonadIO m, Storable a) =>
Ptr a -> a -> m (Ptr b)
fwrite (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Array a -> Ptr a
aEnd Array a
arr) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Applicative m =>
Int -> Stream m a -> Stream m a
D.take Int
limit Stream m a
str
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Array a
arr {aEnd :: Ptr a
aEnd = Ptr a
end}

    where

    fwrite :: Ptr a -> a -> m (Ptr b)
fwrite Ptr a
ptr a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
ptr a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ PTR_NEXT(ptr,a)

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

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

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

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

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

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

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

-- | Copy two arrays into a newly allocated array.
{-# INLINE spliceCopy #-}
spliceCopy :: (MonadIO m, Storable a) => Array a -> Array a -> m (Array a)
spliceCopy :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Array a -> m (Array a)
spliceCopy Array a
arr1 Array a
arr2 = do
    let src1 :: Ptr a
src1 = forall a. Array a -> Ptr a
arrStart Array a
arr1
        src2 :: Ptr a
src2 = forall a. Array a -> Ptr a
arrStart Array a
arr2
        len1 :: Int
len1 = forall a. Array a -> Ptr a
aEnd Array a
arr1 forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
src1
        len2 :: Int
len2 = forall a. Array a -> Ptr a
aEnd Array a
arr2 forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
src2

    Array a
arr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray (Int
len1 forall a. Num a => a -> a -> a
+ Int
len2)
    let dst :: Ptr a
dst = forall a. Array a -> Ptr a
arrStart Array a
arr

    -- XXX Should we use copyMutableByteArray# instead? Is there an overhead to
    -- ccall?
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (forall a b. Ptr a -> Ptr b
castPtr Ptr a
dst) (forall a b. Ptr a -> Ptr b
castPtr Ptr a
src1) Int
len1
        ArrayContents -> IO ()
touch (forall a. Array a -> ArrayContents
arrContents Array a
arr1)
        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (forall a b. Ptr a -> Ptr b
castPtr (Ptr a
dst forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len1)) (forall a b. Ptr a -> Ptr b
castPtr Ptr a
src2) Int
len2
        ArrayContents -> IO ()
touch (forall a. Array a -> ArrayContents
arrContents Array a
arr2)
    forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr { aEnd :: Ptr a
aEnd = Ptr a
dst forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
len1 forall a. Num a => a -> a -> a
+ Int
len2) }

-- | Really really unsafe, appends the second array into the first array. If
-- the first array does not have enough space it may cause silent data
-- corruption or if you are lucky a segfault.
{-# INLINE spliceUnsafe #-}
spliceUnsafe :: MonadIO m => Array a -> (Array a, Int) -> m (Array a)
spliceUnsafe :: forall (m :: * -> *) a.
MonadIO m =>
Array a -> (Array a, Int) -> m (Array a)
spliceUnsafe Array a
dst (Array a
src, Int
srcLen) =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
         let psrc :: Ptr a
psrc = forall a. Array a -> Ptr a
arrStart Array a
src
         let pdst :: Ptr a
pdst = forall a. Array a -> Ptr a
aEnd Array a
dst
         forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
pdst forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
srcLen forall a. Ord a => a -> a -> Bool
<= forall a. Array a -> Ptr a
aBound Array a
dst) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
         Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (forall a b. Ptr a -> Ptr b
castPtr Ptr a
pdst) (forall a b. Ptr a -> Ptr b
castPtr Ptr a
psrc) Int
srcLen
         ArrayContents -> IO ()
touch (forall a. Array a -> ArrayContents
arrContents Array a
src)
         ArrayContents -> IO ()
touch (forall a. Array a -> ArrayContents
arrContents Array a
dst)
         forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Array a
dst {aEnd :: Ptr a
aEnd = Ptr a
pdst forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
srcLen}

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

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

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

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

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

-- | Drops the separator byte
{-# INLINE breakOn #-}
breakOn :: MonadIO m
    => Word8 -> Array Word8 -> m (Array Word8, Maybe (Array Word8))
breakOn :: forall (m :: * -> *).
MonadIO m =>
Word8 -> Array Word8 -> m (Array Word8, Maybe (Array Word8))
breakOn Word8
sep arr :: Array Word8
arr@Array{Ptr Word8
ArrayContents
aBound :: Ptr Word8
aEnd :: Ptr Word8
arrStart :: Ptr Word8
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let p :: Ptr Word8
p = Ptr Word8
arrStart
    Ptr Word8
loc <- Ptr Word8 -> Word8 -> CSize -> IO (Ptr Word8)
c_memchr Ptr Word8
p Word8
sep (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Ptr Word8
aEnd forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p)
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
        if Ptr Word8
loc forall a. Eq a => a -> a -> Bool
== forall {a}. Ptr a
nullPtr
        then (Array Word8
arr, forall a. Maybe a
Nothing)
        else
            ( Array
                { arrContents :: ArrayContents
arrContents = ArrayContents
arrContents
                , arrStart :: Ptr Word8
arrStart = Ptr Word8
arrStart
                , aEnd :: Ptr Word8
aEnd = Ptr Word8
loc
                , aBound :: Ptr Word8
aBound = Ptr Word8
loc
                }
            , forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Array
                    { arrContents :: ArrayContents
arrContents = ArrayContents
arrContents
                    , arrStart :: Ptr Word8
arrStart = Ptr Word8
arrStart forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Ptr Word8
loc forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p forall a. Num a => a -> a -> a
+ Int
1)
                    , aEnd :: Ptr Word8
aEnd = Ptr Word8
aEnd
                    , aBound :: Ptr Word8
aBound = Ptr Word8
aBound
                    }
            )

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

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

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

-- | Cast an @Array a@ into an @Array Word8@.
--
-- /Pre-release/
--
asBytes :: Array a -> Array Word8
asBytes :: forall a. Array a -> Array Word8
asBytes = forall a b. Array a -> Array b
castUnsafe

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

-- | Use an @Array a@ as @Ptr a@.
--
-- /Unsafe/
--
-- /Pre-release/
--
asPtrUnsafe :: MonadIO m => Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe :: forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array{Ptr a
ArrayContents
aBound :: Ptr a
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} Ptr a -> m b
f = do
  b
r <- Ptr a -> m b
f Ptr a
arrStart
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
arrContents
  forall (m :: * -> *) a. Monad m => a -> m a
return b
r

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

-- | Compare if two arrays are equal.
--
-- /Pre-release/
{-# INLINE cmp #-}
cmp :: MonadIO m => Array a -> Array a -> m Bool
cmp :: forall (m :: * -> *) a. MonadIO m => Array a -> Array a -> m Bool
cmp Array a
arr1 Array a
arr2 =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let ptr1 :: Ptr a
ptr1 = forall a. Array a -> Ptr a
arrStart Array a
arr1
        let ptr2 :: Ptr a
ptr2 = forall a. Array a -> Ptr a
arrStart Array a
arr2
        let len1 :: Int
len1 = forall a. Array a -> Ptr a
aEnd Array a
arr1 forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
ptr1
        let len2 :: Int
len2 = forall a. Array a -> Ptr a
aEnd Array a
arr2 forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
ptr2

        if Int
len1 forall a. Eq a => a -> a -> Bool
== Int
len2
        then
            if Ptr a
ptr1 forall a. Eq a => a -> a -> Bool
== Ptr a
ptr2
            then forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
            else do
                Bool
r <- Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
memcmp (forall a b. Ptr a -> Ptr b
castPtr Ptr a
ptr1) (forall a b. Ptr a -> Ptr b
castPtr Ptr a
ptr2) Int
len1
                ArrayContents -> IO ()
touch (forall a. Array a -> ArrayContents
arrContents Array a
arr1)
                ArrayContents -> IO ()
touch (forall a. Array a -> ArrayContents
arrContents Array a
arr2)
                forall (m :: * -> *) a. Monad m => a -> m a
return Bool
r
        else forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

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

-- This is a Storable array, we cannot have unevaluated data in it so this is
-- just a no op.
instance NFData (Array a) where
    {-# INLINE rnf #-}
    rnf :: Array a -> ()
rnf Array {} = ()