@@ -128,9 +128,16 @@ module Streamly.Mem.Array
128128 , writeSlice
129129 , writeSliceRev
130130 -}
131+ -- * Immutable Transformations
132+ , transformWith
133+
134+ -- * Folds
135+ , foldWith
136+ , foldArray
131137 )
132138where
133139
140+import Control.Monad (when)
134141import Control.Monad.IO.Class (MonadIO(..))
135142import Data.Functor.Identity (Identity)
136143import Foreign.ForeignPtr (withForeignPtr)
@@ -169,7 +176,9 @@ newArray len = undefined
169176-- @since 0.7.0
170177{-# INLINE writeN #-}
171178writeN :: (MonadIO m, Storable a) => Int -> SerialT m a -> m (Array a)
172−writeN n m = fromStreamDN n $ D.toStreamD m
179+writeN n m = do
180+ if n < 0 then error "writeN: negative write count specified" else return ()
181+ fromStreamDN n $ D.toStreamD m
173182
174183-------------------------------------------------------------------------------
175184-- Elimination
@@ -386,11 +395,17 @@ _toArraysOf n = FL.lchunksOf n (toArrayN n) FL.toStream
386395-- XXX check if GHC's memory allocator is efficient enough. We can try the C
387396-- malloc to compare against.
388397
398+{-# INLINE bytesToCount #-}
399+bytesToCount :: Storable a => a -> Int -> Int
400+bytesToCount x n =
401+ let elemSize = sizeOf x
402+ in n + elemSize - 1 `div` elemSize
403+
389404{-# INLINE toArrayMinChunk #-}
390405toArrayMinChunk :: forall m a. (MonadIO m, Storable a)
391406 => Int -> Fold m a (Array a)
392407-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n
393−toArrayMinChunk chunkSize = Fold step initial extract
408+toArrayMinChunk elemCount = Fold step initial extract
394409
395410 where
396411
@@ -398,7 +413,9 @@ toArrayMinChunk chunkSize = Fold step initial extract
398413 liftIO $ poke end x
399414 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
400415
401− initial = liftIO $ A.newArray (mkChunkSize chunkSize)
416+ initial = do
417+ when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"
418+ liftIO $ A.newArray elemCount
402419 step arr@(Array _ end bound) x | end == bound = do
403420 arr1 <- liftIO $ reallocDouble 1 arr
404421 insertElem arr1 x
@@ -412,7 +429,7 @@ toArrayMinChunk chunkSize = Fold step initial extract
412429-- @since 0.7.0
413430{-# INLINE toArray #-}
414431toArray :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)
415−toArray = toArrayMinChunk 1024
432+toArray = toArrayMinChunk (bytesToCount (undefined :: a) (mkChunkSize 1024))
416433
417434-- | Convert a stream of arrays into a stream of their elements.
418435--
@@ -467,10 +484,11 @@ _spliceArraysRealloced s = do
467484 return $ dst `plusPtr` len
468485
469486{-# INLINE spliceArraysBuffered #-}
470−spliceArraysBuffered :: (MonadIO m, Storable a)
487+spliceArraysBuffered :: forall m a. (MonadIO m, Storable a)
471488 => SerialT m (Array a) -> m (Array a)
472489spliceArraysBuffered s = do
473− idst <- liftIO $ A.newArray (mkChunkSizeKB 4)
490+ idst <- liftIO $ A.newArray (bytesToCount (undefined :: a)
491+ (mkChunkSizeKB 4))
474492 arr <- S.foldlM' appendArr idst s
475493 liftIO $ shrinkToFit arr
476494
@@ -508,8 +526,40 @@ spliceArrays = spliceArraysBuffered
508526--
509527{-# INLINE write #-}
510528write :: (MonadIO m, Storable a) => SerialT m a -> m (Array a)
511−write m = A.fromStreamD $ D.toStreamD m
529+write = FL.foldl' toArray
530+-- write m = A.fromStreamD $ D.toStreamD m
512531
513532-- XXX efficiently compare two streams of arrays. Two streams can have chunks
514533-- of different sizes, we can handle that in the stream comparison abstraction.
515534-- This could be useful e.g. to fast compare whether two files differ.
535+
536+-------------------------------------------------------------------------------
537+-- Transform via stream operations
538+-------------------------------------------------------------------------------
539+
540+-- for non-length changing operations we can use the original length for
541+-- allocation. If we can predict the length then we can use the prediction for
542+-- new allocation. Otherwise we can use a hint and adjust dynamically.
543+
544+-- | Transform an array into another array using a stream transformation
545+-- operation.
546+--
547+-- @since 0.7.0
548+{-# INLINE transformWith #-}
549+transformWith :: (MonadIO m, Storable a, Storable b)
550+ => (SerialT m a -> SerialT m b) -> Array a -> m (Array b)
551+transformWith f arr = FL.foldl' (toArrayMinChunk (length arr)) $ f (read arr)
552+
553+-- | Fold an array using a 'Fold'.
554+--
555+-- @since 0.7.0
556+{-# INLINE foldArray #-}
557+foldArray :: (MonadIO m, Storable a) => Fold m a b -> Array a -> m b
558+foldArray f arr = FL.foldl' f (read arr)
559+
560+-- | Fold an array using a stream fold operation.
561+--
562+-- @since 0.7.0
563+{-# INLINE foldWith #-}
564+foldWith :: (MonadIO m, Storable a) => (SerialT m a -> m b) -> Array a -> m b
565+foldWith f arr = f (read arr)