{-# 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
--
-- Unboxed pinned mutable array type for 'Storable' types with an option to use
-- foreign (non-GHC) memory allocators. Fulfils the following goals:
--
-- * Random access (array)
-- * Efficient storage (unboxed)
-- * Performance (unboxed access)
-- * Performance - in-place operations (mutable)
-- * Performance - GC (pinned, mutable)
-- * interfacing with OS (pinned)
-- * Fragmentation control (foreign allocators)
--
-- Stream and Fold APIs allow easy, efficient and convenient operations on
-- arrays.

module Streamly.Internal.Data.Array.Foreign.Mut.Type
    (
    -- * Type
    -- $arrayNotes
      Array (..)

    -- * Construction
    , mutableArray
    , unsafeWithNewArray
    , newArray
    , newArrayAligned
    , newArrayAlignedUnmanaged
    , newArrayAlignedAllocWith

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

    -- * Resizing
    , realloc
    , shrinkToFit

    -- * Size
    , length
    , byteLength
    , byteCapacity

    -- * Random access
    , unsafeIndexIO
    , unsafeIndex

    -- * Mutation
    , unsafeWriteIndex
    , unsafeSnoc
    , snoc

    -- * Folding
    , foldl'
    , foldr

    -- * Composable Folds
    , toArrayMinChunk
    , writeNAllocWith
    , writeN
    , writeNUnsafe
    , ArrayUnsafe (..)
    , writeNAligned
    , writeNAlignedUnmanaged
    , write
    , writeAligned

    -- * Unfolds
    , ReadUState
    , read
    , readRev
    , producer
    , flattenArrays
    , flattenArraysRev

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

    -- * Combining
    , spliceWith
    , spliceWithDoubling
    , spliceTwo

    -- * Splitting
    , breakOn
    , splitAt

    -- * Stream of arrays
    , arraysOf
    , bufferChunks
    , writeChunks

    -- * Utilities
    , defaultChunkSize
    , mkChunkSize
    , mkChunkSizeKB
    , bytesToElemCount
    , unsafeInlineIO
    , memcpy
    , memcmp
    )
where

#include "inline.hs"

import Control.Exception (assert)
import Control.DeepSeq (NFData(..))
import Control.Monad (when, void)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Functor.Identity (runIdentity)
#if __GLASGOW_HASKELL__ < 808
import Data.Semigroup (Semigroup(..))
#endif
import Data.Word (Word8)
import Foreign.C.Types (CSize(..), CInt(..))
import Foreign.ForeignPtr (touchForeignPtr)
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import Foreign.Ptr (plusPtr, minusPtr, castPtr, nullPtr)
import Foreign.Storable (Storable(..))
import GHC.Base (nullAddr#, realWorld#, build)
import GHC.Exts (IsList, IsString(..))
import GHC.ForeignPtr (ForeignPtr(..))
import GHC.IO (IO(IO), unsafePerformIO)
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 (adaptState)
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Text.Read (readPrec, readListPrec, readListPrecDefault)

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

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

#if MIN_VERSION_base(4,10,0)
import Foreign.ForeignPtr (plusForeignPtr)
#else
import GHC.Base (Int(..), plusAddr#)
import GHC.ForeignPtr (ForeignPtr(..))
plusForeignPtr :: ForeignPtr a -> Int -> ForeignPtr b
plusForeignPtr (ForeignPtr addr c) (I# d) = ForeignPtr (plusAddr# addr d) c
#endif

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

-- $arrayNotes
--
-- We can use a 'Storable' constraint in the Array type and the constraint can
-- be automatically provided to a function that pattern matches on the Array
-- type. However, it has huge performance cost, so we do not use it.
-- Investigate a GHC improvement possiblity.
--
-- XXX Rename the fields to better names.
--
data Array a =
#ifdef DEVBUILD
    Storable a =>
#endif
    Array
    { Array a -> ForeignPtr a
aStart :: {-# UNPACK #-} !(ForeignPtr 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
    }

{-# INLINE mutableArray #-}
mutableArray ::
#ifdef DEVBUILD
    Storable a =>
#endif
    ForeignPtr a -> Ptr a -> Ptr a -> Array a
mutableArray :: ForeignPtr a -> Ptr a -> Ptr a -> Array a
mutableArray = ForeignPtr a -> Ptr a -> Ptr a -> Array a
forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array

-------------------------------------------------------------------------------
-- Utility functions
-------------------------------------------------------------------------------

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)

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

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

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

{-# INLINE unsafeInlineIO #-}
unsafeInlineIO :: IO a -> a
unsafeInlineIO :: IO a -> a
unsafeInlineIO (IO State# RealWorld -> (# State# RealWorld, a #)
m) = case State# RealWorld -> (# State# RealWorld, a #)
m State# RealWorld
realWorld# of (# State# RealWorld
_, a
r #) -> a
r

{-# 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. Num a => a -> a -> a
+ Int
elemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
elemSize


-- | GHC memory management allocation header overhead
allocOverhead :: Int
allocOverhead :: Int
allocOverhead = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
forall a. Storable a => a -> Int
sizeOf (Int
forall a. HasCallStack => a
undefined :: Int)

mkChunkSize :: Int -> Int
mkChunkSize :: Int -> Int
mkChunkSize Int
n = let size :: Int
size = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
allocOverhead in Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
size Int
0

mkChunkSizeKB :: Int -> Int
mkChunkSizeKB :: Int -> Int
mkChunkSizeKB Int
n = Int -> Int
mkChunkSize (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k)
   where k :: Int
k = Int
1024

-- | Default maximum buffer size in bytes, for reading from and writing to IO
-- devices, the value is 32KB minus GHC allocation overhead, which is a few
-- bytes, so that the actual allocation is 32KB.
defaultChunkSize :: Int
defaultChunkSize :: Int
defaultChunkSize = Int -> Int
mkChunkSizeKB Int
32

-- | Remove the free space from an Array.
shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)
shrinkToFit :: Array a -> IO (Array a)
shrinkToFit arr :: Array a
arr@Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} = 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 start :: Ptr a
start = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
    let used :: Int
used = Ptr a
aEnd Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
start
        waste :: Int
waste = Ptr a
aBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
aEnd
    -- if used == waste == 0 then do not realloc
    -- if the wastage is more than 25% of the array then realloc
    if Int
used 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 -> IO (Array a)
forall a. Storable a => Int -> Array a -> IO (Array a)
realloc Int
used Array a
arr
    else Array a -> IO (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr

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

-- | allocate a new array using the provided allocator function.
{-# INLINE newArrayAlignedAllocWith #-}
newArrayAlignedAllocWith :: forall a. Storable a
    => (Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
newArrayAlignedAllocWith :: (Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
newArrayAlignedAllocWith Int -> Int -> IO (ForeignPtr a)
alloc Int
alignSize Int
count = do
    let size :: Int
size = 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)
    ForeignPtr a
fptr <- Int -> Int -> IO (ForeignPtr a)
alloc Int
size Int
alignSize
    let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
fptr
    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. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array
        { aStart :: ForeignPtr a
aStart = ForeignPtr a
fptr
        , 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
        }

-- | Allocate a new array aligned to the specified alignmend and using
-- unmanaged pinned memory. The memory will not be automatically freed by GHC.
-- This could be useful in allocate once global data structures. Use carefully
-- as incorrect use can lead to memory leak.
{-# INLINE newArrayAlignedUnmanaged #-}
newArrayAlignedUnmanaged :: forall a. Storable a => Int -> Int -> IO (Array a)
newArrayAlignedUnmanaged :: Int -> Int -> IO (Array a)
newArrayAlignedUnmanaged =
    (Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
forall a.
Storable a =>
(Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
newArrayAlignedAllocWith Int -> Int -> IO (ForeignPtr a)
forall a. Int -> Int -> IO (ForeignPtr a)
Malloc.mallocForeignPtrAlignedUnmanagedBytes

{-# INLINE newArrayAligned #-}
newArrayAligned :: forall a. Storable a => Int -> Int -> IO (Array a)
newArrayAligned :: Int -> Int -> IO (Array a)
newArrayAligned = (Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
forall a.
Storable a =>
(Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
newArrayAlignedAllocWith Int -> Int -> IO (ForeignPtr a)
forall a. Int -> Int -> IO (ForeignPtr a)
Malloc.mallocForeignPtrAlignedBytes

-- XXX can unaligned allocation be more efficient when alignment is not needed?
--
-- | Allocate an array that can hold 'count' items.  The memory of the array is
-- uninitialized.
--
-- Note that this is internal routine, the reference to this array cannot be
-- given out until the array has been written to and frozen.
{-# INLINE newArray #-}
newArray :: forall a. Storable a => Int -> IO (Array a)
newArray :: Int -> IO (Array a)
newArray = Int -> Int -> IO (Array a)
forall a. Storable a => Int -> Int -> IO (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.
{-# INLINE unsafeWithNewArray #-}
unsafeWithNewArray :: forall a. Storable a => Int -> (Ptr a -> IO ()) -> IO (Array a)
unsafeWithNewArray :: Int -> (Ptr a -> IO ()) -> IO (Array a)
unsafeWithNewArray Int
count Ptr a -> IO ()
f = do
    Array a
arr <- Int -> IO (Array a)
forall a. Storable a => Int -> IO (Array a)
newArray Int
count
    ForeignPtr a -> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
arr) ((Ptr a -> IO (Array a)) -> IO (Array a))
-> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> Ptr a -> IO ()
f Ptr a
p IO () -> IO (Array a) -> IO (Array a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Array a -> IO (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr

-------------------------------------------------------------------------------
-- snoc
-------------------------------------------------------------------------------

{-# INLINE unsafeWriteIndex #-}
unsafeWriteIndex :: forall a. Storable a => Array a -> Int -> a -> IO (Array a)
unsafeWriteIndex :: Array a -> Int -> a -> IO (Array a)
unsafeWriteIndex arr :: Array a
arr@Array {Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} Int
i a
x =
    ForeignPtr a -> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr a
aStart
        ((Ptr a -> IO (Array a)) -> IO (Array a))
-> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
begin -> do
              Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr a
begin Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (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))) a
x
              Array a -> IO (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr

-- XXX grow the array when we are beyond bound.
--
-- Internal routine for when the array is being created. Appends one item at
-- the end of the array. Useful when sequentially writing a stream to the
-- array.
{-# INLINE unsafeSnoc #-}
unsafeSnoc :: forall a. Storable a => Array a -> a -> IO (Array a)
unsafeSnoc :: Array a -> a -> IO (Array a)
unsafeSnoc arr :: Array a
arr@Array {Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} a
x = do
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ptr a
aEnd Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aBound) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"BUG: unsafeSnoc: writing beyond array bounds"
    Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
aEnd a
x
    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
    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
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)}

{-# INLINE snoc #-}
snoc :: forall a. Storable a => Array a -> a -> IO (Array a)
snoc :: Array a -> a -> IO (Array a)
snoc arr :: Array a
arr@Array {Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} a
x =
    if Ptr a
aEnd Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aBound
    then do
        let oldStart :: Ptr a
oldStart = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
            size :: Int
size = Ptr a
aEnd Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
oldStart
            newSize :: Int
newSize = Int
size 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)
        ForeignPtr a
newPtr <-
            Int -> Int -> IO (ForeignPtr a)
forall a. Int -> Int -> IO (ForeignPtr a)
Malloc.mallocForeignPtrAlignedBytes
                Int
newSize
                (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a))
        ForeignPtr a -> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr a
newPtr ((Ptr a -> IO (Array a)) -> IO (Array a))
-> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
pNew -> do
            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
            Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr a
pNew Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size) a
x
            ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
            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. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array
                    { aStart :: ForeignPtr a
aStart = ForeignPtr a
newPtr
                    , 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
+ a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
                    , aBound :: Ptr a
aBound = Ptr a
pNew Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
newSize
                    }
    else do
        Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
aEnd a
x
        ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
        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
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)}

-------------------------------------------------------------------------------
-- re-allocate
-------------------------------------------------------------------------------

-- | 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 -> Array a -> IO (Array a)
reallocAligned :: Int -> Int -> Array a -> IO (Array a)
reallocAligned Int
alignSize Int
newSize Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} = 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 = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
    let size :: Int
size = Ptr a
aEnd Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
oldStart
    ForeignPtr a
newPtr <- Int -> Int -> IO (ForeignPtr a)
forall a. Int -> Int -> IO (ForeignPtr a)
Malloc.mallocForeignPtrAlignedBytes Int
newSize Int
alignSize
    ForeignPtr a -> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr a
newPtr ((Ptr a -> IO (Array a)) -> IO (Array a))
-> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
pNew -> do
        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
        ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
        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. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array
            { aStart :: ForeignPtr a
aStart = ForeignPtr a
newPtr
            , aEnd :: Ptr a
aEnd   = Ptr a
pNew Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size
            , aBound :: Ptr a
aBound = Ptr a
pNew Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
newSize
            }

-- XXX can unaligned allocation be more efficient when alignment is not needed?
{-# INLINABLE realloc #-}
realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)
realloc :: Int -> Array a -> IO (Array a)
realloc = Int -> Int -> Array a -> IO (Array a)
forall a. Int -> Int -> Array a -> IO (Array a)
reallocAligned (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a))

-------------------------------------------------------------------------------
-- Elimination
-------------------------------------------------------------------------------

-- | Return element at the specified index without checking the bounds.
--
-- Unsafe because it does not check the bounds of the array.
{-# INLINE_NORMAL unsafeIndexIO #-}
unsafeIndexIO :: forall a. Storable a => Array a -> Int -> IO a
unsafeIndexIO :: Array a -> Int -> IO a
unsafeIndexIO Array {Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} Int
i =
        ForeignPtr a -> (Ptr a -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr a
aStart ((Ptr a -> IO a) -> IO a) -> (Ptr a -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> do
        let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            elemOff :: Ptr b
elemOff = Ptr a
p 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 b. Ptr b
elemOff 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 ())
        Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
forall b. Ptr b
elemOff

-- | Return element at the specified index without checking the bounds.
{-# INLINE_NORMAL unsafeIndex #-}
unsafeIndex :: forall a. Storable a => Array a -> Int -> a
unsafeIndex :: Array a -> Int -> a
unsafeIndex Array a
arr Int
i = let !r :: a
r = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ Array a -> Int -> IO a
forall a. Storable a => Array a -> Int -> IO a
unsafeIndexIO Array a
arr Int
i in a
r

-- | /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
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
        len :: Int
len = Ptr a
aEnd Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
p
    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

-- | /O(1)/ Get the length of the array i.e. the number of elements in the
-- array.
--
-- @since 0.7.0
{-# INLINE length #-}
length :: forall a. Storable a => Array a -> Int
length :: Array a -> Int
length Array a
arr = Array a -> Int
forall a. Array a -> Int
byteLength Array a
arr Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)

-- | 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
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
        len :: Int
len = Ptr a
aBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
p
    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

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

data GroupState s start end bound
    = GroupStart s
    | GroupBuffer s start end bound
    | GroupYield start end bound (GroupState s 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 (ForeignPtr a) (Ptr a) (Ptr a)
 -> m (Step
         (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)))
-> GroupState s (ForeignPtr 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 (ForeignPtr a) (Ptr a) (Ptr a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a.
State Stream m a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
step' (s -> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound. s -> GroupState s start end bound
GroupStart s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> m (Step (GroupState s (ForeignPtr 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.fromStreamDArraysOf: the size of "
                 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"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 ForeignPtr 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 a. Storable a => Int -> IO (Array a)
newArray Int
n
        Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)))
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (s
-> ForeignPtr a
-> Ptr a
-> Ptr a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound.
s -> start -> end -> bound -> GroupState s start end bound
GroupBuffer s
st ForeignPtr a
start Ptr a
end Ptr a
bound)

    step' State Stream m a
gst (GroupBuffer s
st ForeignPtr 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
                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 (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)))
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$
                    if Ptr a
forall b. Ptr b
end' Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
bound
                    then GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (ForeignPtr a
-> Ptr a
-> Ptr a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound.
start
-> end
-> bound
-> GroupState s start end bound
-> GroupState s start end bound
GroupYield ForeignPtr a
start Ptr a
forall b. Ptr b
end' Ptr a
bound (s -> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound. s -> GroupState s start end bound
GroupStart s
s))
                    else GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (s
-> ForeignPtr a
-> Ptr a
-> Ptr a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound.
s -> start -> end -> bound -> GroupState s start end bound
GroupBuffer s
s ForeignPtr a
start Ptr a
forall b. Ptr b
end' Ptr a
bound)
            D.Skip s
s -> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)))
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (s
-> ForeignPtr a
-> Ptr a
-> Ptr a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound.
s -> start -> end -> bound -> GroupState s start end bound
GroupBuffer s
s ForeignPtr a
start Ptr a
end Ptr a
bound)
            Step s a
D.Stop -> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)))
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
forall s a. s -> Step s a
D.Skip (ForeignPtr a
-> Ptr a
-> Ptr a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound.
start
-> end
-> bound
-> GroupState s start end bound
-> GroupState s start end bound
GroupYield ForeignPtr a
start Ptr a
end Ptr a
bound GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
forall s start end bound. GroupState s start end bound
GroupFinish)

    step' State Stream m a
_ (GroupYield ForeignPtr a
start Ptr a
end Ptr a
bound GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
next) =
        Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
 -> m (Step
         (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)))
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall a b. (a -> b) -> a -> b
$ Array a
-> GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
-> Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
forall s a. a -> s -> Step s a
D.Yield (ForeignPtr a -> Ptr a -> Ptr a -> Array a
forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array ForeignPtr a
start Ptr a
end Ptr a
bound) GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
next

    step' State Stream m a
_ GroupState s (ForeignPtr a) (Ptr a) (Ptr a)
GroupFinish = Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a)
-> m (Step (GroupState s (ForeignPtr a) (Ptr a) (Ptr a)) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return Step (GroupState s (ForeignPtr 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 bufferChunks #-}
bufferChunks :: (MonadIO m, Storable a) =>
    D.Stream m a -> m (K.Stream m (Array a))
bufferChunks :: Stream m a -> m (Stream m (Array a))
bufferChunks Stream m a
m = (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 (t :: (* -> *) -> * -> *) a (m :: * -> *).
IsStream t =>
a -> t m a -> t m a
K.cons Stream m (Array a)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
IsStream t =>
t m a
K.nil (Stream m (Array a) -> m (Stream m (Array a)))
-> Stream m (Array a) -> m (Stream m (Array a))
forall a b. (a -> b) -> a -> b
$ 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
defaultChunkSize Stream m a
m

-------------------------------------------------------------------------------
-- Streams of arrays - flatten
-------------------------------------------------------------------------------

data ReadUState a = ReadUState
    {-# UNPACK #-} !(ForeignPtr a)  -- foreign ptr with end of array pointer
    {-# UNPACK #-} !(Ptr a)         -- current pointer

-- | Resumable unfold of an array.
--
{-# INLINE_NORMAL producer #-}
producer :: forall m a. (Monad 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.
(Monad m, Storable a) =>
ReadUState a -> m (Step (ReadUState a) a)
step Array a -> m (ReadUState a)
forall (m :: * -> *) a a. Monad m => Array a -> m (ReadUState a)
inject ReadUState a -> m (Array a)
forall (m :: * -> *) a a. Monad m => ReadUState a -> m (Array a)
extract
    where

    inject :: Array a -> m (ReadUState a)
inject (Array (ForeignPtr Addr#
start ForeignPtrContents
contents) (Ptr Addr#
end) Ptr a
_) =
        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
$ ForeignPtr a -> Ptr a -> ReadUState a
forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState (Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
end ForeignPtrContents
contents) (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
start)

    {-# INLINE_LATE step #-}
    step :: ReadUState a -> m (Step (ReadUState a) a)
step (ReadUState fp :: ForeignPtr a
fp@(ForeignPtr Addr#
end ForeignPtrContents
_) Ptr a
p) | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
end =
        let x :: ()
x = IO () -> ()
forall a. IO a -> a
unsafeInlineIO (IO () -> ()) -> IO () -> ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
fp
        in ()
x () -> m (Step (ReadUState a) a) -> m (Step (ReadUState a) a)
`seq` 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 ForeignPtr a
fp Ptr a
p) = do
            -- 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.
            let !x :: a
x = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
            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
                (ForeignPtr a -> Ptr a -> ReadUState a
forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState ForeignPtr a
fp (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)))

    extract :: ReadUState a -> m (Array a)
extract (ReadUState (ForeignPtr Addr#
end ForeignPtrContents
contents) (Ptr Addr#
p)) =
        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
$ ForeignPtr a -> Ptr a -> Ptr a -> Array a
forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array (Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
p ForeignPtrContents
contents) (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
end) (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
end)

-- | Unfold an array into a stream.
--
-- @since 0.7.0
{-# INLINE_NORMAL read #-}
read :: forall m a. (Monad 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.
(Monad 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. (Monad 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.
(Monad 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 ForeignPtr a
fp 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
$ ForeignPtr a -> Ptr a -> ReadUState a
forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState ForeignPtr a
fp Ptr a
forall b. Ptr b
p

    {-# INLINE_LATE step #-}
    step :: ReadUState a -> m (Step (ReadUState a) a)
step (ReadUState fp :: ForeignPtr a
fp@(ForeignPtr Addr#
start ForeignPtrContents
_) Ptr a
p) | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
start =
        let x :: ()
x = IO () -> ()
forall a. IO a -> a
unsafeInlineIO (IO () -> ()) -> IO () -> ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
fp
        in ()
x () -> m (Step (ReadUState a) a) -> m (Step (ReadUState a) a)
`seq` 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 ForeignPtr a
fp Ptr a
p) = do
            -- 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.
            let !x :: a
x = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
            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
                (ForeignPtr a -> Ptr a -> ReadUState a
forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState ForeignPtr a
fp (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))))

data FlattenState s a =
      OuterLoop s
    | InnerLoop s !(ForeignPtr a) !(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 a -> m (Step (FlattenState s a) a))
-> FlattenState s 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 a -> m (Step (FlattenState s a) a)
forall (m :: * -> *) a.
State Stream m a
-> FlattenState s a -> m (Step (FlattenState s a) a)
step' (s -> FlattenState s a
forall s a. s -> FlattenState s a
OuterLoop s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> FlattenState s a -> m (Step (FlattenState s 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 a) a -> m (Step (FlattenState s a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s a) a -> m (Step (FlattenState s a) a))
-> Step (FlattenState s a) a -> m (Step (FlattenState s a) a)
forall a b. (a -> b) -> a -> b
$ case Step s (Array a)
r of
            D.Yield Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} s
s ->
                let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
                in FlattenState s a -> Step (FlattenState s a) a
forall s a. s -> Step s a
D.Skip (s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
forall s a. s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
InnerLoop s
s ForeignPtr a
aStart Ptr a
p Ptr a
aEnd)
            D.Skip s
s -> FlattenState s a -> Step (FlattenState s a) a
forall s a. s -> Step s a
D.Skip (s -> FlattenState s a
forall s a. s -> FlattenState s a
OuterLoop s
s)
            Step s (Array a)
D.Stop -> Step (FlattenState s a) a
forall s a. Step s a
D.Stop

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

    step' State Stream m a
_ (InnerLoop s
st ForeignPtr a
startf 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
                    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
startf
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        Step (FlattenState s a) a -> m (Step (FlattenState s a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s a) a -> m (Step (FlattenState s a) a))
-> Step (FlattenState s a) a -> m (Step (FlattenState s a) a)
forall a b. (a -> b) -> a -> b
$ a -> FlattenState s a -> Step (FlattenState s a) a
forall s a. a -> s -> Step s a
D.Yield a
x (s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
forall s a. s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
InnerLoop s
st ForeignPtr a
startf
                            (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 a -> m (Step (FlattenState s a) a))
-> FlattenState s 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 a -> m (Step (FlattenState s a) a)
forall (m :: * -> *) a.
State Stream m a
-> FlattenState s a -> m (Step (FlattenState s a) a)
step' (s -> FlattenState s a
forall s a. s -> FlattenState s a
OuterLoop s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> FlattenState s a -> m (Step (FlattenState s 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 a) a -> m (Step (FlattenState s a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s a) a -> m (Step (FlattenState s a) a))
-> Step (FlattenState s a) a -> m (Step (FlattenState s a) a)
forall a b. (a -> b) -> a -> b
$ case Step s (Array a)
r of
            D.Yield Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} 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))
                -- XXX we do not need aEnd
                in FlattenState s a -> Step (FlattenState s a) a
forall s a. s -> Step s a
D.Skip (s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
forall s a. s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
InnerLoop s
s ForeignPtr a
aStart Ptr a
forall b. Ptr b
p Ptr a
aEnd)
            D.Skip s
s -> FlattenState s a -> Step (FlattenState s a) a
forall s a. s -> Step s a
D.Skip (s -> FlattenState s a
forall s a. s -> FlattenState s a
OuterLoop s
s)
            Step s (Array a)
D.Stop -> Step (FlattenState s a) a
forall s a. Step s a
D.Stop

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

    step' State Stream m a
_ (InnerLoop s
st ForeignPtr a
startf 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
                    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
startf
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        Step (FlattenState s a) a -> m (Step (FlattenState s a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FlattenState s a) a -> m (Step (FlattenState s a) a))
-> Step (FlattenState s a) a -> m (Step (FlattenState s a) a)
forall a b. (a -> b) -> a -> b
$ a -> FlattenState s a -> Step (FlattenState s a) a
forall s a. a -> s -> Step s a
D.Yield a
x (s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
forall s a. s -> ForeignPtr a -> Ptr a -> Ptr a -> FlattenState s a
InnerLoop s
st ForeignPtr a
startf
                            (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))) Ptr a
end)

-------------------------------------------------------------------------------
-- 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 :: (a -> b -> b) -> b -> Array a -> b
toListFB a -> b -> b
c b
n Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} = Ptr a -> b
go (ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart)
    where

    go :: Ptr a -> b
go Ptr a
p | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd = b
n
    go Ptr a
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.
        let !x :: a
x = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> 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
                    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in a -> b -> b
c a
x (Ptr a -> b
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)))

-- | Convert an 'Array' into a list.
--
-- @since 0.7.0
{-# INLINE toList #-}
toList :: Storable a => Array a -> [a]
toList :: Array a -> [a]
toList Array a
s = (forall b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> (a -> b -> b) -> b -> Array a -> b
forall a b. Storable a => (a -> b -> b) -> b -> Array a -> b
toListFB a -> b -> b
c b
n Array a
s)

-- | 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. (Monad m, Storable a) => Array a -> D.Stream m a
toStreamD :: Array a -> Stream m a
toStreamD Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
    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.
Monad m =>
p -> Ptr a -> m (Step (Ptr b) a)
step 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. 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 = do
        -- 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.
        let !x :: a
x = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> 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
                    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        Step (Ptr b) a -> m (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Ptr b) a -> m (Step (Ptr b) a))
-> Step (Ptr b) a -> m (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
x (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 t m a. (K.IsStream t, Storable a) => Array a -> t m a
toStreamK :: Array a -> t m a
toStreamK Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
    in Ptr a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
IsStream t =>
Ptr a -> t m a
go Ptr a
p

    where

    go :: Ptr a -> t m a
go Ptr a
p | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd = t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
IsStream t =>
t m a
K.nil
         | Bool
otherwise =
        -- See Note in toStreamD.
        let !x :: a
x = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> 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
                    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in a
x a -> t m a -> t m a
forall (t :: (* -> *) -> * -> *) a (m :: * -> *).
IsStream t =>
a -> t m a -> t m a
`K.cons` Ptr a -> t 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. (Monad m, Storable a) => Array a -> D.Stream m a
toStreamDRev :: Array a -> Stream m a
toStreamDRev Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    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.
Monad m =>
p -> Ptr a -> m (Step (Ptr b) a)
step Ptr a
forall b. Ptr b
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
< ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart = 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 = do
        -- See comments in toStreamD for why we use unsafeInlineIO
        let !x :: a
x = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> 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
                    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        Step (Ptr b) a -> m (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Ptr b) a -> m (Step (Ptr b) a))
-> Step (Ptr b) a -> m (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
x (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 t m a. (K.IsStream t, Storable a) => Array a -> t m a
toStreamKRev :: Array a -> t m a
toStreamKRev Array {Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    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 -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
IsStream t =>
Ptr a -> t m a
go Ptr a
forall b. Ptr b
p

    where

    go :: Ptr a -> t m a
go Ptr a
p | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart = t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
IsStream t =>
t m a
K.nil
         | Bool
otherwise =
        let !x :: a
x = IO a -> a
forall a. IO a -> a
unsafeInlineIO (IO a -> a) -> IO a -> 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
                    ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
aStart
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in a
x a -> t m a -> t m a
forall (t :: (* -> *) -> * -> *) a (m :: * -> *).
IsStream t =>
a -> t m a -> t m a
`K.cons` Ptr a -> t 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
-------------------------------------------------------------------------------

-- | Strict left fold of an array.
{-# INLINE_NORMAL foldl' #-}
foldl' :: forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
foldl' :: (b -> a -> b) -> b -> Array a -> b
foldl' b -> a -> b
f b
z Array a
arr = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> Identity b -> b
forall a b. (a -> b) -> a -> b
$ (b -> a -> b) -> b -> Stream Identity a -> Identity 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 Identity a -> Identity b)
-> Stream Identity a -> Identity b
forall a b. (a -> b) -> a -> b
$ Array a -> Stream Identity a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
toStreamD Array a
arr

-- | Right fold of an array.
{-# INLINE_NORMAL foldr #-}
foldr :: Storable a => (a -> b -> b) -> b -> Array a -> b
foldr :: (a -> b -> b) -> b -> Array a -> b
foldr a -> b -> b
f b
z Array a
arr = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> Identity b -> b
forall a b. (a -> b) -> a -> b
$ (a -> b -> b) -> b -> Stream Identity a -> Identity 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 Identity a -> Identity b)
-> Stream Identity a -> Identity b
forall a b. (a -> b) -> a -> b
$ Array a -> Stream Identity a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
toStreamD Array a
arr

-------------------------------------------------------------------------------
-- Write Folds to fold streams into arrays
-------------------------------------------------------------------------------

{-# INLINE_NORMAL writeNAllocWith #-}
writeNAllocWith :: forall m a. (MonadIO m, Storable a)
    => (Int -> IO (Array a)) -> Int -> Fold m a (Array a)
writeNAllocWith :: (Int -> IO (Array a)) -> Int -> Fold m a (Array a)
writeNAllocWith Int -> IO (Array a)
alloc Int
n = (Array a -> a -> m (Step (Array a) (Array a)))
-> m (Step (Array a) (Array a))
-> (Array a -> m (Array a))
-> Fold m a (Array a)
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold Array a -> a -> m (Step (Array a) (Array a))
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> a -> m (Step (Array a) (Array a))
step m (Step (Array a) (Array a))
forall b. m (Step (Array a) b)
initial Array a -> m (Array a)
forall a. a -> m a
extract

    where

    initial :: m (Step (Array a) b)
initial = Array a -> Step (Array a) b
forall s b. s -> Step s b
FL.Partial (Array a -> Step (Array a) b)
-> m (Array a) -> m (Step (Array a) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Array a) -> m (Array a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> IO (Array a)
alloc (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
n Int
0))
    step :: Array a -> a -> m (Step (Array a) (Array a))
step arr :: Array a
arr@(Array ForeignPtr a
_ Ptr a
end Ptr a
bound) a
_ | Ptr a
end Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
bound = Step (Array a) (Array a) -> m (Step (Array a) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Array a) (Array a) -> m (Step (Array a) (Array a)))
-> Step (Array a) (Array a) -> m (Step (Array a) (Array a))
forall a b. (a -> b) -> a -> b
$ Array a -> Step (Array a) (Array a)
forall s b. b -> Step s b
FL.Done Array a
arr
    step (Array ForeignPtr 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
        Step (Array a) (Array a) -> m (Step (Array a) (Array a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Array a) (Array a) -> m (Step (Array a) (Array a)))
-> Step (Array a) (Array a) -> m (Step (Array a) (Array a))
forall a b. (a -> b) -> a -> b
$ Array a -> Step (Array a) (Array a)
forall s b. s -> Step s b
FL.Partial (Array a -> Step (Array a) (Array a))
-> Array a -> Step (Array a) (Array a)
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> Ptr a -> Ptr a -> Array a
forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array ForeignPtr 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)) Ptr a
bound
    -- XXX note that shirkToFit does not maintain alignment, in case we are
    -- using aligned allocation.
    extract :: a -> m a
extract = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return -- liftIO . shrinkToFit

-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
-- 'Array'.
--
-- @writeN n = Fold.take n writeNUnsafe@
--
-- @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 -> IO (Array a)) -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> IO (Array a)) -> Int -> Fold m a (Array a)
writeNAllocWith Int -> IO (Array a)
forall a. Storable a => Int -> IO (Array a)
newArray

-- | @writeNAligned alignment n@ folds a maximum of @n@ elements from the input
-- stream to an 'Array' aligned to the given size.
--
-- /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
alignSize = (Int -> IO (Array a)) -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> IO (Array a)) -> Int -> Fold m a (Array a)
writeNAllocWith (Int -> Int -> IO (Array a)
forall a. Storable a => Int -> Int -> IO (Array a)
newArrayAligned Int
alignSize)

-- | @writeNAlignedUnmanaged n@ folds a maximum of @n@ elements from the input
-- stream to an 'Array' aligned to the given size and using unmanaged memory.
-- This could be useful to allocate memory that we need to allocate only once
-- in the lifetime of the program.
--
-- /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
alignSize =
    (Int -> IO (Array a)) -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
(Int -> IO (Array a)) -> Int -> Fold m a (Array a)
writeNAllocWith (Int -> Int -> IO (Array a)
forall a. Storable a => Int -> Int -> IO (Array a)
newArrayAlignedUnmanaged Int
alignSize)

data ArrayUnsafe a = ArrayUnsafe
    {-# UNPACK #-} !(ForeignPtr a) -- first address
    {-# UNPACK #-} !(Ptr a)        -- first unused address

-- | 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.
--
-- @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
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 ArrayUnsafe a -> m (Array a)
forall (m :: * -> *) a. Monad m => ArrayUnsafe a -> m (Array a)
extract

    where

    initial :: m (Step (ArrayUnsafe a) b)
initial = do
        (Array ForeignPtr a
start Ptr a
end Ptr a
_) <- 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 a. Storable a => Int -> IO (Array a)
newArray (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
n Int
0)
        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
$ ForeignPtr a -> Ptr a -> ArrayUnsafe a
forall a. ForeignPtr a -> Ptr a -> ArrayUnsafe a
ArrayUnsafe ForeignPtr a
start Ptr a
end

    step :: ArrayUnsafe a -> a -> m (Step (ArrayUnsafe a) b)
step (ArrayUnsafe ForeignPtr 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
        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
$ ForeignPtr a -> Ptr a -> ArrayUnsafe a
forall a. ForeignPtr a -> Ptr a -> ArrayUnsafe a
ArrayUnsafe ForeignPtr 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))

    extract :: ArrayUnsafe a -> m (Array a)
extract (ArrayUnsafe ForeignPtr a
start Ptr a
end) = 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
$ ForeignPtr a -> Ptr a -> Ptr a -> Array a
forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array ForeignPtr a
start Ptr a
end Ptr a
end -- liftIO . shrinkToFit

-- XXX Buffer to a list instead?
--
-- | Buffer a stream into a stream of arrays.
--
-- @writeChunks = Fold.many Fold.toStream (Array.writeN n)@
--
-- See 'bufferChunks'.
--
-- /Unimplemented/
--
{-# INLINE_NORMAL writeChunks #-}
writeChunks :: -- (MonadIO m, Storable a) =>
    Int -> Fold m a (D.Stream m (Array a))
writeChunks :: Int -> Fold m a (Stream m (Array a))
writeChunks = Int -> Fold m a (Stream m (Array a))
forall a. HasCallStack => a
undefined -- Fold.many Fold.toStream (Array.writeN n)

-- XXX Compare toArrayMinChunk with fromStreamD which uses an array of streams
-- implementation. We can write this using writeChunks above if that is faster.
-- If toArrayMinChunk 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 seems 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.

{-# INLINE_NORMAL toArrayMinChunk #-}
toArrayMinChunk :: forall m a. (MonadIO m, Storable a)
    => Int -> Int -> Fold m a (Array a)
-- toArrayMinChunk n = FL.rmapM spliceArrays $ toArraysOf n
toArrayMinChunk :: Int -> Int -> Fold m a (Array a)
toArrayMinChunk Int
alignSize 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 ForeignPtr 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
        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
$ ForeignPtr a -> Ptr a -> Ptr a -> Array a
forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array ForeignPtr 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)) 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]
"toArrayMinChunk: 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 a. Storable a => Int -> Int -> IO (Array a)
newArrayAligned Int
alignSize Int
elemCount
    step :: Array a -> a -> m (Array a)
step arr :: Array a
arr@(Array ForeignPtr a
start Ptr a
end Ptr a
bound) a
x | Ptr a
end Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
bound = do
        let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
start
            oldSize :: Int
oldSize = Ptr a
end Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
p
            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 -> Array a -> IO (Array a)
forall a. Int -> Int -> Array a -> IO (Array a)
reallocAligned Int
alignSize 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 a. Storable a => Array a -> IO (Array a)
shrinkToFit

-- | Fold the whole input to a single array.
--
-- /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 -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
toArrayMinChunk (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a))
                        (a -> Int -> Int
forall a. Storable a => a -> Int -> Int
bytesToElemCount (a
forall a. HasCallStack => a
undefined :: a)
                        (Int -> Int
mkChunkSize Int
1024))

-- | Like 'write' but the array memory is aligned according to the specified
-- alignment size. This could be useful when we have specific alignment, for
-- example, cache aligned arrays for lookup table etc.
--
-- /Caution! Do not use this on infinite streams./
--
-- @since 0.7.0
{-# INLINE writeAligned #-}
writeAligned :: forall m a. (MonadIO m, Storable a)
    => Int -> Fold m a (Array a)
writeAligned :: Int -> Fold m a (Array a)
writeAligned Int
alignSize =
    Int -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
toArrayMinChunk Int
alignSize
                    (a -> Int -> Int
forall a. Storable a => a -> Int -> Int
bytesToElemCount (a
forall a. HasCallStack => a
undefined :: a)
                    (Int -> Int
mkChunkSize Int
1024))

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

-- | Use the 'writeN' fold instead.
--
-- @fromStreamDN n = D.fold (writeN n)@
--
{-# INLINE_NORMAL fromStreamDN #-}
fromStreamDN :: forall m a. (MonadIO m, Storable a)
    => Int -> D.Stream m a -> m (Array a)
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 a. Storable a => Int -> IO (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. Monad 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 :: Storable a => Int -> [a] -> Array a
fromListN :: Int -> [a] -> Array a
fromListN Int
n [a]
xs = IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ Int -> Stream IO a -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> m (Array a)
fromStreamDN Int
n (Stream IO a -> IO (Array a)) -> Stream IO a -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ [a] -> Stream IO a
forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

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

-- 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 = do
    Stream m (Array a)
buffered <- Stream m a -> m (Stream m (Array a))
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m a -> m (Stream m (Array a))
bufferChunks Stream m a
m
    Int
len <- (Int -> Int -> Int) -> Int -> Stream m Int -> m Int
forall (t :: (* -> *) -> * -> *) (m :: * -> *) b a.
(IsStream t, Monad m) =>
(b -> a -> b) -> b -> t 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 (t :: (* -> *) -> * -> *) a b (m :: * -> *).
IsStream t =>
(a -> b) -> t m a -> t m b
K.map Array a -> Int
forall a. Storable a => Array a -> Int
length Stream m (Array a)
buffered)
    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.
(Monad 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. Monad m => Stream m a -> Stream m a
D.fromStreamK Stream m (Array a)
buffered

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

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

-- | Copy two arrays into a newly allocated array.
{-# INLINE spliceTwo #-}
spliceTwo :: (MonadIO m, Storable a) => Array a -> Array a -> m (Array a)
spliceTwo :: Array a -> Array a -> m (Array a)
spliceTwo Array a
arr1 Array a
arr2 = do
    let src1 :: Ptr a
src1 = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
arr1)
        src2 :: Ptr a
src2 = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart 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 a. Storable a => Int -> IO (Array a)
newArray (Int
len1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len2)
    let dst :: Ptr a
dst = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart 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
        ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart 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
        ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart 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) }

-- | Splice an array into a pre-reserved mutable array.  The user must ensure
-- that there is enough space in the mutable array, otherwise the splicing
-- fails.
{-# INLINE spliceWith #-}
spliceWith :: (MonadIO m) => Array a -> Array a -> m (Array a)
spliceWith :: Array a -> Array a -> m (Array a)
spliceWith dst :: Array a
dst@(Array ForeignPtr a
_ Ptr a
end Ptr a
bound) Array a
src =
    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 srcLen :: Int
srcLen = Array a -> Int
forall a. Array a -> Int
byteLength Array a
src
        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 [Char] -> IO (Array a)
forall a. HasCallStack => [Char] -> a
error
                 [Char]
"Bug: spliceWith: Not enough space in the target array"
        else ForeignPtr a -> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
dst) ((Ptr a -> IO (Array a)) -> IO (Array a))
-> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
_ ->
                ForeignPtr a -> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
src) ((Ptr a -> IO (Array a)) -> IO (Array a))
-> (Ptr a -> IO (Array a)) -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
psrc -> do
                     let pdst :: Ptr a
pdst = Array a -> Ptr a
forall a. Array a -> Ptr a
aEnd Array a
dst
                     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
                     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}

-- | Splice a new array into a preallocated mutable array, doubling the space
-- if there is no space in the target array.
{-# INLINE spliceWithDoubling #-}
spliceWithDoubling :: (MonadIO m, Storable a)
    => Array a -> Array a -> m (Array a)
spliceWithDoubling :: Array a -> Array a -> m (Array a)
spliceWithDoubling dst :: Array a
dst@(Array ForeignPtr a
start Ptr a
end Ptr a
bound) Array a
src  = do
    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` ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart 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 oldStart :: Ptr a
oldStart = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
start
                oldSize :: Int
oldSize = Ptr a
end Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
oldStart
                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
oldSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 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
$ Int -> Array a -> IO (Array a)
forall a. Storable a => Int -> Array a -> IO (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 -> m (Array a)
forall (m :: * -> *) a.
MonadIO m =>
Array a -> Array a -> m (Array a)
spliceWith Array a
dst1 Array a
src

-------------------------------------------------------------------------------
-- 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
ForeignPtr Word8
aBound :: Ptr Word8
aEnd :: Ptr Word8
aStart :: ForeignPtr Word8
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} = 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 = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
aStart
    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 b. Ptr b
nullPtr
        then (Array Word8
arr, Maybe (Array Word8)
forall a. Maybe a
Nothing)
        else
            ( Array :: forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array
                { aStart :: ForeignPtr Word8
aStart = ForeignPtr Word8
aStart
                , 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. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array
                    { aStart :: ForeignPtr Word8
aStart = ForeignPtr Word8
aStart ForeignPtr Word8 -> Int -> ForeignPtr Word8
forall a b. ForeignPtr a -> Int -> ForeignPtr b
`plusForeignPtr` (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
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    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 = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off
                in ( Array :: forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array
                  { aStart :: ForeignPtr a
aStart = ForeignPtr a
aStart
                  , aEnd :: Ptr a
aEnd = Ptr a
forall b. Ptr b
p
                  , aBound :: Ptr a
aBound = Ptr a
forall b. Ptr b
p
                  }
                , Array :: forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array
                  { aStart :: ForeignPtr a
aStart = ForeignPtr a
aStart ForeignPtr a -> Int -> ForeignPtr a
forall a b. ForeignPtr a -> Int -> ForeignPtr b
`plusForeignPtr` Int
off
                  , aEnd :: Ptr a
aEnd = Ptr a
aEnd
                  , aBound :: Ptr a
aBound = Ptr a
aBound
                  }
                )

-------------------------------------------------------------------------------
-- Instances
-------------------------------------------------------------------------------

instance (Show a, Storable a) => Show (Array a) where
    {-# INLINE showsPrec #-}
    showsPrec :: Int -> Array a -> [Char] -> [Char]
showsPrec Int
_ = [a] -> [Char] -> [Char]
forall a. Show a => a -> [Char] -> [Char]
shows ([a] -> [Char] -> [Char])
-> (Array a -> [a]) -> Array a -> [Char] -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> [a]
forall a. Storable a => Array a -> [a]
toList

instance (Storable a, Read a, Show a) => Read (Array a) where
    {-# INLINE readPrec #-}
    readPrec :: ReadPrec (Array a)
readPrec = do
          [a]
xs <- ReadPrec [a]
forall a. Read a => ReadPrec a
readPrec
          Array a -> ReadPrec (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> Array a
forall a. Storable a => [a] -> Array a
fromList [a]
xs)
    readListPrec :: ReadPrec [Array a]
readListPrec = ReadPrec [Array a]
forall a. Read a => ReadPrec [a]
readListPrecDefault

instance (a ~ Char) => IsString (Array a) where
    {-# INLINE fromString #-}
    fromString :: [Char] -> Array a
fromString = [Char] -> Array a
forall a. Storable a => [a] -> Array a
fromList

-- GHC versions 8.0 and below cannot derive IsList
instance Storable a => IsList (Array a) where
    type (Item (Array a)) = a
    {-# INLINE fromList #-}
    fromList :: [Item (Array a)] -> Array a
fromList = [Item (Array a)] -> Array a
forall a. Storable a => [a] -> Array a
fromList
    {-# INLINE fromListN #-}
    fromListN :: Int -> [Item (Array a)] -> Array a
fromListN = Int -> [Item (Array a)] -> Array a
forall a. Storable a => Int -> [a] -> Array a
fromListN
    {-# INLINE toList #-}
    toList :: Array a -> [Item (Array a)]
toList = Array a -> [Item (Array a)]
forall a. Storable a => Array a -> [a]
toList

{-# INLINE arrcmp #-}
arrcmp :: Array a -> Array a -> Bool
arrcmp :: Array a -> Array a -> Bool
arrcmp Array a
arr1 Array a
arr2 =
    let !res :: Bool
res = IO Bool -> Bool
forall a. IO a -> a
unsafeInlineIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
            let ptr1 :: Ptr a
ptr1 = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (ForeignPtr a -> Ptr a) -> ForeignPtr a -> Ptr a
forall a b. (a -> b) -> a -> b
$ Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
arr1
            let ptr2 :: Ptr a
ptr2 = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (ForeignPtr a -> Ptr a) -> ForeignPtr a -> Ptr a
forall a b. (a -> b) -> a -> b
$ Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart 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 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
                ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr (ForeignPtr a -> IO ()) -> ForeignPtr a -> IO ()
forall a b. (a -> b) -> a -> b
$ Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
arr1
                ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr (ForeignPtr a -> IO ()) -> ForeignPtr a -> IO ()
forall a b. (a -> b) -> a -> b
$ Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart 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
    in Bool
res

-- XXX we are assuming that Storable equality means element equality. This may
-- or may not be correct? arrcmp is 40% faster compared to stream equality.
instance (Storable a, Eq a) => Eq (Array a) where
    {-# INLINE (==) #-}
    == :: Array a -> Array a -> Bool
(==) = Array a -> Array a -> Bool
forall a. Array a -> Array a -> Bool
arrcmp
    -- arr1 == arr2 = runIdentity $ D.eqBy (==) (toStreamD arr1) (toStreamD arr2)

instance (Storable a, NFData a) => NFData (Array a) where
    {-# INLINE rnf #-}
    rnf :: Array a -> ()
rnf = (() -> a -> ()) -> () -> Array a -> ()
forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
foldl' (\()
_ a
x -> a -> ()
forall a. NFData a => a -> ()
rnf a
x) ()

instance (Storable a, Ord a) => Ord (Array a) where
    {-# INLINE compare #-}
    compare :: Array a -> Array a -> Ordering
compare Array a
arr1 Array a
arr2 = IO Ordering -> Ordering
forall a. IO a -> a
unsafePerformIO (IO Ordering -> Ordering) -> IO Ordering -> Ordering
forall a b. (a -> b) -> a -> b
$
        (a -> a -> Ordering) -> Stream IO a -> Stream IO a -> IO Ordering
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
D.cmpBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Array a -> Stream IO a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
toStreamD Array a
arr1) (Array a -> Stream IO a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
toStreamD Array a
arr2)

    -- Default definitions defined in base do not have an INLINE on them, so we
    -- replicate them here with an INLINE.
    {-# INLINE (<) #-}
    Array a
x < :: Array a -> Array a -> Bool
<  Array a
y = case Array a -> Array a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
LT -> Bool
True;  Ordering
_ -> Bool
False }

    {-# INLINE (<=) #-}
    Array a
x <= :: Array a -> Array a -> Bool
<= Array a
y = case Array a -> Array a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
GT -> Bool
False; Ordering
_ -> Bool
True }

    {-# INLINE (>) #-}
    Array a
x > :: Array a -> Array a -> Bool
>  Array a
y = case Array a -> Array a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
GT -> Bool
True;  Ordering
_ -> Bool
False }

    {-# INLINE (>=) #-}
    Array a
x >= :: Array a -> Array a -> Bool
>= Array a
y = case Array a -> Array a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
LT -> Bool
False; Ordering
_ -> Bool
True }

    -- These two default methods use '<=' rather than 'compare'
    -- because the latter is often more expensive
    {-# INLINE max #-}
    max :: Array a -> Array a -> Array a
max Array a
x Array a
y = if Array a
x Array a -> Array a -> Bool
forall a. Ord a => a -> a -> Bool
<= Array a
y then Array a
y else Array a
x

    {-# INLINE min #-}
    min :: Array a -> Array a -> Array a
min Array a
x Array a
y = if Array a
x Array a -> Array a -> Bool
forall a. Ord a => a -> a -> Bool
<= Array a
y then Array a
x else Array a
y

#ifdef DEVBUILD
-- Definitions using the Storable constraint from the Array type. These are to
-- make the Foldable instance possible though it is much slower (7x slower).
--
{-# INLINE_NORMAL toStreamD_ #-}
toStreamD_ :: forall m a. MonadIO m => Int -> Array a -> D.Stream m a
toStreamD_ size Array{..} =
    let p = unsafeForeignPtrToPtr aStart
    in D.Stream step p

    where

    {-# INLINE_LATE step #-}
    step _ p | p == aEnd = return D.Stop
    step _ p = do
        x <- liftIO $ do
                    r <- peek p
                    touchForeignPtr aStart
                    return r
        return $ D.Yield x (p `plusPtr` size)

{-# INLINE_NORMAL _foldr #-}
_foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
_foldr f z arr@Array {..} =
    let !n = sizeOf (undefined :: a)
    in unsafePerformIO $ D.foldr f z $ toStreamD_ n arr

-- | Note that the 'Foldable' instance is 7x slower than the direct
-- operations.
instance Foldable Array where
  foldr = _foldr
#endif

-------------------------------------------------------------------------------
-- Semigroup and Monoid
-------------------------------------------------------------------------------

-- Note: we cannot splice the second array into the free space of the first
-- array because that would require a monadic API due to mutation.
-- | Copies the two arrays into a newly allocated array.
instance Storable a => Semigroup (Array a) where
    {-# INLINE (<>) #-}
    Array a
arr1 <> :: Array a -> Array a -> Array a
<> Array a
arr2 = IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ Array a -> Array a -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Array a -> m (Array a)
spliceTwo Array a
arr1 Array a
arr2

nullForeignPtr :: ForeignPtr a
nullForeignPtr :: ForeignPtr a
nullForeignPtr = Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
nullAddr# ([Char] -> ForeignPtrContents
forall a. HasCallStack => [Char] -> a
error [Char]
"nullForeignPtr")

nil ::
#ifdef DEVBUILD
    Storable a =>
#endif
    Array a
nil :: Array a
nil = ForeignPtr a -> Ptr a -> Ptr a -> Array a
forall a. ForeignPtr a -> Ptr a -> Ptr a -> Array a
Array ForeignPtr a
forall a. ForeignPtr a
nullForeignPtr (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
nullAddr#) (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
nullAddr#)

instance Storable a => Monoid (Array a) where
    {-# INLINE mempty #-}
    mempty :: Array a
mempty = Array a
forall a. Array a
nil
    {-# INLINE mappend #-}
    mappend :: Array a -> Array a -> Array a
mappend = Array a -> Array a -> Array a
forall a. Semigroup a => a -> a -> a
(<>)