#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Array.Foreign.Type
-- Copyright   : (c) 2020 Composewell Technologies
--
-- License     : BSD3-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- See notes in "Streamly.Internal.Data.Array.Foreign.Mut.Type"
--
module Streamly.Internal.Data.Array.Foreign.Type
    (
    -- $arrayNotes
      Array (..)

    -- * Freezing and Thawing
    , unsafeFreeze
    , unsafeFreezeWithShrink
    , unsafeThaw

    -- * Construction
    , splice

    , fromPtr
    , fromAddr#
    , fromCString#
    , fromList
    , fromListN
    , fromListRev
    , fromListRevN
    , fromStreamDN
    , fromStreamD

    -- * Split
    , breakOn

    -- * Elimination
    , unsafeIndexIO
    , unsafeIndex
    , byteLength
    , length

    , foldl'
    , foldr
    , splitAt

    , readRev
    , toStreamD
    , toStreamDRev
    , toStreamK
    , toStreamKRev
    , toStream
    , toStreamRev
    , toList

    -- * Folds
    , writeWith
    , writeN
    , writeNUnsafe
    , MA.ArrayUnsafe (..)
    , writeNAligned
    , writeNAlignedUnmanaged
    , write

    -- * Streams of arrays
    , arraysOf
    , bufferChunks
    , flattenArrays
    , flattenArraysRev
    )
where

import Control.Exception (assert)
import Control.DeepSeq (NFData(..))
import Control.Monad.IO.Class (MonadIO(..))
import Data.Functor.Identity (Identity(..))
import Data.Word (Word8)
import Foreign.C.String (CString)
import Foreign.C.Types (CSize(..))
import Foreign.Ptr (plusPtr, castPtr)
import Foreign.Storable (Storable(..))
import GHC.Base (Addr#, nullAddr#, build)
import GHC.Exts (IsList, IsString(..))

import GHC.IO (unsafePerformIO)
import GHC.Ptr (Ptr(..))
import Streamly.Internal.Data.Array.Foreign.Mut.Type
    (ArrayContents, ReadUState(..), touch)
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Stream.Serial (SerialT(..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Text.Read (readPrec, readListPrec, readListPrecDefault)

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

import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
import qualified GHC.Exts as Exts

import Streamly.Internal.System.IO (unsafeInlineIO, defaultChunkSize)

#if __GLASGOW_HASKELL__ < 808
import Data.Semigroup (Semigroup(..))
#endif

#ifdef DEVBUILD
import qualified Data.Foldable as F
#endif

--
-- $setup
-- >>> :m
-- >>> :set -XMagicHash
-- >>> import Prelude hiding (length, foldr, read, unlines, splitAt)
-- >>> import Streamly.Internal.Data.Array.Foreign as Array

-------------------------------------------------------------------------------
-- 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 -> ArrayContents
arrContents ::
#ifndef USE_FOREIGN_PTR
        {-# UNPACK #-}
#endif
            !ArrayContents -- ^ first address
    , Array a -> Ptr a
arrStart :: {-# UNPACK #-} !(Ptr a) -- start address
    , Array a -> Ptr a
aEnd   :: {-# UNPACK #-} !(Ptr a)        -- first unused addres
    }

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

foreign import ccall unsafe "string.h strlen" c_strlen
    :: CString -> IO CSize

-------------------------------------------------------------------------------
-- Freezing and Thawing
-------------------------------------------------------------------------------

-- XXX For debugging we can track slices/references through a weak IORef.  Then
-- trigger a GC after freeze/thaw and assert that there are no references
-- remaining.

-- | Makes an immutable array using the underlying memory of the mutable
-- array.
--
-- Please make sure that there are no other references to the mutable array
-- lying around, so that it is never used after freezing it using
-- /unsafeFreeze/.  If the underlying array is mutated, the immutable promise
-- is lost.
--
-- /Pre-release/
{-# INLINE unsafeFreeze #-}
unsafeFreeze :: MA.Array a -> Array a
unsafeFreeze :: Array a -> Array a
unsafeFreeze (MA.Array ArrayContents
ac Ptr a
as Ptr a
ae Ptr a
_) = ArrayContents -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array ArrayContents
ac Ptr a
as Ptr a
ae

-- | Similar to 'unsafeFreeze' but uses 'MA.rightSize' on the mutable array
-- first.
{-# INLINE unsafeFreezeWithShrink #-}
unsafeFreezeWithShrink :: Storable a => MA.Array a -> Array a
unsafeFreezeWithShrink :: Array a -> Array a
unsafeFreezeWithShrink Array a
arr = 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
$ do
  MA.Array ArrayContents
ac Ptr a
as Ptr a
ae Ptr a
_ <- Array a -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> m (Array a)
MA.rightSize Array a
arr
  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
$ ArrayContents -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array ArrayContents
ac Ptr a
as Ptr a
ae

-- | Makes a mutable array using the underlying memory of the immutable array.
--
-- Please make sure that there are no other references to the immutable array
-- lying around, so that it is never used after thawing it using /unsafeThaw/.
-- If the resulting array is mutated, any references to the older immutable
-- array are mutated as well.
--
-- /Pre-release/
{-# INLINE unsafeThaw #-}
unsafeThaw :: Array a -> MA.Array a
unsafeThaw :: Array a -> Array a
unsafeThaw (Array ArrayContents
ac Ptr a
as Ptr a
ae) = ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Ptr a -> Array a
MA.Array ArrayContents
ac Ptr a
as Ptr a
ae Ptr a
ae

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

-- Splice two immutable arrays creating a new array.
{-# INLINE splice #-}
splice :: (MonadIO m, Storable a) => Array a -> Array a -> m (Array a)
splice :: Array a -> Array a -> m (Array a)
splice Array a
arr1 Array a
arr2 =
    Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Array a -> Array a) -> m (Array a) -> m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array a -> Array a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Array a -> m (Array a)
MA.splice (Array a -> Array a
forall a. Array a -> Array a
unsafeThaw Array a
arr1) (Array a -> Array a
forall a. Array a -> Array a
unsafeThaw Array a
arr2)

-- | Create an 'Array' of the given number of elements of type @a@ from a read
-- only pointer @Ptr a@.  The pointer is not freed when the array is garbage
-- collected. This API is unsafe for the following reasons:
--
-- 1. The pointer must point to static pinned memory or foreign memory that
-- does not require freeing..
-- 2. The pointer must be legally accessible upto the given length.
-- 3. To guarantee that the array is immutable, the contents of the address
-- must be guaranteed to not change.
--
-- /Unsafe/
--
-- /Pre-release/
--
{-# INLINE fromPtr #-}
fromPtr ::
#ifdef DEVBUILD
    Storable a =>
#endif
    Int -> Ptr a -> Array a
fromPtr :: Int -> Ptr a -> Array a
fromPtr Int
n Ptr a
ptr = IO (Array a) -> Array a
forall a. IO a -> a
unsafeInlineIO (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ do
    let end :: Ptr b
end = Ptr a
ptr Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
n
    Array a -> IO (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> IO (Array a)) -> Array a -> IO (Array a)
forall a b. (a -> b) -> a -> b
$ Array :: forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array
        { arrContents :: ArrayContents
arrContents = ArrayContents
MA.nilArrayContents
        , arrStart :: Ptr a
arrStart = Ptr a
ptr
        , aEnd :: Ptr a
aEnd = Ptr a
forall b. Ptr b
end
        }

-- XXX when converting an array of Word8 from a literal string we can simply
-- refer to the literal string. Is it possible to write rules such that
-- fromList Word8 can be rewritten so that GHC does not first convert the
-- literal to [Char] and then we convert it back to an Array Word8?
--
-- TBD: We can also add template haskell quasiquotes to specify arrays of other
-- literal types. TH will encode them into a string literal and we read that as
-- an array of the required type. With template Haskell we can provide a safe
-- version of fromString#.
--
-- | Create an @Array Word8@ of the given length from a static, read only
-- machine address 'Addr#'. See 'fromPtr' for safety caveats.
--
-- A common use case for this API is to create an array from a static unboxed
-- string literal. GHC string literals are of type 'Addr#', and must contain
-- characters that can be encoded in a byte i.e. characters or literal bytes in
-- the range from 0-255.
--
-- >>> import Data.Word (Word8)
-- >>> Array.fromAddr# 5 "hello world!"# :: Array Word8
-- [104,101,108,108,111]
--
-- >>> Array.fromAddr# 3 "\255\NUL\255"# :: Array Word8
-- [255,0,255]
--
-- /See also: 'fromString#'/
--
-- /Unsafe/
--
-- /Time complexity: O(1)/
--
-- /Pre-release/
--
{-# INLINE fromAddr# #-}
fromAddr# ::
#ifdef DEVBUILD
    Storable a =>
#endif
    Int -> Addr# -> Array a
fromAddr# :: Int -> Addr# -> Array a
fromAddr# Int
n Addr#
addr# = Int -> Ptr a -> Array a
forall a. Int -> Ptr a -> Array a
fromPtr Int
n (Ptr Any -> Ptr a
forall a b. Ptr a -> Ptr b
castPtr (Ptr Any -> Ptr a) -> Ptr Any -> Ptr a
forall a b. (a -> b) -> a -> b
$ Addr# -> Ptr Any
forall a. Addr# -> Ptr a
Ptr Addr#
addr#)

-- | Generate a byte array from an 'Addr#' that contains a sequence of NUL
-- (@0@) terminated bytes. The array would not include the NUL byte. The
-- address must be in static read-only memory and must be legally accessible up
-- to and including the first NUL byte.
--
-- An unboxed string literal (e.g. @"hello"#@) is a common example of an
-- 'Addr#' in static read only memory. It represents the UTF8 encoded sequence
-- of bytes terminated by a NUL byte (a 'CString') corresponding to the
-- given unicode string.
--
-- >>> Array.fromCString# "hello world!"#
-- [104,101,108,108,111,32,119,111,114,108,100,33]
--
-- >>> Array.fromCString# "\255\NUL\255"#
-- [255]
--
-- /See also: 'fromAddr#'/
--
-- /Unsafe/
--
-- /Time complexity: O(n) (computes the length of the string)/
--
-- /Pre-release/
--
{-# INLINE fromCString# #-}
fromCString# :: Addr# -> Array Word8
fromCString# :: Addr# -> Array Word8
fromCString# Addr#
addr# = do
    let cstr :: Ptr a
cstr = Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
addr#
        len :: CSize
len = IO CSize -> CSize
forall a. IO a -> a
unsafeInlineIO (IO CSize -> CSize) -> IO CSize -> CSize
forall a b. (a -> b) -> a -> b
$ CString -> IO CSize
c_strlen CString
forall b. Ptr b
cstr
    Int -> Ptr Word8 -> Array Word8
forall a. Int -> Ptr a -> Array a
fromPtr (CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
len) (Ptr Any -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr Any
forall b. Ptr b
cstr)

-- | 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 (Streamly.Memory.Array)/
--
-- @since 0.8.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
$ Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Array a -> Array a) -> IO (Array a) -> IO (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [a] -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> [a] -> m (Array a)
MA.fromListN Int
n [a]
xs

-- XXX We can possibly have a direction flag in the array to reverse it without
-- actually doing anything. With that we can just do "reverse . fromList". But
-- it may complicate all the APIs as all reads of the array will have to handle
-- the flag.
--
-- | Create an 'Array' from the first N elements of a list in reverse order.
-- The array is allocated to size N, if the list terminates before N elements
-- then the array may hold less than N elements.
--
-- /Unimplemented/
{-# INLINABLE fromListRevN #-}
fromListRevN :: {- Storable a => -} Int -> [a] -> Array a
fromListRevN :: Int -> [a] -> Array a
fromListRevN Int
_n [a]
_xs = Array a
forall a. HasCallStack => a
undefined

-- | Create an 'Array' from a list. The list must be of finite size.
--
-- /Since 0.7.0 (Streamly.Memory.Array)/
--
-- @since 0.8.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
$ Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Array a -> Array a) -> IO (Array a) -> IO (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> IO (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
[a] -> m (Array a)
MA.fromList [a]
xs

-- | Create an 'Array' from a list in reverse order. The list must be of finite
-- size.
--
-- /Unimplemented/
{-# INLINABLE fromListRev #-}
fromListRev :: {- Storable a => -} [a] -> Array a
fromListRev :: [a] -> Array a
fromListRev [a]
_xs = Array a
forall a. HasCallStack => a
undefined

{-# 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 = Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Array a -> Array a) -> m (Array a) -> m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Stream m a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> m (Array a)
MA.fromStreamDN Int
limit Stream m a
str

{-# INLINE_NORMAL fromStreamD #-}
fromStreamD :: forall m a. (MonadIO m, Storable a)
    => D.Stream m a -> m (Array a)
fromStreamD :: Stream m a -> m (Array a)
fromStreamD Stream m a
str = Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Array a -> Array a) -> m (Array a) -> m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stream m a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m a -> m (Array a)
MA.fromStreamD Stream m a
str

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

{-# 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 a (m :: * -> *). a -> Stream m a -> Stream m a
K.cons Stream m (Array a)
forall (m :: * -> *) a. Stream m a
K.nil (Stream m (Array a) -> m (Stream m (Array a)))
-> Stream m (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

-- | @arraysOf n stream@ groups the input stream into a stream of
-- arrays of size n.
{-# INLINE_NORMAL arraysOf #-}
arraysOf :: forall m a. (MonadIO m, Storable a)
    => Int -> D.Stream m a -> D.Stream m (Array a)
arraysOf :: Int -> Stream m a -> Stream m (Array a)
arraysOf Int
n Stream m a
str = (Array a -> Array a) -> Stream m (Array a) -> Stream m (Array a)
forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
D.map Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Stream m (Array a) -> Stream m (Array a))
-> Stream m (Array a) -> 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)
MA.arraysOf Int
n Stream m a
str

-- | 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 = Stream m (Array a) -> Stream m a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m (Array a) -> Stream m a
MA.flattenArrays (Stream m (Array a) -> Stream m a)
-> (Stream m (Array a) -> Stream m (Array a))
-> Stream m (Array a)
-> Stream m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array a -> Array a) -> Stream m (Array a) -> Stream m (Array a)
forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
D.map Array a -> Array a
forall a. Array a -> Array a
unsafeThaw

-- | 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 = Stream m (Array a) -> Stream m a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m (Array a) -> Stream m a
MA.flattenArraysRev (Stream m (Array a) -> Stream m a)
-> (Stream m (Array a) -> Stream m (Array a))
-> Stream m (Array a)
-> Stream m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array a -> Array a) -> Stream m (Array a) -> Stream m (Array a)
forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
D.map Array a -> Array a
forall a. Array a -> Array a
unsafeThaw

-- 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 Array Word8
arr = do
  (Array Word8
a, Maybe (Array Word8)
b) <- Word8 -> Array Word8 -> m (Array Word8, Maybe (Array Word8))
forall (m :: * -> *).
MonadIO m =>
Word8 -> Array Word8 -> m (Array Word8, Maybe (Array Word8))
MA.breakOn Word8
sep (Array Word8 -> Array Word8
forall a. Array a -> Array a
unsafeThaw Array Word8
arr)
  (Array Word8, Maybe (Array Word8))
-> m (Array Word8, Maybe (Array Word8))
forall (m :: * -> *) a. Monad m => a -> m a
return (Array Word8 -> Array Word8
forall a. Array a -> Array a
unsafeFreeze Array Word8
a, Array Word8 -> Array Word8
forall a. Array a -> Array a
unsafeFreeze (Array Word8 -> Array Word8)
-> Maybe (Array Word8) -> Maybe (Array Word8)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Array Word8)
b)

-------------------------------------------------------------------------------
-- 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 a
arr = Array a -> Int -> IO a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Int -> m a
MA.getIndexUnsafe (Array a -> Array a
forall a. Array a -> Array a
unsafeThaw Array a
arr)

-- | 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 a -> Int
forall a. Array a -> Int
MA.byteLength (Array a -> Int) -> (Array a -> Array a) -> Array a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array a
forall a. Array a -> Array a
unsafeThaw

-- | /O(1)/ Get the length of the array i.e. the number of elements in the
-- array.
--
-- /Since 0.7.0 (Streamly.Memory.Array)/
--
-- @since 0.8.0
{-# INLINE length #-}
length :: forall a. Storable a => Array a -> Int
length :: Array a -> Int
length Array a
arr = Array a -> Int
forall a. Storable a => Array a -> Int
MA.length (Array a -> Array a
forall a. Array a -> Array a
unsafeThaw Array a
arr)

-- | Unfold an array into a stream in reverse order.
--
-- @since 0.8.0
{-# 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 ArrayContents
contents Ptr a
start Ptr a
end) =
        let p :: Ptr b
p = Ptr a
end Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
         in ReadUState a -> m (ReadUState a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ReadUState a -> m (ReadUState a))
-> ReadUState a -> m (ReadUState a)
forall a b. (a -> b) -> a -> b
$ ArrayContents -> Ptr a -> Ptr a -> ReadUState a
forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
start Ptr a
forall b. Ptr b
p

    {-# INLINE_LATE step #-}
    step :: ReadUState a -> m (Step (ReadUState a) a)
step (ReadUState ArrayContents
contents Ptr a
start Ptr a
p) | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Ptr a
start =
        let x :: ()
x = IO () -> ()
forall a. IO a -> a
unsafeInlineIO (IO () -> ()) -> IO () -> ()
forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
        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 ArrayContents
contents Ptr a
start 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
            let cur :: Ptr b
cur = Ptr a
p Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
            Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ReadUState a) a -> m (Step (ReadUState a) a))
-> Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall a b. (a -> b) -> a -> b
$ a -> ReadUState a -> Step (ReadUState a) a
forall s a. a -> s -> Step s a
D.Yield a
x (ArrayContents -> Ptr a -> Ptr a -> ReadUState a
forall a. ArrayContents -> Ptr a -> Ptr a -> ReadUState a
ReadUState ArrayContents
contents Ptr a
start Ptr a
forall b. Ptr b
cur)


{-# 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
ArrayContents
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = (State Stream m a -> Ptr a -> m (Step (Ptr a) a))
-> Ptr a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a -> Ptr a -> m (Step (Ptr a) a)
forall (m :: * -> *) p b.
Monad m =>
p -> Ptr a -> m (Step (Ptr b) a)
step Ptr a
arrStart

    where

    {-# INLINE_LATE step #-}
    step :: p -> Ptr a -> m (Step (Ptr b) a)
step p
_ Ptr a
p | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = Step (Ptr b) a -> m (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Ptr b) a
forall s a. Step s a
D.Stop
    step p
_ Ptr a
p = 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
                    ArrayContents -> IO ()
touch ArrayContents
arrContents
                    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 m a. Storable a => Array a -> K.Stream m a
toStreamK :: Array a -> Stream m a
toStreamK Array{Ptr a
ArrayContents
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = Ptr a -> Stream m a
forall (m :: * -> *). Ptr a -> Stream m a
go Ptr a
arrStart

    where

    go :: Ptr a -> Stream m a
go Ptr a
p | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = Stream m a
forall (m :: * -> *) a. Stream m a
K.nil
         | Bool
otherwise =
        -- 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
                    ArrayContents -> IO ()
touch ArrayContents
arrContents
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in a
x a -> Stream m a -> Stream m a
forall a (m :: * -> *). a -> Stream m a -> Stream m a
`K.cons` Ptr a -> Stream m a
go (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))

{-# 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
ArrayContents
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} =
    let p :: Ptr b
p = Ptr a
aEnd Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
    in (State Stream m a -> Ptr a -> m (Step (Ptr a) a))
-> Ptr a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a -> Ptr a -> m (Step (Ptr a) a)
forall (m :: * -> *) p b.
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
< Ptr a
arrStart = Step (Ptr b) a -> m (Step (Ptr b) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Ptr b) a
forall s a. Step s a
D.Stop
    step p
_ Ptr a
p = 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
                    ArrayContents -> IO ()
touch ArrayContents
arrContents
                    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 m a. Storable a => Array a -> K.Stream m a
toStreamKRev :: Array a -> Stream m a
toStreamKRev Array {Ptr a
ArrayContents
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} =
    let p :: Ptr b
p = Ptr a
aEnd Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
    in Ptr a -> Stream m a
forall (m :: * -> *). Ptr a -> Stream m a
go Ptr a
forall b. Ptr b
p

    where

    go :: Ptr a -> Stream m a
go Ptr a
p | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Ptr a
arrStart = Stream m a
forall (m :: * -> *) a. Stream m a
K.nil
         | Bool
otherwise =
        let !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
                    ArrayContents -> IO ()
touch ArrayContents
arrContents
                    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
        in a
x a -> Stream m a -> Stream m a
forall a (m :: * -> *). a -> Stream m a -> Stream m a
`K.cons` Ptr a -> Stream m a
go (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)))

-- | Convert an 'Array' into a stream.
--
-- /Pre-release/
{-# INLINE_EARLY toStream #-}
toStream :: (Monad m, Storable a) => Array a -> SerialT m a
toStream :: Array a -> SerialT m a
toStream = Stream m a -> SerialT m a
forall (m :: * -> *) a. Stream m a -> SerialT m a
SerialT (Stream m a -> SerialT m a)
-> (Array a -> Stream m a) -> Array a -> SerialT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream m a -> Stream m a
forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
D.toStreamK (Stream m a -> Stream m a)
-> (Array a -> Stream m a) -> Array a -> Stream m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
toStreamD
-- XXX add fallback to StreamK rule
-- {-# RULES "Streamly.Array.read fallback to StreamK" [1]
--     forall a. S.readK (read a) = K.fromArray a #-}

-- | Convert an 'Array' into a stream in reverse order.
--
-- /Pre-release/
{-# INLINE_EARLY toStreamRev #-}
toStreamRev :: (Monad m, Storable a) => Array a -> SerialT m a
toStreamRev :: Array a -> SerialT m a
toStreamRev = Stream m a -> SerialT m a
forall (m :: * -> *) a. Stream m a -> SerialT m a
SerialT (Stream m a -> SerialT m a)
-> (Array a -> Stream m a) -> Array a -> SerialT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream m a -> Stream m a
forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
D.toStreamK (Stream m a -> Stream m a)
-> (Array a -> Stream m a) -> Array a -> Stream m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
toStreamDRev

-- XXX add fallback to StreamK rule
-- {-# RULES "Streamly.Array.readRev fallback to StreamK" [1]
--     forall a. S.toStreamK (readRev a) = K.revFromArray a #-}

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

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

-- | 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 Array a
arr = (Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze Array a
a, Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze Array a
b)
  where
    (Array a
a, Array a
b) = Int -> Array a -> (Array a, Array a)
forall a. Storable a => Int -> Array a -> (Array a, Array a)
MA.splitAt Int
i (Array a -> Array a
forall a. Array a -> Array a
unsafeThaw Array a
arr)

-- 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
ArrayContents
aEnd :: Ptr a
arrStart :: Ptr a
arrContents :: ArrayContents
aEnd :: forall a. Array a -> Ptr a
arrStart :: forall a. Array a -> Ptr a
arrContents :: forall a. Array a -> ArrayContents
..} = Ptr a -> b
go Ptr a
arrStart
    where

    go :: Ptr a -> b
go Ptr a
p | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr a
aEnd) (Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
aEnd) = 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
                    ArrayContents -> IO ()
touch ArrayContents
arrContents
                    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 (Streamly.Memory.Array)/
--
-- @since 0.8.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)

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

-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
-- 'Array'.
--
-- /Since 0.7.0 (Streamly.Memory.Array)/
--
-- @since 0.8.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 = (Array a -> Array a) -> Fold m a (Array a) -> Fold m a (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Fold m a (Array a) -> Fold m a (Array a))
-> (Int -> Fold m a (Array a)) -> Int -> Fold m a (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Fold m a (Array a)
MA.writeN

-- | @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 = (Array a -> Array a) -> Fold m a (Array a) -> Fold m a (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Fold m a (Array a) -> Fold m a (Array a))
-> (Int -> Fold m a (Array a)) -> Int -> Fold m a (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
MA.writeNAligned 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 =
    (Array a -> Array a) -> Fold m a (Array a) -> Fold m a (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Fold m a (Array a) -> Fold m a (Array a))
-> (Int -> Fold m a (Array a)) -> Int -> Fold m a (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
MA.writeNAlignedUnmanaged Int
alignSize

-- | 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 = Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Array a -> Array a) -> Fold m a (Array a) -> Fold m a (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Fold m a (Array a)
MA.writeNUnsafe Int
n

{-# INLINE_NORMAL writeWith #-}
writeWith :: forall m a. (MonadIO m, Storable a)
    => Int -> Fold m a (Array a)
-- writeWith n = FL.rmapM spliceArrays $ toArraysOf n
writeWith :: Int -> Fold m a (Array a)
writeWith Int
elemCount = Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze (Array a -> Array a) -> Fold m a (Array a) -> Fold m a (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Fold m a (Array a)
MA.writeWith Int
elemCount

-- | Fold the whole input to a single array.
--
-- /Caution! Do not use this on infinite streams./
--
-- /Since 0.7.0 (Streamly.Memory.Array)/
--
-- @since 0.8.0
{-# INLINE write #-}
write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)
write :: Fold m a (Array a)
write = (Array a -> Array a) -> Fold m a (Array a) -> Fold m a (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array a -> Array a
forall a. Array a -> Array a
unsafeFreeze Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Fold m a (Array a)
MA.write

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

instance (Show a, Storable a) => Show (Array a) where
    {-# INLINE showsPrec #-}
    showsPrec :: Int -> Array a -> ShowS
showsPrec Int
_ = [a] -> ShowS
forall a. Show a => a -> ShowS
shows ([a] -> ShowS) -> (Array a -> [a]) -> Array a -> ShowS
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 = [a] -> Array a
forall a. Storable a => [a] -> Array a
fromList ([a] -> Array a) -> ReadPrec [a] -> ReadPrec (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReadPrec [a]
forall a. Read a => ReadPrec a
readPrec
    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 :: String -> Array a
fromString = String -> 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

-- 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
arr1 == :: Array a -> Array a -> Bool
== Array a
arr2 = IO Bool -> Bool
forall a. IO a -> a
unsafeInlineIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$! Array a -> Array a
forall a. Array a -> Array a
unsafeThaw Array a
arr1 Array a -> Array a -> IO Bool
forall (m :: * -> *) a. MonadIO m => Array a -> Array a -> m Bool
`MA.cmp` Array a -> Array a
forall a. Array a -> Array a
unsafeThaw Array a
arr2

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

instance (Storable a, Ord a) => Ord (Array a) where
    {-# INLINE compare #-}
    compare :: Array a -> Array a -> Ordering
compare Array a
arr1 Array a
arr2 = Identity Ordering -> Ordering
forall a. Identity a -> a
runIdentity (Identity Ordering -> Ordering) -> Identity Ordering -> Ordering
forall a b. (a -> b) -> a -> b
$
        (a -> a -> Ordering)
-> Stream Identity a -> Stream Identity a -> Identity 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 Identity a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
toStreamD Array a
arr1) (Array a -> Stream Identity 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{..} = D.Stream step arrStart

    where

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

{-

XXX Why isn't Storable implicit? This does not compile unless I use the Storable
contraint.

{-# INLINE_NORMAL _foldr #-}
_foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
_foldr f z arr =
    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
-------------------------------------------------------------------------------

instance Storable a => Semigroup (Array a) where
    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)
splice Array a
arr1 Array a
arr2

nil ::
#ifdef DEVBUILD
    Storable a =>
#endif
    Array a
nil :: Array a
nil = ArrayContents -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array ArrayContents
MA.nilArrayContents (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
    mempty :: Array a
mempty = Array a
forall a. Array a
nil
    mappend :: Array a -> Array a -> Array a
mappend = Array a -> Array a -> Array a
forall a. Semigroup a => a -> a -> a
(<>)