@@ -28,6 +28,7 @@ where
2828
2929import Control.Monad.IO.Class (MonadIO(..))
3030import Control.Monad (when)
31+import Control.Monad.Catch (MonadThrow)
3132import Data.Bifunctor (first)
3233import Foreign.Storable (Storable(..))
3334import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..))
@@ -40,6 +41,7 @@ import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
4041import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MArray
4142import qualified Streamly.Internal.Data.Fold.Type as FL
4243import qualified Streamly.Internal.Data.Stream.StreamD as D
44+import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
4345
4446-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
4547-- @n@ elements each.
@@ -194,28 +196,64 @@ compact :: (MonadIO m, Storable a)
194196compact n (SerialT xs) =
195197 SerialT $ D.toStreamK $ packArraysChunksOf n (D.fromStreamK xs)
196198
197−-- See lpackArraysChunksOf/packArraysChunksOf to implement this.
198−--
199199-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
200200-- maximum specified size. Note that if a single array is bigger than the
201201-- specified size we do not split it to fit. When we coalesce multiple arrays
202202-- if the size would exceed the specified size we do not coalesce therefore the
203203-- actual array size may be less than the specified chunk size.
204204--
205−-- /Unimplemented/
206−{-# INLINE_NORMAL compactLEFold #-}
207−compactLEFold :: -- (MonadIO m, Storable a) =>
208− Int -> Fold m (Array a) (Array a)
209−compactLEFold = undefined
205+-- /Internal/
206+{-# INLINE_NORMAL compactLEParserD #-}
207+compactLEParserD ::
208+ forall m a. (MonadThrow m, MonadIO m, Storable a)
209+ => Int -> ParserD.Parser m (Array a) (Array a)
210+compactLEParserD n = ParserD.Parser step initial extract
211+
212+ where
213+
214+ nBytes = n * sizeOf (undefined :: a)
215+
216+ initial =
217+ return
218+ $ if n <= 0
219+ then error
220+ $ functionPath
221+ ++ ": the size of arrays ["
222+ ++ show n ++ "] must be a natural number"
223+ else ParserD.IPartial Nothing
224+
225+ step Nothing arr =
226+ return
227+ $ let len = MArray.byteLength arr
228+ in if len >= nBytes
229+ then ParserD.Done 0 arr
230+ else ParserD.Partial 0 (Just arr)
231+ step (Just buf) arr =
232+ let len = MArray.byteLength buf + MArray.byteLength arr
233+ in if len > nBytes
234+ then return $ ParserD.Done 1 buf
235+ else do
236+ buf1 <-
237+ if MArray.byteCapacity buf < nBytes
238+ then liftIO $ MArray.realloc nBytes buf
239+ else return buf
240+ buf2 <- MArray.splice buf1 arr
241+ return $ ParserD.Partial 0 (Just buf2)
242+
243+ extract Nothing = return MArray.nil
244+ extract (Just buf) = return buf
245+
246+ functionPath =
247+ "Streamly.Internal.Data.Array.Stream.Mut.Foreign.compactLEParserD"
210248
211249-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
212250-- maximum specified size in bytes.
213251--
214252-- /Internal/
215−compactLE :: (MonadIO m {-, Storable a-}) =>
253+compactLE :: (MonadThrow m, MonadIO m, Storable a) =>
216254 Int -> SerialT m (Array a) -> SerialT m (Array a)
217255compactLE n (SerialT xs) =
218− SerialT $ D.toStreamK $ D.foldMany (compactLEFold n) (D.fromStreamK xs)
256+ SerialT $ D.toStreamK $ D.parseMany (compactLEParserD n) (D.fromStreamK xs)
219257
220258-- | Like 'compactLE' but generates arrays of exactly equal to the size
221259-- specified except for the last array in the stream which could be shorter.