@@ -132,7 +132,7 @@ module Streamly.Internal.Data.MutArray.Generic
132132 -- multidimensional array representations.
133133
134134 -- ** Construct from streams
135− -- , chunksOf
135+ , chunksOf
136136 -- , arrayStreamKFromStreamD
137137 -- , writeChunks
138138
@@ -182,6 +182,7 @@ import Streamly.Internal.Data.Fold.Type (Fold(..))
182182import Streamly.Internal.Data.Producer.Type (Producer (..))
183183import Streamly.Internal.Data.Unfold.Type (Unfold(..))
184184import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
185+import Streamly.Internal.Data.SVar.Type (adaptState)
185186
186187import qualified Streamly.Internal.Data.Fold.Type as FL
187188import qualified Streamly.Internal.Data.Producer as Producer
@@ -263,6 +264,21 @@ nil = new 0
263264-- Random writes
264265-------------------------------------------------------------------------------
265266
267+-- | Write the given element to the given index of the 'MutableArray#'. Does not
268+-- check if the index is out of bounds of the array.
269+--
270+-- /Pre-release/
271+{-# INLINE putIndexUnsafeWith #-}
272+putIndexUnsafeWith :: MonadIO m => Int -> MutableArray# RealWorld a -> a -> m ()
273+putIndexUnsafeWith n _arrContents# x =
274+ liftIO
275+ $ IO
276+ $ \s# ->
277+ case n of
278+ I# n# ->
279+ let s1# = writeArray# _arrContents# n# x s#
280+ in (# s1#, () #)
281+
266282-- | Write the given element to the given index of the array. Does not check if
267283-- the index is out of bounds of the array.
268284--
@@ -271,13 +287,7 @@ nil = new 0
271287putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()
272288putIndexUnsafe i MutArray {..} x =
273289 assert (i >= 0 && i < arrLen)
274− (liftIO
275− $ IO
276− $ \s# ->
277− case i + arrStart of
278− I# n# ->
279− let s1# = writeArray# arrContents# n# x s#
280− in (# s1#, () #))
290+ putIndexUnsafeWith (i + arrStart) arrContents# x
281291
282292invalidIndex :: String -> Int -> a
283293invalidIndex label i =
@@ -666,6 +676,75 @@ fromPureStream :: MonadIO m => Stream Identity a -> m (MutArray a)
666676fromPureStream xs =
667677 liftIO $ D.fold write $ D.morphInner (return . runIdentity) xs
668678
679+-------------------------------------------------------------------------------
680+-- Chunking
681+-------------------------------------------------------------------------------
682+
683+data GroupState s a start end bound
684+ = GroupStart s
685+ | GroupBuffer s (MutableArray# RealWorld a) start end bound
686+ | GroupYield
687+ (MutableArray# RealWorld a)
688+ start
689+ end
690+ bound
691+ (GroupState s a start end bound)
692+ | GroupFinish
693+
694+-- | @chunksOf n stream@ groups the input stream into a stream of
695+-- arrays of size n.
696+--
697+-- @chunksOf n = foldMany (MutArray.writeN n)@
698+--
699+-- /Pre-release/
700+{-# INLINE_NORMAL chunksOf #-}
701+chunksOf :: forall m a. MonadIO m
702+ => Int -> D.Stream m a -> D.Stream m (MutArray a)
703+-- XXX the idiomatic implementation leads to large regression in the D.reverse'
704+-- benchmark. It seems it has difficulty producing optimized code when
705+-- converting to StreamK. Investigate GHC optimizations.
706+-- chunksOf n = D.foldMany (writeN n)
707+chunksOf n (D.Stream step state) =
708+ D.Stream step' (GroupStart state)
709+
710+ where
711+
712+ -- start is always 0
713+ -- end and len are always equal
714+
715+ {-# INLINE_LATE step' #-}
716+ step' _ (GroupStart st) = do
717+ when (n <= 0) $
718+ -- XXX we can pass the module string from the higher level API
719+ error $ "Streamly.Internal.Data.Array.Generic.Mut.Type.chunksOf: "
720+ ++ "the size of arrays [" ++ show n
721+ ++ "] must be a natural number"
722+ (MutArray contents start end bound :: MutArray a) <- new n
723+ return $ D.Skip (GroupBuffer st contents start end bound)
724+
725+ step' gst (GroupBuffer st contents start end bound) = do
726+ r <- step (adaptState gst) st
727+ case r of
728+ D.Yield x s -> do
729+ putIndexUnsafeWith end contents x
730+ let end1 = end + 1
731+ return $
732+ if end1 >= bound
733+ then D.Skip
734+ (GroupYield
735+ contents start end1 bound (GroupStart s))
736+ else D.Skip (GroupBuffer s contents start end1 bound)
737+ D.Skip s ->
738+ return $ D.Skip (GroupBuffer s contents start end bound)
739+ D.Stop ->
740+ return
741+ $ D.Skip (GroupYield contents start end bound GroupFinish)
742+
743+ step' _ (GroupYield contents start end bound next) =
744+ return $ D.Yield (MutArray contents start end bound) next
745+
746+ step' _ GroupFinish = return D.Stop
747+
669748-------------------------------------------------------------------------------
670749-- Unfolds
671750-------------------------------------------------------------------------------