@@ -92,7 +92,6 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
9292
9393 -- * Stream of arrays
9494 , arraysOf
95− , bufferChunks
9695 , writeChunks
9796
9897 -- * Utilities
@@ -527,10 +526,10 @@ arraysOf n (D.Stream step state) =
527526
528527-- XXX buffer to a list instead?
529528-- | Buffer the stream into arrays in memory.
530−{-# INLINE bufferChunks #-}
531−bufferChunks :: (MonadIO m, Storable a) =>
529+{-# INLINE arrayStreamKFromStreamD #-}
530+arrayStreamKFromStreamD :: (MonadIO m, Storable a) =>
532531 D.Stream m a -> m (K.Stream m (Array a))
533−bufferChunks m = D.foldr K.cons K.nil $ arraysOf defaultChunkSize m
532+arrayStreamKFromStreamD m = D.foldr K.cons K.nil $ arraysOf defaultChunkSize m
534533
535534-------------------------------------------------------------------------------
536535-- Streams of arrays - Flattening
@@ -898,17 +897,19 @@ writeNUnsafe = writeNUnsafeWith newArray
898897--
899898-- | Buffer a stream into a stream of arrays.
900899--
901−-- @writeChunks = Fold.many Fold.toStream (Array.writeN n)@
900+-- >>> writeChunks n = FL.many (writeN n) FL.toStreamK
902901--
903902-- Breaking an array into an array stream can be useful to consume a large
904903-- array sequentially such that memory of the array is released incrementatlly.
905904--
905+-- See also: 'arrayStreamKFromStreamD'.
906+--
906907-- /Unimplemented/
907908--
908909{-# INLINE_NORMAL writeChunks #-}
909−writeChunks :: -- (MonadIO m, Storable a) =>
910− Int -> Fold m a (D.Stream m (Array a))
911−writeChunks = undefined -- Fold.many Fold.toStream (Array.writeN n)
910+writeChunks :: (MonadIO m, Storable a) =>
911+ Int -> Fold m a (K.Stream n (Array a))
912+writeChunks n = FL.many (writeN n) FL.toStreamK
912913
913914-- XXX Compare toArrayMinChunk with fromStreamD which uses an array of streams
914915-- implementation. We can write this using writeChunks above if that is faster.
@@ -1018,6 +1019,20 @@ fromListN n xs = fromStreamDN n $ D.fromList xs
10181019-- convert stream to a single array
10191020-------------------------------------------------------------------------------
10201021
1022+{-# INLINE arrayStreamKLength #-}
1023+arrayStreamKLength :: (Monad m, Storable a) => K.Stream m (Array a) -> m Int
1024+arrayStreamKLength as = K.foldl' (+) 0 (K.map length as)
1025+
1026+-- | Convert an array stream to an array. Note that this requires peak memory
1027+-- that is double the size of the array stream.
1028+--
1029+{-# INLINE fromArrayStreamK #-}
1030+fromArrayStreamK :: (Storable a, MonadIO m) =>
1031+ K.Stream m (Array a) -> m (Array a)
1032+fromArrayStreamK as = do
1033+ len <- arrayStreamKLength as
1034+ fromStreamDN len $ D.unfoldMany read $ D.fromStreamK as
1035+
10211036-- CAUTION: a very large number (millions) of arrays can degrade performance
10221037-- due to GC overhead because we need to buffer the arrays before we flatten
10231038-- all the arrays.
@@ -1035,10 +1050,7 @@ fromListN n xs = fromStreamDN n $ D.fromList xs
10351050--
10361051{-# INLINE fromStreamD #-}
10371052fromStreamD :: (MonadIO m, Storable a) => D.Stream m a -> m (Array a)
1038−fromStreamD m = do
1039− buffered <- bufferChunks m
1040− len <- K.foldl' (+) 0 (K.map length buffered)
1041− fromStreamDN len $ D.unfoldMany read $ D.fromStreamK buffered
1053+fromStreamD m = arrayStreamKFromStreamD m >>= fromArrayStreamK
10421054
10431055-- | Create an 'Array' from a list. The list must be of finite size.
10441056--