@@ -99,7 +99,7 @@ module Streamly.Mem.Array
9999 -- Folds
100100 , toArrayN
101101 -- , toArrays
102− -- , toArray
102+ , toArray
103103
104104 -- Streams
105105 , arraysOf
@@ -132,8 +132,9 @@ module Streamly.Mem.Array
132132where
133133
134134import Control.Monad.IO.Class (MonadIO(..))
135−import Data.Functor.Identity
135+import Data.Functor.Identity (Identity)
136136import Foreign.ForeignPtr (withForeignPtr)
137+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
137138import Foreign.Ptr (minusPtr, plusPtr, castPtr)
138139import Foreign.Storable (Storable(..))
139140import Prelude hiding (length, null, last, map, (!!), read)
@@ -141,11 +142,13 @@ import Prelude hiding (length, null, last, map, (!!), read)
141142import Streamly.Mem.Array.Types hiding (flattenArrays, newArray)
142143import Streamly.Streams.Serial (SerialT)
143144import Streamly.Streams.StreamK.Type (IsStream)
144−import Streamly.Fold (Fold, lchunksOf, toStream)
145+import Streamly.Fold.Types (Fold(..))
145146
147+import qualified Streamly.Fold as FL
146148import qualified Streamly.Mem.Array.Types as A
147149import qualified Streamly.Prelude as S
148150import qualified Streamly.Streams.StreamD as D
151+import qualified Streamly.Streams.Prelude as P
149152
150153-------------------------------------------------------------------------------
151154-- Construction
@@ -353,11 +356,63 @@ writeSliceRev arr i len s = undefined
353356-- Streams of Arrays
354357-------------------------------------------------------------------------------
355358
359+-- exponentially increasing sizes of the chunks upto the max limit.
360+-- XXX this will be easier to implement with parsers/terminating folds
361+-- With this we should be able to reduce the number of chunks/allocations.
362+-- The reallocation/copy based toArray can also be implemented using this.
363+--
364+{-
365+{-# INLINE toArraysInRange #-}
366+toArraysInRange :: (IsStream t, MonadIO m, Storable a)
367+ => Int -> Int -> Fold m (Array a) b -> Fold m a b
368+toArraysInRange low high (Fold step initial extract) =
369+-}
370+
356371-- | Fold the input to a pure buffered stream (List) of arrays.
357−{-# INLINE _toArrays #-}
358−_toArrays :: (MonadIO m, Storable a)
372+{-# INLINE _toArraysOf #-}
373+_toArraysOf :: (MonadIO m, Storable a)
359374 => Int -> Fold m a (SerialT Identity (Array a))
360−_toArrays n = lchunksOf n (toArrayN n) toStream
375+_toArraysOf n = FL.lchunksOf n (toArrayN n) FL.toStream
376+
377+-- XXX The realloc based implementation needs to make one extra copy if we use
378+-- shrinkToFit. On the other hand, the stream of arrays implementation may
379+-- buffer the array chunk pointers in memory but it does not have to shrink as
380+-- we know the exact size in the end. However, memory copying does not seems to
381+-- be as expensive as the allocations. Therefore, we need to reduce the number
382+-- of allocations instead. Also, the size of allocations matters, right sizing
383+-- an allocation even at the cost of copying sems to help. Should be measured
384+-- on a big stream with heavy calls to toArray to see the effect.
385+--
386+-- XXX check if GHC's memory allocator is efficient enough. We can try the C
387+-- malloc to compare against.
388+
389+{-# INLINE toArrayMinChunk #-}
390+toArrayMinChunk :: forall m a. (MonadIO m, Storable a)
391+ => Int -> Fold m a (Array a)
392+-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n
393+toArrayMinChunk chunkSize = Fold step initial extract
394+
395+ where
396+
397+ insertElem (Array start end bound) x = do
398+ liftIO $ poke end x
399+ return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
400+
401+ initial = liftIO $ A.newArray (mkChunkSize chunkSize)
402+ step arr@(Array _ end bound) x | end == bound = do
403+ arr1 <- liftIO $ reallocDouble 1 arr
404+ insertElem arr1 x
405+ step arr x = insertElem arr x
406+ extract = liftIO . shrinkToFit
407+
408+-- | Fold the whole input to a single array.
409+--
410+-- /Caution! Do not use this on infinite streams./
411+--
412+-- @since 0.7.0
413+{-# INLINE toArray #-}
414+toArray :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)
415+toArray = toArrayMinChunk 1024
361416
362417-- | Convert a stream of arrays into a stream of their elements.
363418--
@@ -389,14 +444,14 @@ arraysOf :: (IsStream t, MonadIO m, Storable a)
389444arraysOf n str =
390445 D.fromStreamD $ fromStreamDArraysOf n (D.toStreamD str)
391446
392−-- | Given a stream of arrays, splice them all together to generate a single
393−-- array. The stream must be /finite/.
394−--
395−-- @since 0.7.0
396−{-# INLINABLE spliceArrays #-}
397−spliceArrays :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
398−spliceArrays s = do
399− buffered <- S.foldr S.cons S.nil s
447+-- XXX Both of these implementations of splicing seem to perform equally well.
448+-- We need to perform benchmarks over a range of sizes though.
449+
450+{-# INLINE _spliceArraysRealloced #-}
451+_spliceArraysRealloced :: (MonadIO m, Storable a)
452+ => SerialT m (Array a) -> m (Array a)
453+_spliceArraysRealloced s = do
454+ buffered <- P.foldr S.cons S.nil s
400455 len <- S.sum (S.map length buffered)
401456
402457 arr <- liftIO $ A.newArray len
@@ -411,6 +466,38 @@ spliceArrays s = do
411466 memcpy (castPtr dst) (castPtr src) len
412467 return $ dst `plusPtr` len
413468
469+{-# INLINE spliceArraysBuffered #-}
470+spliceArraysBuffered :: (MonadIO m, Storable a)
471+ => SerialT m (Array a) -> m (Array a)
472+spliceArraysBuffered s = do
473+ idst <- liftIO $ A.newArray (mkChunkSizeKB 4)
474+ arr <- S.foldlM' appendArr idst s
475+ liftIO $ shrinkToFit arr
476+
477+ where
478+
479+ appendArr dst@(Array _ end bound) src = liftIO $ do
480+ let srcLen = aEnd src `minusPtr` unsafeForeignPtrToPtr (aStart src)
481+ dst1 <-
482+ if end `plusPtr` srcLen >= bound
483+ then reallocDouble srcLen dst
484+ else return dst
485+
486+ withForeignPtr (aStart dst1) $ \_ -> do
487+ withForeignPtr (aStart src) $ \psrc -> do
488+ let pdst = aEnd dst1
489+ memcpy (castPtr pdst) (castPtr psrc) srcLen
490+ return $ dst1 { aEnd = pdst `plusPtr` srcLen }
491+
492+-- | Given a stream of arrays, splice them all together to generate a single
493+-- array. The stream must be /finite/.
494+--
495+-- @since 0.7.0
496+{-# INLINE spliceArrays #-}
497+spliceArrays :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
498+spliceArrays = spliceArraysBuffered
499+-- spliceArrays = _spliceArraysRealloced
500+
414501-- | Create an 'Array' from a stream. This is useful when we want to create a
415502-- single array from a stream of unknown size. 'writeN' is at least twice
416503-- as efficient when the size is already known.