{-# 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
    , unsafeWithArrayContents
    , 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

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

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

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

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

    -- * Eliminating and Reading

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

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

    -- experimental
    , producer

    -- ** Random reads
    , getIndex
    , getIndexUnsafe
    , getIndices
    -- , 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
    , memcpy
    , memcmp
    , c_memchr
    )
where

#include "inline.hs"

#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 ((.&.))
#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.SVar.Type (adaptState)
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

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

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 :: Storable a => a -> Int -> Int
bytesToElemCount :: a -> Int -> Int
bytesToElemCount a
x Int
n =
    let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf a
x
    in Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
elemSize

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

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

-------------------------------------------------------------------------------
-- 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) =
    (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case MutableByteArray# RealWorld -> State# RealWorld -> State# RealWorld
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
_ = [Char] -> ArrayContents
forall a. HasCallStack => [Char] -> a
error [Char]
"Unsupported foreign ptr"
arrayToFptrContents :: ArrayContents -> ForeignPtrContents
arrayToFptrContents (ArrayContents MutableByteArray# RealWorld
contents) = MutableByteArray# RealWorld -> ForeignPtrContents
PlainPtr MutableByteArray# RealWorld
contents
#endif

-- | Similar to unsafeWithForeignPtr.
{-# INLINE unsafeWithArrayContents #-}
unsafeWithArrayContents :: MonadIO m =>
    ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents :: ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents ArrayContents
contents Ptr a
ptr Ptr a -> m b
f = do
  b
r <- Ptr a -> m b
f Ptr a
ptr
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
  b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
r

-------------------------------------------------------------------------------
-- 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
    { Array a -> ArrayContents
arrContents :: UNPACKIF !ArrayContents
    , Array a -> Ptr a
arrStart :: {-# UNPACK #-} !(Ptr a)      -- ^ first address
    , Array a -> Ptr a
aEnd   :: {-# UNPACK #-} !(Ptr a)        -- ^ first unused address
    , Array a -> Ptr a
aBound :: {-# UNPACK #-} !(Ptr a)        -- ^ first address beyond allocated memory
    }

-- | @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 :: ForeignPtr a -> Ptr a -> Ptr a -> Array a
fromForeignPtrUnsafe fp :: ForeignPtr a
fp@(ForeignPtr Addr#
start ForeignPtrContents
contents) Ptr a
end Ptr a
bound =
    Bool -> Array a -> Array a
forall a. HasCallStack => Bool -> a -> a
assert (ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
fp Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
end Bool -> Bool -> Bool
&& Ptr a
end Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
bound)
           (ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array (ForeignPtrContents -> ArrayContents
fptrToArrayContents ForeignPtrContents
contents) (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
start) Ptr a
end Ptr a
bound)

-------------------------------------------------------------------------------
-- 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 :: (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 = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) Int
0
    (ArrayContents
contents, Ptr a
p) <- Int -> Int -> m (ArrayContents, Ptr a)
alloc Int
size Int
alignSize
    Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> m (Array a)) -> Array a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
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 Ptr a -> Int -> Ptr a
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 :: Int -> Int -> IO (ArrayContents, Ptr a)
newAlignedArrayContents Int
size Int
_align | Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
  [Char] -> IO (ArrayContents, Ptr a)
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"newAlignedArrayContents: size must be >= 0"
newAlignedArrayContents (I# Int#
size) (I# Int#
align) = (State# RealWorld
 -> (# State# RealWorld, (ArrayContents, Ptr a) #))
-> IO (ArrayContents, Ptr a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld
  -> (# State# RealWorld, (ArrayContents, Ptr a) #))
 -> IO (ArrayContents, Ptr a))
-> (State# RealWorld
    -> (# State# RealWorld, (ArrayContents, Ptr a) #))
-> IO (ArrayContents, Ptr a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
    case Int#
-> Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
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 = Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr (ByteArray# -> Addr#
byteArrayContents# (MutableByteArray# RealWorld -> ByteArray#
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, Ptr a
forall a. Ptr a
p) #)
#endif

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

-- | 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 :: Int -> Int -> m (Array a)
newArrayAlignedUnmanaged Int
_align Int
count = do
    let size :: Int
size = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) Int
0
    Ptr a
p <- IO (Ptr a) -> m (Ptr a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Ptr a) -> m (Ptr a)) -> IO (Ptr a) -> m (Ptr a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (Ptr a)
forall a. Int -> IO (Ptr a)
mallocBytes Int
size
    Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> m (Array a)) -> Array a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
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 Ptr a -> Int -> Ptr a
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 :: Int -> Int -> m (Array a)
newArrayAligned = (Int -> Int -> m (ArrayContents, Ptr a))
-> Int -> Int -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> Int -> m (ArrayContents, Ptr a))
-> Int -> Int -> m (Array a)
newArrayWith (\Int
s Int
a -> IO (ArrayContents, Ptr a) -> m (ArrayContents, Ptr a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (ArrayContents, Ptr a) -> m (ArrayContents, Ptr a))
-> IO (ArrayContents, Ptr a) -> m (ArrayContents, Ptr a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IO (ArrayContents, Ptr a)
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 :: Int -> m (Array a)
newArray = Int -> Int -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAligned (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => 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 :: Int -> (Ptr a -> m ()) -> m (Array a)
withNewArrayUnsafe Int
count Ptr a -> m ()
f = do
    Array a
arr <- Int -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray Int
count
    ArrayContents -> Ptr a -> (Ptr a -> m (Array a)) -> m (Array a)
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
arr) (Array a -> Ptr a
forall a. Array a -> Ptr a
arrStart Array a
arr)
        ((Ptr a -> m (Array a)) -> m (Array a))
-> (Ptr a -> m (Array a)) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> Ptr a -> m ()
f Ptr a
p m () -> m (Array a) -> m (Array a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr

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

-- | Write an input stream of (index, value) pairs to an array. Throws an
-- error if any index is out of bounds.
--
-- /Unimplemented/
{-# INLINE putIndices #-}
putIndices :: Array a -> Fold m (Int, a) ()
putIndices :: Array a -> Fold m (Int, a) ()
putIndices = Array a -> Fold m (Int, a) ()
forall a. HasCallStack => a
undefined

-- | 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)
    => Array a -> Int -> a -> m ()
putIndexUnsafe :: Array a -> Int -> a -> m ()
putIndexUnsafe 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
..} Int
i a
x =
    ArrayContents -> Ptr a -> (Ptr a -> m ()) -> m ()
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents ArrayContents
arrContents Ptr a
arrStart ((Ptr a -> m ()) -> m ()) -> (Ptr a -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            elemPtr :: Ptr b
elemPtr = Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
        Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr Any
forall a. Ptr a
elemPtr Ptr Any -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
forall a. Ptr a
elemPtr a
x

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

{-# INLINE putIndexPtr #-}
putIndexPtr :: forall m a. (MonadIO m, Storable a) =>
    Ptr a -> Ptr a -> Int -> a -> m ()
putIndexPtr :: Ptr a -> Ptr a -> Int -> a -> m ()
putIndexPtr Ptr a
ptr Ptr a
end Int
i a
x = do
    let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        elemPtr :: Ptr b
elemPtr = Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr Any
forall a. Ptr a
elemPtr Ptr Any -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
end
    then IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
forall a. Ptr a
elemPtr a
x
    else [Char] -> Int -> m ()
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 arr ix (const (val, ()))
-- >>> f = Array.putIndices
-- >>> putIndex arr ix val = Stream.fold (f arr) (Stream.fromPure (ix, val))
--
-- /Pre-release/
{-# INLINE putIndex #-}
putIndex :: (MonadIO m, Storable a) => Array a -> Int -> a -> m ()
putIndex :: Array a -> Int -> a -> m ()
putIndex Array a
arr Int
i a
x =
    ArrayContents -> Ptr a -> (Ptr a -> m ()) -> m ()
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
arr) (Array a -> Ptr a
forall a. Array a -> Ptr a
arrStart Array a
arr)
        ((Ptr a -> m ()) -> m ()) -> (Ptr a -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> Ptr a -> Ptr a -> Int -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> a -> m ()
putIndexPtr Ptr a
p (Array a -> Ptr a
forall a. Array a -> Ptr a
aEnd Array a
arr) Int
i a
x

-- | Modify a given index of an array using a modifier function.
--
-- /Pre-release/
modifyIndexUnsafe :: forall m a b. (MonadIO m, Storable a) =>
    Array a -> Int -> (a -> (a, b)) -> m b
modifyIndexUnsafe :: Array a -> Int -> (a -> (a, b)) -> m b
modifyIndexUnsafe 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
..} Int
i a -> (a, b)
f = do
    IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> (Ptr a -> IO b) -> IO b
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents ArrayContents
arrContents Ptr a
arrStart ((Ptr a -> IO b) -> IO b) -> (Ptr a -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            elemPtr :: Ptr b
elemPtr = Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
        Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr Any
forall a. Ptr a
elemPtr Ptr Any -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
        a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
forall a. Ptr a
elemPtr
        let (a
x, b
res) = a -> (a, b)
f a
r
        Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
forall a. Ptr a
elemPtr a
x
        b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
res

-- | Modify a given index of an array using a modifier function.
--
-- /Pre-release/
modifyIndex :: forall m a b. (MonadIO m, Storable a) =>
    Array a -> Int -> (a -> (a, b)) -> m b
modifyIndex :: Array a -> Int -> (a -> (a, b)) -> m b
modifyIndex 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
..} Int
i a -> (a, b)
f = do
    IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> (Ptr a -> IO b) -> IO b
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents ArrayContents
arrContents Ptr a
arrStart ((Ptr a -> IO b) -> IO b) -> (Ptr a -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            elemPtr :: Ptr b
elemPtr = Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
        if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr Any
forall a. Ptr a
elemPtr Ptr Any -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd
        then do
            a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
forall a. Ptr a
elemPtr
            let (a
x, b
res) = a -> (a, b)
f a
r
            Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
forall a. Ptr a
elemPtr a
x
            b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
res
        else [Char] -> Int -> IO b
forall a. [Char] -> Int -> a
invalidIndex [Char]
"modifyIndex" Int
i

-- | Modify the array indices generated by the supplied unfold.
--
-- /Pre-release/
modifyIndices :: -- forall m a b. (MonadIO m, Storable a) =>
    Unfold m (Array a) Int -> Array a -> (a -> a) -> m ()
modifyIndices :: Unfold m (Array a) Int -> Array a -> (a -> a) -> m ()
modifyIndices = Unfold m (Array a) Int -> Array a -> (a -> a) -> m ()
forall a. HasCallStack => a
undefined

-- | Modify each element of an array using the supplied modifier function.
--
-- /Unimplemented/
modify :: -- forall m a b. (MonadIO m, Storable a) =>
    Array a -> (a -> a) -> m ()
modify :: Array a -> (a -> a) -> m ()
modify = Array a -> (a -> a) -> m ()
forall a. HasCallStack => a
undefined

-- | Swap the elements at two indices.
--
-- /Pre-release/
swapIndices :: -- (MonadIO m, Storable a) =>
    Array a -> Int -> Int -> m ()
swapIndices :: Array a -> Int -> Int -> m ()
swapIndices = Array a -> Int -> Int -> m ()
forall a. HasCallStack => a
undefined

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

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

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

-- | Round up an array larger than 'largeObjectThreshold' to use the whole
-- block.
{-# INLINE roundUpLargeArray #-}
roundUpLargeArray :: Int -> Int
roundUpLargeArray :: Int -> Int
roundUpLargeArray Int
size =
    if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
largeObjectThreshold
    then
        Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert
            (Int
blockSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& ((Int
blockSize Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
blockSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0))
            ((Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Num a => a -> a
negate Int
blockSize)
    else Int
size

{-
roundUpToPower2 :: Int -> Int
roundUpToPower2 = undefined
-}

-- | @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 :: a -> Int -> Int
allocBytesToBytes a
_ Int
n =
    Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int -> Int
arrayPayloadSize Int
n) (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: 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 :: a -> Int -> Int
allocBytesToElemCount a
x Int
bytes =
    let n :: Int
n = a -> Int -> Int
forall a. Storable a => a -> Int -> Int
bytesToElemCount a
x (a -> Int -> Int
forall a. Storable a => a -> Int -> Int
allocBytesToBytes a
x Int
bytes)
     in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1) Int
n

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

-------------------------------------------------------------------------------
-- 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 :: 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 = IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ do
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
newEnd Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aBound) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
aEnd a
x
    ArrayContents -> IO ()
touch ArrayContents
arrContents
    Array a -> IO (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> IO (Array a)) -> Array a -> IO (Array a)
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 :: 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
..} =
    Ptr a -> Array a -> a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Array a -> a -> m (Array a)
snocNewEnd (Ptr a
aEnd Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) Array 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 :: 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 = IO (Maybe (Array a)) -> m (Maybe (Array a))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (Array a)) -> m (Maybe (Array a)))
-> IO (Maybe (Array a)) -> m (Maybe (Array a))
forall a b. (a -> b) -> a -> b
$ do
    let newEnd :: Ptr b
newEnd = Ptr a
aEnd Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
    if Ptr a
forall a. Ptr a
newEnd Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aBound
    then Array a -> Maybe (Array a)
forall a. a -> Maybe a
Just (Array a -> Maybe (Array a))
-> IO (Array a) -> IO (Maybe (Array a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> Array a -> a -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Array a -> a -> m (Array a)
snocNewEnd Ptr a
forall a. Ptr a
newEnd Array a
arr a
x
    else Maybe (Array a) -> IO (Maybe (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Array a)
forall a. Maybe a
Nothing

reallocWith :: forall m a. (MonadIO m , Storable a) =>
       String
    -> (Int -> Int)
    -> Int
    -> Array a
    -> m (Array a)
reallocWith :: [Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
reallocWith [Char]
label Int -> Int
sizer Int
reqSize Array a
arr = do
    let oldSize :: Int
oldSize = Array a -> Ptr a
forall a. Array a -> Ptr a
aEnd Array a
arr Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Array a -> Ptr a
forall a. Array a -> Ptr a
arrStart Array a
arr
        newSize :: Int
newSize = Int -> Int
sizer Int
oldSize
        safeSize :: Int
safeSize = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
newSize (Int
oldSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
reqSize)
        rounded :: Int
rounded = Int -> Int
roundUpLargeArray Int
safeSize
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
newSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
oldSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
reqSize Bool -> Bool -> Bool
|| [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
badSize) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
rounded Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
safeSize) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Int -> Array a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m (Array a)
realloc Int
rounded Array a
arr

    where

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

-- 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 :: (Int -> Int) -> Array a -> a -> m (Array a)
snocWithRealloc Int -> Int
sizer Array a
arr a
x = do
    let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
    Array a
arr1 <- IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ [Char] -> (Int -> Int) -> Int -> Array a -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
[Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
reallocWith [Char]
"snocWith" Int -> Int
sizer Int
elemSize Array a
arr
    Array a -> a -> m (Array a)
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 :: (Int -> Int) -> Array a -> a -> m (Array a)
snocWith Int -> Int
allocSize Array a
arr a
x = IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ do
    let newEnd :: Ptr b
newEnd = Array a -> Ptr a
forall a. Array a -> Ptr a
aEnd Array a
arr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
    if Ptr a
forall a. Ptr a
newEnd Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Array a -> Ptr a
forall a. Array a -> Ptr a
aBound Array a
arr
    then Ptr a -> Array a -> a -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Array a -> a -> m (Array a)
snocNewEnd Ptr a
forall a. Ptr a
newEnd Array a
arr a
x
    else (Int -> Int) -> Array a -> a -> IO (Array a)
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 :: Array a -> a -> m (Array a)
snocLinear = (Int -> Int) -> Array a -> a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> Int) -> Array a -> a -> m (Array a)
snocWith (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int -> Int
forall a. Storable a => a -> Int -> Int
allocBytesToBytes (a
forall a. HasCallStack => a
undefined :: a) Int
arrayChunkBytes)

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

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

-- XXX See if resizing can be implemented by reading the old array as a stream
-- and then using writeN to the new array.
--
-- | Reallocate the array to the specified size in bytes. If the size is less
-- than the original array the array gets truncated.
{-# NOINLINE reallocAligned #-}
reallocAligned :: Int -> Int -> Int -> Array a -> IO (Array a)
reallocAligned :: Int -> Int -> Int -> Array a -> IO (Array a)
reallocAligned Int
elemSize Int
alignSize Int
newSize 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
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
aEnd Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aBound) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let oldStart :: Ptr a
oldStart = Ptr a
arrStart
        oldSize :: Int
oldSize = Ptr a
aEnd Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
oldStart
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
oldSize Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
elemSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    (ArrayContents
contents, Ptr a
pNew) <- Int -> Int -> IO (ArrayContents, Ptr a)
forall a. Int -> Int -> IO (ArrayContents, Ptr a)
newAlignedArrayContents Int
newSize Int
alignSize
    let size :: Int
size = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
oldSize Int
newSize
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
size Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
elemSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
pNew) (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
oldStart) Int
size
    ArrayContents -> IO ()
touch ArrayContents
arrContents
    Array a -> IO (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> IO (Array a)) -> Array a -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array
        { arrStart :: Ptr a
arrStart = Ptr a
pNew
        , arrContents :: ArrayContents
arrContents = ArrayContents
contents
        , aEnd :: Ptr a
aEnd   = Ptr a
pNew Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
size Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
elemSize))
        , aBound :: Ptr a
aBound = Ptr a
pNew Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
newSize
        }

{-# INLINABLE realloc #-}
realloc :: forall m a. (MonadIO m, Storable a) => Int -> Array a -> m (Array a)
realloc :: Int -> Array a -> m (Array a)
realloc Int
i Array a
arr =
    IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
        (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Array a -> IO (Array a)
forall a. Int -> Int -> Int -> Array a -> IO (Array a)
reallocAligned
            (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a)) Int
i Array a
arr

-- | Change the reserved memory 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).
--
-- /Unimplemented/
{-# INLINE resize #-}
resize :: -- (MonadIO m, Storable a) =>
    Int -> Array a -> m (Array a)
resize :: Int -> Array a -> m (Array a)
resize = Int -> Array a -> m (Array a)
forall a. HasCallStack => a
undefined

-- | Like 'resize' but if the capacity is more than 'largeObjectThreshold' then
-- it is rounded up to the closest power of 2.
--
-- /Unimplemented/
{-# INLINE resizeExp #-}
resizeExp :: -- (MonadIO m, Storable a) =>
    Int -> Array a -> m (Array a)
resizeExp :: Int -> Array a -> m (Array a)
resizeExp = Int -> Array a -> m (Array a)
forall a. HasCallStack => a
undefined

-- | 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 :: 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
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
aEnd Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aBound) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let start :: Ptr a
start = Ptr a
arrStart
        len :: Int
len = Ptr a
aEnd Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
start
        capacity :: Int
capacity = Ptr a
aBound Ptr a -> Ptr a -> Int
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 Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
aEnd
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
target Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    -- We trade off some wastage (25%) to avoid reallocations and copying.
    if Int
target Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
capacity Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
waste
    then Int -> Array a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Array a -> m (Array a)
realloc Int
target Array a
arr
    else Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr

-------------------------------------------------------------------------------
-- Reducing the length
-------------------------------------------------------------------------------

-- XXX Either slice the array or stream it and write it out to a new array?
--
-- | Drop the last n elements of the array to reduce the length by n. The
-- capacity is reallocated using the user supplied function.
--
-- /Unimplemented/
{-# INLINE truncateWith #-}
truncateWith :: -- (MonadIO m, Storable a) =>
    Int -> (Int -> Int) -> Array a -> m (Array a)
truncateWith :: Int -> (Int -> Int) -> Array a -> m (Array a)
truncateWith = Int -> (Int -> Int) -> Array a -> m (Array a)
forall a. HasCallStack => a
undefined

-- | Drop the last n elements of the array to reduce the length by n.
--
-- The capacity is rounded to 1K or 4K if the length is more than the GHC large
-- block threshold.
--
-- /Unimplemented/
{-# INLINE truncate #-}
truncate :: -- (MonadIO m, Storable a) =>
    Int -> Array a -> m (Array a)
truncate :: Int -> Array a -> m (Array a)
truncate = Int -> Array a -> m (Array a)
forall a. HasCallStack => a
undefined

-- | Like 'truncate' but the capacity is rounded to the closest power of 2.
--
-- /Unimplemented/
{-# INLINE truncateExp #-}
truncateExp :: -- (MonadIO m, Storable a) =>
    Int -> Array a -> m (Array a)
truncateExp :: Int -> Array a -> m (Array a)
truncateExp = Int -> Array a -> m (Array a)
forall a. HasCallStack => a
undefined

-------------------------------------------------------------------------------
-- 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) => Array a -> Int -> m a
getIndexUnsafe :: Array a -> Int -> m a
getIndexUnsafe 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
..} Int
i =
    ArrayContents -> Ptr a -> (Ptr a -> m a) -> m a
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents ArrayContents
arrContents Ptr a
arrStart ((Ptr a -> m a) -> m a) -> (Ptr a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            elemPtr :: Ptr b
elemPtr = Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
        Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr Any
forall a. Ptr a
elemPtr Ptr Any -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
        IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
forall a. Ptr a
elemPtr

{-# INLINE getIndexPtr #-}
getIndexPtr :: forall m a. (MonadIO m, Storable a) =>
    Ptr a -> Ptr a -> Int -> m a
getIndexPtr :: Ptr a -> Ptr a -> Int -> m a
getIndexPtr Ptr a
ptr Ptr a
end Int
i = do
    let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        elemPtr :: Ptr b
elemPtr = Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr Any
forall a. Ptr a
elemPtr Ptr Any -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
end
    then IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
forall a. Ptr a
elemPtr
    else [Char] -> Int -> m a
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) => Array a -> Int -> m a
getIndex :: Array a -> Int -> m a
getIndex Array a
arr Int
i =
    ArrayContents -> Ptr a -> (Ptr a -> m a) -> m a
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
arr) (Array a -> Ptr a
forall a. Array a -> Ptr a
arrStart Array a
arr)
        ((Ptr a -> m a) -> m a) -> (Ptr a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> Ptr a -> Ptr a -> Int -> m a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtr Ptr a
p (Array a -> Ptr a
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 :: Ptr a -> Ptr a -> Int -> m a
getIndexPtrRev Ptr a
ptr Ptr a
end Int
i = do
    let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        elemPtr :: Ptr b
elemPtr = Ptr a
end Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr a
forall a. Ptr a
elemPtr Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
ptr
    then IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
forall a. Ptr a
elemPtr
    else [Char] -> Int -> m a
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) => Array a -> Int -> m a
getIndexRev :: Array a -> Int -> m a
getIndexRev Array a
arr Int
i =
    ArrayContents -> Ptr a -> (Ptr a -> m a) -> m a
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
arr) (Array a -> Ptr a
forall a. Array a -> Ptr a
arrStart Array a
arr)
        ((Ptr a -> m a) -> m a) -> (Ptr a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> Ptr a -> Ptr a -> Int -> m a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtrRev Ptr a
p (Array a -> Ptr a
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 getIndices #-}
getIndices :: (MonadIO m, Storable a) =>
    Unfold m (Array a) Int -> Unfold m (Array a) a
getIndices :: Unfold m (Array a) Int -> Unfold m (Array a) a
getIndices (Unfold s -> m (Step s Int)
stepi Array a -> m s
injecti) = (GetIndicesState ArrayContents (Ptr a) (Ptr a) s
 -> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a))
-> (Array a -> m (GetIndicesState ArrayContents (Ptr a) (Ptr a) s))
-> Unfold m (Array a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold GetIndicesState ArrayContents (Ptr a) (Ptr a) s
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
forall a.
Storable a =>
GetIndicesState ArrayContents (Ptr a) (Ptr a) s
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
step Array a -> m (GetIndicesState ArrayContents (Ptr a) (Ptr a) s)
forall a.
Array a -> m (GetIndicesState ArrayContents (Ptr a) (Ptr a) s)
inject

    where

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

    {-# 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 <- s -> m (Step s Int)
stepi s
st
        case Step s Int
r of
            D.Yield Int
i s
s -> do
                a
x <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Ptr a -> Ptr a -> Int -> IO a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Ptr a -> Ptr a -> Int -> m a
getIndexPtr Ptr a
start Ptr a
end Int
i
                Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
 -> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a))
-> Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
forall a b. (a -> b) -> a -> b
$ a
-> GetIndicesState ArrayContents (Ptr a) (Ptr a) s
-> Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
forall s a. a -> s -> Step s a
D.Yield a
x (ArrayContents
-> Ptr a
-> Ptr a
-> s
-> GetIndicesState ArrayContents (Ptr a) (Ptr a) s
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 -> Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
 -> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a))
-> Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
forall a b. (a -> b) -> a -> b
$ GetIndicesState ArrayContents (Ptr a) (Ptr a) s
-> Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
forall s a. s -> Step s a
D.Skip (ArrayContents
-> Ptr a
-> Ptr a
-> s
-> GetIndicesState ArrayContents (Ptr a) (Ptr a) s
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
                IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
                Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
-> m (Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (GetIndicesState ArrayContents (Ptr a) (Ptr a) s) a
forall s a. Step s a
D.Stop

-------------------------------------------------------------------------------
-- 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 :: Int -> Int -> Array a -> Array a
getSliceUnsafe Int
index Int
len (Array ArrayContents
contents Ptr a
start Ptr a
e Ptr a
_) =
    let size :: Int
size = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        fp1 :: Ptr b
fp1 = Ptr a
start Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size)
        end :: Ptr b
end = Ptr Any
forall a. Ptr a
fp1 Ptr Any -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size)
     in Bool -> Array a -> Array a
forall a. HasCallStack => Bool -> a -> a
assert
            (Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr a
forall a. Ptr a
end Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
e)
            (ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
forall a. Ptr a
fp1 Ptr a
forall a. Ptr a
end Ptr a
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 :: Int -> Int -> Array a -> Array a
getSlice Int
index Int
len (Array ArrayContents
contents Ptr a
start Ptr a
e Ptr a
_) =
    let size :: Int
size = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        fp1 :: Ptr b
fp1 = Ptr a
start Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size)
        end :: Ptr b
end = Ptr Any
forall a. Ptr a
fp1 Ptr Any -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size)
     in if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Ptr a
forall a. Ptr a
end Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
e
        then ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
forall a. Ptr a
fp1 Ptr a
forall a. Ptr a
end Ptr a
forall a. Ptr a
end
        else [Char] -> Array a
forall a. HasCallStack => [Char] -> a
error
                ([Char] -> Array a) -> [Char] -> Array a
forall a b. (a -> b) -> a -> b
$ [Char]
"getSlice: invalid slice, index "
                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
index [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" length " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
len

-------------------------------------------------------------------------------
-- 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.
--
-- /Unimplemented/
{-# INLINE reverse #-}
reverse :: Array a -> m Bool
reverse :: Array a -> m Bool
reverse = Array a -> m Bool
forall a. HasCallStack => a
undefined

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

-- | Partition an array into two halves using a partitioning predicate. The
-- first half retains values where the predicate is 'False' and the second half
-- retains values where the predicate is 'True'.
--
-- /Unimplemented/
{-# INLINE partitionBy #-}
partitionBy :: (a -> Bool) -> Array a -> m (Array a, Array a)
partitionBy :: (a -> Bool) -> Array a -> m (Array a, Array a)
partitionBy = (a -> Bool) -> Array a -> m (Array a, Array a)
forall a. HasCallStack => a
undefined

-- | 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 (Array a)
shuffleBy :: (a -> a -> m Bool) -> Array a -> Array a -> m (Array a)
shuffleBy = (a -> a -> m Bool) -> Array a -> Array a -> m (Array a)
forall a. HasCallStack => a
undefined

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

-- | @mergeBy level merge array@ performs a pairwise bottom up fold recursively
-- merging the pairs using the supplied merge function. Level indicates the
-- level in the tree where the fold would stop.
--
-- This performs a random shuffle if the shuffle 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)) -> Array a -> m (Array a)
mergeBy :: Int
-> (Array a -> Array a -> m (Array a)) -> Array a -> m (Array a)
mergeBy = Int
-> (Array a -> Array a -> m (Array a)) -> Array a -> m (Array a)
forall a. HasCallStack => a
undefined

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

-- | /O(1)/ Get the byte length of the array.
--
-- @since 0.7.0
{-# INLINE byteLength #-}
byteLength :: Array a -> Int
byteLength :: 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 Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
arrStart
    in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) Int
len

-- Note: try to avoid the use of length in performance sensitive internal
-- routines as it involves a costly 'div' operation. Instead use the end ptr
-- 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 :: Array a -> Int
length Array a
arr =
    let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        blen :: Int
blen = Array a -> Int
forall a. Array a -> Int
byteLength Array a
arr
     in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
blen Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
elemSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (Int
blen Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
elemSize)

-- | Get the total capacity of an array. An array may have space reserved
-- beyond the current used length of the array.
--
-- /Pre-release/
{-# INLINE byteCapacity #-}
byteCapacity :: Array a -> Int
byteCapacity :: 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 Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
arrStart
    in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) Int
len

-- | The remaining capacity in the array for appending more elements without
-- reallocation.
--
-- /Pre-release/
{-# INLINE bytesFree #-}
bytesFree :: Array a -> Int
bytesFree :: 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 Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
aEnd
    in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) Int
n

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

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

-- | @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 :: 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) =
    (State Stream m (Array a)
 -> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
 -> m (Step
         (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)))
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> Stream m (Array a)
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m (Array a)
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
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' (s -> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
            -- XXX we can pass the module string from the higher level API
            [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"Streamly.Internal.Data.Array.Foreign.Mut.Type.arraysOf: "
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"the size of arrays [" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"] must be a natural number"
        Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound <- IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray Int
n
        Step (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step
   (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)))
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (s
-> ArrayContents
-> Ptr a
-> Ptr a
-> Ptr a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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 (State Stream m a -> State Stream m a
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
                IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ArrayContents -> IO ()
touch ArrayContents
contents
                let end' :: Ptr b
end' = Ptr a
end Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
                Step (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step
   (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)))
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$
                    if Ptr a
forall a. Ptr a
end' Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
bound
                    then GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip
                            (ArrayContents
-> Ptr a
-> Ptr a
-> Ptr a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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
forall a. Ptr a
end' Ptr a
bound (s -> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
forall s contents start end bound.
s -> GroupState s contents start end bound
GroupStart s
s))
                    else GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (s
-> ArrayContents
-> Ptr a
-> Ptr a
-> Ptr a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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
forall a. Ptr a
end' Ptr a
bound)
            D.Skip s
s ->
                Step (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step
   (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)))
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (s
-> ArrayContents
-> Ptr a
-> Ptr a
-> Ptr a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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 ->
                Step (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return
                    (Step
   (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)))
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (ArrayContents
-> Ptr a
-> Ptr a
-> Ptr a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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 GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
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) =
        Step (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step
   (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)))
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ Array a
-> GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)
-> Step
     (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
forall s a. a -> s -> Step s a
D.Yield (ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
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 = Step (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step
        (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return Step (GroupState s ArrayContents (Ptr a) (Ptr a) (Ptr a)) (Array a)
forall s a. Step s a
D.Stop

-- XXX buffer to a list instead?
-- | Buffer the stream into arrays in memory.
{-# INLINE arrayStreamKFromStreamD #-}
arrayStreamKFromStreamD :: forall m a. (MonadIO m, Storable a) =>
    D.Stream m a -> m (K.Stream m (Array a))
arrayStreamKFromStreamD :: Stream m a -> m (Stream m (Array a))
arrayStreamKFromStreamD =
    let n :: Int
n = a -> Int -> Int
forall a. Storable a => a -> Int -> Int
allocBytesToElemCount (a
forall a. HasCallStack => a
undefined :: a) Int
defaultChunkSize
     in (Array a -> Stream m (Array a) -> Stream m (Array a))
-> Stream m (Array a)
-> Stream m (Array a)
-> m (Stream m (Array a))
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr Array a -> Stream m (Array a) -> Stream m (Array a)
forall a (m :: * -> *). a -> Stream m a -> Stream m a
K.cons Stream m (Array a)
forall (m :: * -> *) a. Stream m a
K.nil (Stream m (Array a) -> m (Stream m (Array a)))
-> (Stream m a -> Stream m (Array a))
-> Stream m a
-> m (Stream m (Array a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Stream m a -> Stream m (Array a)
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 :: Stream m (Array a) -> Stream m a
flattenArrays (D.Stream State Stream m (Array a) -> s -> m (Step s (Array a))
step s
state) = (State Stream m a
 -> FlattenState s ArrayContents a
 -> m (Step (FlattenState s ArrayContents a) a))
-> FlattenState s ArrayContents a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a
-> FlattenState s ArrayContents a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a.
State Stream m a
-> FlattenState s ArrayContents a
-> m (Step (FlattenState s ArrayContents a) a)
step' (s -> FlattenState s ArrayContents a
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 (State Stream m a -> State Stream m (Array a)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s ArrayContents a) a
 -> m (Step (FlattenState s ArrayContents a) a))
-> Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
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 ->
                FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. s -> Step s a
D.Skip (s
-> ArrayContents
-> Ptr a
-> Ptr a
-> FlattenState s ArrayContents a
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 -> FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. s -> Step s a
D.Skip (s -> FlattenState s ArrayContents a
forall s contents a. s -> FlattenState s contents a
OuterLoop s
s)
            Step s (Array a)
D.Stop -> Step (FlattenState s ArrayContents a) a
forall s a. Step s a
D.Stop

    step' State Stream m a
_ (InnerLoop s
st ArrayContents
_ Ptr a
p Ptr a
end) | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
end) (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
end) =
        Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s ArrayContents a) a
 -> m (Step (FlattenState s ArrayContents a) a))
-> Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall a b. (a -> b) -> a -> b
$ FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. s -> Step s a
D.Skip (FlattenState s ArrayContents a
 -> Step (FlattenState s ArrayContents a) a)
-> FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall a b. (a -> b) -> a -> b
$ s -> FlattenState s ArrayContents a
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 <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ do
                    a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
                    ArrayContents -> IO ()
touch ArrayContents
contents
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s ArrayContents a) a
 -> m (Step (FlattenState s ArrayContents a) a))
-> Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall a b. (a -> b) -> a -> b
$ a
-> FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. a -> s -> Step s a
D.Yield a
x (s
-> ArrayContents
-> Ptr a
-> Ptr a
-> FlattenState s ArrayContents a
forall s contents a.
s -> contents -> Ptr a -> Ptr a -> FlattenState s contents a
InnerLoop s
st ArrayContents
contents
                            (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) Ptr 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 :: Stream m (Array a) -> Stream m a
flattenArraysRev (D.Stream State Stream m (Array a) -> s -> m (Step s (Array a))
step s
state) = (State Stream m a
 -> FlattenState s ArrayContents a
 -> m (Step (FlattenState s ArrayContents a) a))
-> FlattenState s ArrayContents a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a
-> FlattenState s ArrayContents a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a.
State Stream m a
-> FlattenState s ArrayContents a
-> m (Step (FlattenState s ArrayContents a) a)
step' (s -> FlattenState s ArrayContents a
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 (State Stream m a -> State Stream m (Array a)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s ArrayContents a) a
 -> m (Step (FlattenState s ArrayContents a) a))
-> Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
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 a
aEnd Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
                 in FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. s -> Step s a
D.Skip (s
-> ArrayContents
-> Ptr a
-> Ptr a
-> FlattenState s ArrayContents a
forall s contents a.
s -> contents -> Ptr a -> Ptr a -> FlattenState s contents a
InnerLoop s
s ArrayContents
arrContents Ptr a
forall a. Ptr a
p Ptr a
arrStart)
            D.Skip s
s -> FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. s -> Step s a
D.Skip (s -> FlattenState s ArrayContents a
forall s contents a. s -> FlattenState s contents a
OuterLoop s
s)
            Step s (Array a)
D.Stop -> Step (FlattenState s ArrayContents a) a
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 Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Ptr a
start =
        Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s ArrayContents a) a
 -> m (Step (FlattenState s ArrayContents a) a))
-> Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall a b. (a -> b) -> a -> b
$ FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. s -> Step s a
D.Skip (FlattenState s ArrayContents a
 -> Step (FlattenState s ArrayContents a) a)
-> FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall a b. (a -> b) -> a -> b
$ s -> FlattenState s ArrayContents a
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 <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ do
                    a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
                    ArrayContents -> IO ()
touch ArrayContents
contents
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        let cur :: Ptr b
cur = Ptr a
p Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
        Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s ArrayContents a) a
 -> m (Step (FlattenState s ArrayContents a) a))
-> Step (FlattenState s ArrayContents a) a
-> m (Step (FlattenState s ArrayContents a) a)
forall a b. (a -> b) -> a -> b
$ a
-> FlattenState s ArrayContents a
-> Step (FlattenState s ArrayContents a) a
forall s a. a -> s -> Step s a
D.Yield a
x (s
-> ArrayContents
-> Ptr a
-> Ptr a
-> FlattenState s ArrayContents a
forall s contents a.
s -> contents -> Ptr a -> Ptr a -> FlattenState s contents a
InnerLoop s
st ArrayContents
contents Ptr a
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 :: Array a -> ReadUState a
toReadUState (Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
_) = ArrayContents -> Ptr a -> Ptr a -> ReadUState 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 :: Producer m (Array a) a
producer = (ReadUState a -> m (Step (ReadUState a) a))
-> (Array a -> m (ReadUState a))
-> (ReadUState a -> m (Array a))
-> Producer m (Array a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer ReadUState a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
ReadUState a -> m (Step (ReadUState a) a)
step (ReadUState a -> m (ReadUState a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ReadUState a -> m (ReadUState a))
-> (Array a -> ReadUState a) -> Array a -> m (ReadUState a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> ReadUState a
forall a. Array a -> ReadUState a
toReadUState) ReadUState a -> m (Array a)
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)
        | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
cur Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
end) (Ptr a
cur Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
end) = do
            IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
            Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ReadUState a) a
forall s a. Step s a
D.Stop
    step (ReadUState ArrayContents
contents Ptr a
end Ptr a
cur) = do
            !a
x <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
cur
            let cur1 :: Ptr b
cur1 = Ptr a
cur Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ReadUState a) a -> m (Step (ReadUState a) a))
-> Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall a b. (a -> b) -> a -> b
$ a -> ReadUState a -> Step (ReadUState a) a
forall s a. a -> s -> Step s a
D.Yield a
x (ArrayContents -> Ptr a -> Ptr a -> ReadUState a
forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
end Ptr a
forall a. Ptr a
cur1)

    extract :: ReadUState a -> m (Array a)
extract (ReadUState ArrayContents
contents Ptr a
end Ptr a
cur) = Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> m (Array a)) -> Array a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
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 :: Unfold m (Array a) a
read = Producer m (Array a) a -> Unfold m (Array a) a
forall (m :: * -> *) a b. Producer m a b -> Unfold m a b
Producer.simplify Producer m (Array a) a
forall (m :: * -> *) a.
(MonadIO m, 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 :: Unfold m (Array a) a
readRev = (ReadUState a -> m (Step (ReadUState a) a))
-> (Array a -> m (ReadUState a)) -> Unfold m (Array a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ReadUState a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
ReadUState a -> m (Step (ReadUState a) a)
step Array a -> m (ReadUState a)
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 a
end Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
         in ReadUState a -> m (ReadUState a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ReadUState a -> m (ReadUState a))
-> ReadUState a -> m (ReadUState a)
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> Ptr a -> ReadUState a
forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
start Ptr a
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 Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Ptr a
start = do
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
        Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ReadUState a) a
forall s a. Step s a
D.Stop
    step (ReadUState ArrayContents
contents Ptr a
start Ptr a
p) = do
            a
x <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
            let cur :: Ptr b
cur = Ptr a
p Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
            Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ReadUState a) a -> m (Step (ReadUState a) a))
-> Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall a b. (a -> b) -> a -> b
$ a -> ReadUState a -> Step (ReadUState a) a
forall s a. a -> s -> Step s a
D.Yield a
x (ArrayContents -> Ptr a -> Ptr a -> ReadUState a
forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
start Ptr a
forall a. Ptr a
cur)

-------------------------------------------------------------------------------
-- 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 (p `plusPtr` sizeOf (undefined :: 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 :: 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
..} = IO [a] -> m [a]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [a] -> m [a]) -> IO [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO [a]
go Ptr a
arrStart
    where

    go :: Ptr a -> IO [a]
go Ptr a
p | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
    go Ptr a
p = do
        a
x <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        ArrayContents -> IO ()
touch ArrayContents
arrContents
        (:) a
x ([a] -> [a]) -> IO [a] -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> IO [a]
go (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: 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 :: 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
..} = (State Stream m a -> Ptr a -> m (Step (Ptr a) a))
-> Ptr a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a -> Ptr a -> m (Step (Ptr a) a)
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 | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = Step (Ptr b) a -> m (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Ptr b) a
forall s a. Step s a
D.Stop
    step p
_ Ptr a
p = IO (Step (Ptr b) a) -> m (Step (Ptr b) a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Step (Ptr b) a) -> m (Step (Ptr b) a))
-> IO (Step (Ptr b) a) -> m (Step (Ptr b) a)
forall a b. (a -> b) -> a -> b
$ do
        a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        ArrayContents -> IO ()
touch ArrayContents
arrContents
        Step (Ptr b) a -> IO (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Ptr b) a -> IO (Step (Ptr b) a))
-> Step (Ptr b) a -> IO (Step (Ptr b) a)
forall a b. (a -> b) -> a -> b
$ a -> Ptr b -> Step (Ptr b) a
forall s a. a -> s -> Step s a
D.Yield a
r (Ptr a
p Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))

{-# INLINE toStreamK #-}
toStreamK :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a
toStreamK :: 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
..} = Ptr a -> Stream m a
forall (m :: * -> *). MonadIO m => Ptr a -> Stream m a
go Ptr a
arrStart

    where

    go :: Ptr a -> Stream m a
go Ptr a
p | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = Stream m a
forall (m :: * -> *) a. Stream m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = do
              a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
              ArrayContents -> IO ()
touch ArrayContents
arrContents
              a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO a
elemM m a -> Stream m a -> Stream m a
forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
`K.consM` Ptr a -> Stream m a
go (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: 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 :: 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 a
aEnd Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
    in (State Stream m a -> Ptr a -> m (Step (Ptr a) a))
-> Ptr a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a -> Ptr a -> m (Step (Ptr a) a)
forall (m :: * -> *) p b.
MonadIO m =>
p -> Ptr a -> m (Step (Ptr b) a)
step Ptr a
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 Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Ptr a
arrStart = Step (Ptr b) a -> m (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Ptr b) a
forall s a. Step s a
D.Stop
    step p
_ Ptr a
p = IO (Step (Ptr b) a) -> m (Step (Ptr b) a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Step (Ptr b) a) -> m (Step (Ptr b) a))
-> IO (Step (Ptr b) a) -> m (Step (Ptr b) a)
forall a b. (a -> b) -> a -> b
$ do
        a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        ArrayContents -> IO ()
touch ArrayContents
arrContents
        Step (Ptr b) a -> IO (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Ptr b) a -> IO (Step (Ptr b) a))
-> Step (Ptr b) a -> IO (Step (Ptr b) a)
forall a b. (a -> b) -> a -> b
$ a -> Ptr b -> Step (Ptr b) a
forall s a. a -> s -> Step s a
D.Yield a
r (Ptr a
p Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)))

{-# INLINE toStreamKRev #-}
toStreamKRev :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a
toStreamKRev :: 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 a
aEnd Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
    in Ptr a -> Stream m a
forall (m :: * -> *). MonadIO m => Ptr a -> Stream m a
go Ptr a
forall a. Ptr a
p

    where

    go :: Ptr a -> Stream m a
go Ptr a
p | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Ptr a
arrStart = Stream m a
forall (m :: * -> *) a. Stream m a
K.nil
         | Bool
otherwise =
        let elemM :: IO a
elemM = do
              a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
              ArrayContents -> IO ()
touch ArrayContents
arrContents
              a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO a
elemM m a -> Stream m a -> Stream m a
forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
`K.consM` Ptr a -> Stream m a
go (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: 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' :: (b -> a -> b) -> b -> Array a -> m b
foldl' b -> a -> b
f b
z Array a
arr = (b -> a -> b) -> b -> Stream m a -> m b
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
D.foldl' b -> a -> b
f b
z (Stream m a -> m b) -> Stream m a -> m b
forall a b. (a -> b) -> a -> b
$ Array a -> Stream m a
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 :: (a -> b -> b) -> b -> Array a -> m b
foldr a -> b -> b
f b
z Array a
arr = (a -> b -> b) -> b -> Stream m a -> m b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr a -> b -> b
f b
z (Stream m a -> m b) -> Stream m a -> m b
forall a b. (a -> b) -> a -> b
$ Array a -> Stream m a
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 :: Array a -> ArrayUnsafe a
toArrayUnsafe (Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
_) =
    ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe 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 :: ArrayUnsafe a -> Array a
fromArrayUnsafe (ArrayUnsafe ArrayContents
contents Ptr a
start Ptr a
end) =
         ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
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 :: m (Array a) -> Int -> Fold m a (Array a)
appendNUnsafe m (Array a)
action Int
n =
    (ArrayUnsafe a -> Array a)
-> Fold m a (ArrayUnsafe a) -> Fold m a (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ArrayUnsafe a -> Array a
forall a. ArrayUnsafe a -> Array a
fromArrayUnsafe (Fold m a (ArrayUnsafe a) -> Fold m a (Array a))
-> Fold m a (ArrayUnsafe a) -> Fold m a (Array a)
forall a b. (a -> b) -> a -> b
$ (ArrayUnsafe a -> a -> m (ArrayUnsafe a))
-> m (ArrayUnsafe a) -> Fold m a (ArrayUnsafe a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' ArrayUnsafe a -> a -> m (ArrayUnsafe a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
ArrayUnsafe a -> a -> m (ArrayUnsafe a)
step m (ArrayUnsafe a)
initial

    where

    initial :: m (ArrayUnsafe a)
initial = do
        Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
        arr :: Array a
arr@(Array ArrayContents
_ Ptr a
_ Ptr a
end Ptr a
bound) <- m (Array a)
action
        let free :: Int
free = Ptr a
bound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
end
            elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            needed :: Int
needed = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
elemSize
        -- XXX We can also reallocate if the array has too much free space,
        -- otherwise we lose that space.
        Array a
arr1 <-
            if Int
free Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
needed
            then ([Char] -> (Int -> Int) -> Int -> Array a -> m (Array a))
-> [Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
forall a. a -> a
noinline [Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
[Char] -> (Int -> Int) -> Int -> Array a -> m (Array a)
reallocWith [Char]
"appendNUnsafeWith" (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
needed) Int
needed Array a
arr
            else Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr
        ArrayUnsafe a -> m (ArrayUnsafe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayUnsafe a -> m (ArrayUnsafe a))
-> ArrayUnsafe a -> m (ArrayUnsafe a)
forall a b. (a -> b) -> a -> b
$ Array a -> ArrayUnsafe a
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
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ArrayContents -> IO ()
touch ArrayContents
contents
        let end1 :: Ptr b
end1 = Ptr a
end Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        ArrayUnsafe a -> m (ArrayUnsafe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayUnsafe a -> m (ArrayUnsafe a))
-> ArrayUnsafe a -> m (ArrayUnsafe a)
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe a
forall a. ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe a
ArrayUnsafe ArrayContents
contents Ptr a
start Ptr a
forall a. Ptr a
end1

-- | 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 :: m (Array a) -> Int -> Fold m a (Array a)
appendN m (Array a)
initial Int
n = Int -> Fold m a (Array a) -> Fold m a (Array a)
forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n (m (Array a) -> Int -> Fold m a (Array a)
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 :: (Int -> Int) -> m (Array a) -> Fold m a (Array a)
appendWith Int -> Int
sizer = (Array a -> a -> m (Array a)) -> m (Array a) -> Fold m a (Array a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' ((Int -> Int) -> Array a -> a -> m (Array a)
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 :: m (Array a) -> Fold m a (Array a)
append = (Int -> Int) -> m (Array a) -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> Int) -> m (Array a) -> Fold m a (Array a)
appendWith (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2)

-- | @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 :: (Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWith Int -> m (Array a)
alloc Int
n = Int -> Fold m a (Array a) -> Fold m a (Array a)
forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n ((Int -> m (Array a)) -> Int -> Fold m a (Array a)
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 :: Int -> Fold m a (Array a)
writeN = (Int -> m (Array a)) -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWith Int -> m (Array a)
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 :: Int -> Int -> Fold m a (Array a)
writeNAligned Int
align = (Int -> m (Array a)) -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWith (Int -> Int -> m (Array a)
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 :: Int -> Int -> Fold m a (Array a)
writeNAlignedUnmanaged Int
align = (Int -> m (Array a)) -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWith (Int -> Int -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAlignedUnmanaged Int
align)

-- 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 :: (Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWithUnsafe Int -> m (Array a)
alloc Int
n = (ArrayUnsafe a -> a -> m (Step (ArrayUnsafe a) (Array a)))
-> m (Step (ArrayUnsafe a) (Array a))
-> (ArrayUnsafe a -> m (Array a))
-> Fold m a (Array a)
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold ArrayUnsafe a -> a -> m (Step (ArrayUnsafe a) (Array a))
forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
ArrayUnsafe a -> a -> m (Step (ArrayUnsafe a) b)
step m (Step (ArrayUnsafe a) (Array a))
forall b. m (Step (ArrayUnsafe a) b)
initial (Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> m (Array a))
-> (ArrayUnsafe a -> Array a) -> ArrayUnsafe a -> m (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ArrayUnsafe a -> Array a
forall a. ArrayUnsafe a -> Array a
fromArrayUnsafe)

    where

    initial :: m (Step (ArrayUnsafe a) b)
initial = ArrayUnsafe a -> Step (ArrayUnsafe a) b
forall s b. s -> Step s b
FL.Partial (ArrayUnsafe a -> Step (ArrayUnsafe a) b)
-> (Array a -> ArrayUnsafe a) -> Array a -> Step (ArrayUnsafe a) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> ArrayUnsafe a
forall a. Array a -> ArrayUnsafe a
toArrayUnsafe (Array a -> Step (ArrayUnsafe a) b)
-> m (Array a) -> m (Step (ArrayUnsafe a) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (Array a)
alloc (Int -> Int -> Int
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
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ArrayContents -> IO ()
touch ArrayContents
contents
        Step (ArrayUnsafe a) b -> m (Step (ArrayUnsafe a) b)
forall (m :: * -> *) a. Monad m => a -> m a
return
          (Step (ArrayUnsafe a) b -> m (Step (ArrayUnsafe a) b))
-> Step (ArrayUnsafe a) b -> m (Step (ArrayUnsafe a) b)
forall a b. (a -> b) -> a -> b
$ ArrayUnsafe a -> Step (ArrayUnsafe a) b
forall s b. s -> Step s b
FL.Partial
          (ArrayUnsafe a -> Step (ArrayUnsafe a) b)
-> ArrayUnsafe a -> Step (ArrayUnsafe a) b
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe a
forall a. ArrayContents -> Ptr a -> Ptr a -> ArrayUnsafe a
ArrayUnsafe ArrayContents
contents Ptr a
start (Ptr a
end Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: 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 :: Int -> Fold m a (Array a)
writeNUnsafe = (Int -> m (Array a)) -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> m (Array a)) -> Int -> Fold m a (Array a)
writeNWithUnsafe Int -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> m (Array a)
newArray

-- 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 :: Int -> Fold m a (Stream n (Array a))
writeChunks Int
n = Fold m a (Array a)
-> Fold m (Array a) (Stream n (Array a))
-> Fold m a (Stream n (Array a))
forall (m :: * -> *) a b c.
Monad m =>
Fold m a b -> Fold m b c -> Fold m a c
FL.many (Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Fold m a (Array a)
writeN Int
n) Fold m (Array a) (Stream n (Array a))
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 :: Int -> Fold m a (Array a)
writeWith Int
elemCount =
    (Array a -> m (Array a))
-> Fold m a (Array a) -> Fold m a (Array a)
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 (Fold m a (Array a) -> Fold m a (Array a))
-> Fold m a (Array a) -> Fold m a (Array a)
forall a b. (a -> b) -> a -> b
$ (Array a -> a -> m (Array a)) -> m (Array a) -> Fold m a (Array a)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' Array a -> a -> m (Array a)
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
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
end a
x
        let end1 :: Ptr b
end1 = Ptr a
end Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> m (Array a)) -> Array a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
start Ptr a
forall a. Ptr a
end1 Ptr a
bound

    initial :: m (Array a)
initial = do
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
elemCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"writeWith: elemCount is negative"
        IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> m (Array a)
newArrayAligned (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => 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 a
end Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a) Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
> Ptr a
bound = do
        let oldSize :: Int
oldSize = Ptr a
end Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
start
            newSize :: Int
newSize = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
oldSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) Int
1
        Array a
arr1 <-
            IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
                (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Array a -> IO (Array a)
forall a. Int -> Int -> Int -> Array a -> IO (Array a)
reallocAligned
                    (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
                    (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a))
                    Int
newSize
                    Array a
arr
        Array a -> a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> a -> m (Array a)
insertElem Array a
arr1 a
x
    step Array a
arr a
x = Array a -> a -> m (Array a)
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 = IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array a) -> m (Array a))
-> (Array a -> IO (Array a)) -> Array a -> m (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> IO (Array a)
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 :: Fold m a (Array a)
write = Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Fold m a (Array a)
writeWith (a -> Int -> Int
forall a. Storable a => a -> Int -> Int
allocBytesToElemCount (a
forall a. HasCallStack => a
undefined :: a) Int
arrayChunkBytes)

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

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

    where

    fwrite :: Ptr a -> a -> m (Ptr b)
fwrite Ptr a
ptr a
x = do
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
ptr a
x
        Ptr b -> m (Ptr b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ptr b -> m (Ptr b)) -> Ptr b -> m (Ptr b)
forall a b. (a -> b) -> a -> b
$ Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: 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 :: Int -> [a] -> m (Array a)
fromListN Int
n [a]
xs = Int -> Stream m a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> m (Array a)
fromStreamDN Int
n (Stream m a -> m (Array a)) -> Stream m a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ [a] -> Stream m a
forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

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

{-# INLINE arrayStreamKLength #-}
arrayStreamKLength :: (Monad m, Storable a) => K.Stream m (Array a) -> m Int
arrayStreamKLength :: Stream m (Array a) -> m Int
arrayStreamKLength Stream m (Array a)
as = (Int -> Int -> Int) -> Int -> Stream m Int -> m Int
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
K.foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ((Array a -> Int) -> Stream m (Array a) -> Stream m Int
forall a b (m :: * -> *). (a -> b) -> Stream m a -> Stream m b
K.map Array a -> Int
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 :: Stream m (Array a) -> m (Array a)
fromArrayStreamK Stream m (Array a)
as = do
    Int
len <- Stream m (Array a) -> m Int
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Stream m (Array a) -> m Int
arrayStreamKLength Stream m (Array a)
as
    Int -> Stream m a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> m (Array a)
fromStreamDN Int
len (Stream m a -> m (Array a)) -> Stream m a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Unfold m (Array a) a -> Stream m (Array a) -> Stream m a
forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> Stream m a -> Stream m b
D.unfoldMany Unfold m (Array a) a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Unfold m (Array a) a
read (Stream m (Array a) -> Stream m a)
-> Stream m (Array a) -> Stream m a
forall a b. (a -> b) -> a -> b
$ Stream m (Array a) -> Stream m (Array a)
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 :: Stream m a -> m (Array a)
fromStreamD Stream m a
m = Stream m a -> m (Stream m (Array a))
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m a -> m (Stream m (Array a))
arrayStreamKFromStreamD Stream m a
m m (Stream m (Array a))
-> (Stream m (Array a) -> m (Array a)) -> m (Array a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Stream m (Array a) -> m (Array a)
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
{-# INLINABLE fromList #-}
fromList :: (MonadIO m, Storable a) => [a] -> m (Array a)
fromList :: [a] -> m (Array a)
fromList [a]
xs = Stream m a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m a -> m (Array a)
fromStreamD (Stream m a -> m (Array a)) -> Stream m a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ [a] -> Stream m a
forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

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

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

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

    -- XXX Should we use copyMutableByteArray# instead? Is there an overhead to
    -- ccall?
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
dst) (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
src1) Int
len1
        ArrayContents -> IO ()
touch (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
arr1)
        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (Ptr Any -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr (Ptr a
dst Ptr a -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len1)) (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
src2) Int
len2
        ArrayContents -> IO ()
touch (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
arr2)
    Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr { aEnd :: Ptr a
aEnd = Ptr a
dst Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
len1 Int -> Int -> Int
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 :: Array a -> (Array a, Int) -> m (Array a)
spliceUnsafe Array a
dst (Array a
src, Int
srcLen) =
    IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array a) -> m (Array a)) -> IO (Array a) -> m (Array a)
forall a b. (a -> b) -> a -> b
$ do
         let psrc :: Ptr a
psrc = Array a -> Ptr a
forall a. Array a -> Ptr a
arrStart Array a
src
         let pdst :: Ptr a
pdst = Array a -> Ptr a
forall a. Array a -> Ptr a
aEnd Array a
dst
         Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
pdst Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
srcLen Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Array a -> Ptr a
forall a. Array a -> Ptr a
aBound Array a
dst) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
         Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memcpy (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
pdst) (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
psrc) Int
srcLen
         ArrayContents -> IO ()
touch (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
src)
         ArrayContents -> IO ()
touch (Array a -> ArrayContents
forall a. Array a -> ArrayContents
arrContents Array a
dst)
         Array a -> IO (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> IO (Array a)) -> Array a -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ Array a
dst {aEnd :: Ptr a
aEnd = Ptr a
pdst Ptr a -> Int -> Ptr a
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 :: (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)
-}
    Bool -> m () -> m ()
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
end Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
bound) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
    let srcLen :: Int
srcLen = Array a -> Ptr a
forall a. Array a -> Ptr a
aEnd Array a
src Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Array a -> Ptr a
forall a. Array a -> Ptr a
arrStart Array a
src

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

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

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

-- | Drops the separator byte
{-# INLINE breakOn #-}
breakOn :: MonadIO m
    => Word8 -> Array Word8 -> m (Array Word8, Maybe (Array Word8))
breakOn :: 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
..} = IO (Array Word8, Maybe (Array Word8))
-> m (Array Word8, Maybe (Array Word8))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array Word8, Maybe (Array Word8))
 -> m (Array Word8, Maybe (Array Word8)))
-> IO (Array Word8, Maybe (Array Word8))
-> m (Array Word8, Maybe (Array Word8))
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 (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CSize) -> Int -> CSize
forall a b. (a -> b) -> a -> b
$ Ptr Word8
aEnd Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p)
    (Array Word8, Maybe (Array Word8))
-> IO (Array Word8, Maybe (Array Word8))
forall (m :: * -> *) a. Monad m => a -> m a
return ((Array Word8, Maybe (Array Word8))
 -> IO (Array Word8, Maybe (Array Word8)))
-> (Array Word8, Maybe (Array Word8))
-> IO (Array Word8, Maybe (Array Word8))
forall a b. (a -> b) -> a -> b
$
        if Ptr Word8
loc Ptr Word8 -> Ptr Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word8
forall a. Ptr a
nullPtr
        then (Array Word8
arr, Maybe (Array Word8)
forall a. Maybe a
Nothing)
        else
            ( Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
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
                }
            , Array Word8 -> Maybe (Array Word8)
forall a. a -> Maybe a
Just (Array Word8 -> Maybe (Array Word8))
-> Array Word8 -> Maybe (Array Word8)
forall a b. (a -> b) -> a -> b
$ Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array
                    { arrContents :: ArrayContents
arrContents = ArrayContents
arrContents
                    , arrStart :: Ptr Word8
arrStart = Ptr Word8
arrStart Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Ptr Word8
loc Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p Int -> Int -> Int
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 :: 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 = Array a -> Int
forall a. Storable a => Array a -> Int
length Array a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    in  if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        then [Char] -> (Array a, Array a)
forall a. HasCallStack => [Char] -> a
error [Char]
"sliceAt: negative array index"
        else if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxIndex
             then [Char] -> (Array a, Array a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (Array a, Array a)) -> [Char] -> (Array a, Array a)
forall a b. (a -> b) -> a -> b
$ [Char]
"sliceAt: specified array index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i
                        [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is beyond the maximum index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
maxIndex
             else let off :: Int
off = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
                      p :: Ptr b
p = Ptr a
arrStart Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off
                in ( Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array
                  { arrContents :: ArrayContents
arrContents = ArrayContents
arrContents
                  , arrStart :: Ptr a
arrStart = Ptr a
arrStart
                  , aEnd :: Ptr a
aEnd = Ptr a
forall a. Ptr a
p
                  , aBound :: Ptr a
aBound = Ptr a
forall a. Ptr a
p
                  }
                , Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array
                  { arrContents :: ArrayContents
arrContents = ArrayContents
arrContents
                  , arrStart :: Ptr a
arrStart = Ptr a
arrStart Ptr a -> Int -> Ptr a
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 :: Array a -> Array b
castUnsafe (Array ArrayContents
contents Ptr a
start Ptr a
end Ptr a
bound) =
    ArrayContents -> Ptr b -> Ptr b -> Ptr b -> Array b
forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents (Ptr a -> Ptr b
forall a b. Ptr a -> Ptr b
castPtr Ptr a
start) (Ptr a -> Ptr b
forall a b. Ptr a -> Ptr b
castPtr Ptr a
end) (Ptr a -> Ptr b
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 :: Array a -> Array Word8
asBytes = Array a -> Array Word8
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 :: Array a -> Maybe (Array b)
cast Array a
arr =
    let len :: Int
len = Array a -> Int
forall a. Array a -> Int
byteLength Array a
arr
        r :: Int
r = Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` b -> Int
forall a. Storable a => a -> Int
sizeOf (b
forall a. HasCallStack => a
undefined :: b)
     in if Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
        then Maybe (Array b)
forall a. Maybe a
Nothing
        else Array b -> Maybe (Array b)
forall a. a -> Maybe a
Just (Array b -> Maybe (Array b)) -> Array b -> Maybe (Array b)
forall a b. (a -> b) -> a -> b
$ Array a -> Array b
forall a b. Array a -> Array b
castUnsafe Array a
arr

-- | Use an @Array a@ as @Ptr b@.
--
-- /Unsafe/
--
-- /Pre-release/
--
asPtrUnsafe :: Array a -> (Ptr b -> IO c) -> IO c
asPtrUnsafe :: Array a -> (Ptr b -> IO c) -> IO c
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 b -> IO c
act = do
    ArrayContents -> Ptr a -> (Ptr a -> IO c) -> IO c
forall (m :: * -> *) a b.
MonadIO m =>
ArrayContents -> Ptr a -> (Ptr a -> m b) -> m b
unsafeWithArrayContents ArrayContents
arrContents Ptr a
arrStart ((Ptr a -> IO c) -> IO c) -> (Ptr a -> IO c) -> IO c
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> Ptr b -> IO c
act (Ptr a -> Ptr b
forall a b. Ptr a -> Ptr b
castPtr Ptr a
ptr)

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

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

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