@@ -38,7 +38,7 @@ module Streamly.Mem.Array.Types
3838 , fromStreamDArraysOf
3939 , flattenArrays
4040 , flattenArraysRev
41− , coalesceChunksOf
41+ , packArraysChunksOf
4242 , groupIOVecsOf
4343 , splitOn
4444
@@ -485,12 +485,49 @@ fromStreamDN limit str = do
485485 liftIO $ poke ptr x
486486 return $ ptr `plusPtr` sizeOf (undefined :: a)
487487
488+data GroupState s start end bound
489+ = GroupStart s
490+ | GroupBuffer s start end bound
491+ | GroupYield start end bound (GroupState s start end bound)
492+ | GroupFinish
493+
488494-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
489495-- arrays of size n.
490−{-# INLINE fromStreamDArraysOf #-}
496+{-# INLINE_NORMAL fromStreamDArraysOf #-}
491497fromStreamDArraysOf :: (MonadIO m, Storable a)
492498 => Int -> D.Stream m a -> D.Stream m (Array a)
493−fromStreamDArraysOf n str = D.groupsOf n (toArrayN n) str
499+-- fromStreamDArraysOf n str = D.groupsOf n (toArrayN n) str
500+fromStreamDArraysOf n (D.Stream step state) =
501+ D.Stream step' (GroupStart state)
502+
503+ where
504+
505+ {-# INLINE_LATE step' #-}
506+ step' _ (GroupStart st) = do
507+ when (n <= 0) $
508+ -- XXX we can pass the module string from the higher level API
509+ error $ "Streamly.Mem.Array.Types.fromStreamDArraysOf: the size of "
510+ ++ "arrays [" ++ show n ++ "] must be a natural number"
511+ Array start end bound <- liftIO $ newArray n
512+ return $ D.Skip (GroupBuffer st start end bound)
513+
514+ step' gst (GroupBuffer st start end bound) = do
515+ r <- step (adaptState gst) st
516+ case r of
517+ D.Yield x s -> do
518+ liftIO $ poke end x
519+ let end' = end `plusPtr` 1
520+ return $
521+ if end' >= bound
522+ then D.Skip (GroupYield start end' bound (GroupStart s))
523+ else D.Skip (GroupBuffer s start end' bound)
524+ D.Skip s -> return $ D.Skip (GroupBuffer s start end bound)
525+ D.Stop -> return $ D.Skip (GroupYield start end bound GroupFinish)
526+
527+ step' _ (GroupYield start end bound next) =
528+ return $ D.Yield (Array start end bound) next
529+
530+ step' _ GroupFinish = return D.Stop
494531
495532-- XXX concatMap does not seem to have the best possible performance so we have
496533-- a custom way to concat arrays.
@@ -833,15 +870,20 @@ data SpliceState s arr
833870-- maximum specified size.
834871--
835872-- @since 0.7.0
836−{-# INLINE_NORMAL coalesceChunksOf #-}
837−coalesceChunksOf :: (MonadIO m, Storable a)
873+{-# INLINE_NORMAL packArraysChunksOf #-}
874+packArraysChunksOf :: (MonadIO m, Storable a)
838875 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
839−coalesceChunksOf n (D.Stream step state) = D.Stream step' (SpliceInitial state)
876+packArraysChunksOf n (D.Stream step state) =
877+ D.Stream step' (SpliceInitial state)
840878
841879 where
842880
843881 {-# INLINE_LATE step' #-}
844882 step' gst (SpliceInitial st) = do
883+ when (n <= 0) $
884+ -- XXX we can pass the module string from the higher level API
885+ error $ "Streamly.Mem.Array.Types.packArraysChunksOf: the size of "
886+ ++ "arrays [" ++ show n ++ "] must be a natural number"
845887 r <- step gst st
846888 case r of
847889 D.Yield arr s -> return $
@@ -858,7 +900,8 @@ coalesceChunksOf n (D.Stream step state) = D.Stream step' (SpliceInitial state)
858900 D.Yield arr s -> do
859901 let len = byteLength buf + byteLength arr
860902 if len > n
861− then return $ D.Skip (SpliceYielding buf (SpliceBuffering s arr))
903+ then return $
904+ D.Skip (SpliceYielding buf (SpliceBuffering s arr))
862905 else do
863906 buf' <- if byteCapacity buf < n
864907 then liftIO $ realloc n buf
@@ -893,6 +936,14 @@ groupIOVecsOf n maxIOVLen (D.Stream step state) =
893936
894937 {-# INLINE_LATE step' #-}
895938 step' gst (GatherInitial st) = do
939+ when (n <= 0) $
940+ -- XXX we can pass the module string from the higher level API
941+ error $ "Streamly.Mem.Array.Types.groupIOVecsOf: the size of "
942+ ++ "groups [" ++ show n ++ "] must be a natural number"
943+ when (maxIOVLen <= 0) $
944+ -- XXX we can pass the module string from the higher level API
945+ error $ "Streamly.Mem.Array.Types.groupIOVecsOf: the number of "
946+ ++ "IOVec entries [" ++ show n ++ "] must be a natural number"
896947 r <- step (adaptState gst) st
897948 case r of
898949 D.Yield arr s -> do