@@ -475,43 +475,13 @@ take n (Stream step state) = n `seq` Stream step' (state, 0)
475475-- Grouping/Splitting
476476------------------------------------------------------------------------------
477477
478−{-# INLINE_LATE foldOneGroup #-}
479−foldOneGroup
480− :: Monad m
481− => Int
482− -> Fold m a b
483− -> a
484− -> State K.Stream m a
485− -> (State K.Stream m a -> s -> m (Step s a))
486− -> s
487− -> m (b, Maybe s)
488−foldOneGroup n (Fold fstep begin done) x gst step state = do
489− acc0 <- begin
490− acc <- fstep acc0 x
491− go SPEC state acc 1
478+-- s = stream state, fs = fold state
479+data GroupState s fs
480+ = GroupStart s
481+ | GroupBuffer s fs Int
482+ | GroupYield fs (GroupState s fs)
483+ | GroupFinish
492484
493− where
494−
495− -- XXX is it strict enough?
496− go !_ st !acc i | i < n = do
497− r <- step gst st
498− case r of
499− Yield y s -> do
500− acc1 <- fstep acc y
501− go SPEC s acc1 (i + 1)
502− Skip s -> go SPEC s acc i
503− Stop -> do
504− res <- done acc
505− return (res, Nothing)
506− go !_ st acc _ = do
507− r <- done acc
508− return (r, Just st)
509−
510−-- groupsOf takes 11 sec to write the file from a stream, whereas just
511−-- reading in the same file as a stream and folding each element of the stream
512−-- using (+) takes just 1.5 sec, so we are still pretty slow (7x slow), there
513−-- is scope to make it faster. There is a possibility of better fusion here.
514−--
515485{-# INLINE_NORMAL groupsOf #-}
516486groupsOf
517487 :: Monad m
@@ -519,31 +489,38 @@ groupsOf
519489 -> Fold m a b
520490 -> Stream m a
521491 -> Stream m b
522−groupsOf n f (Stream step state) =
523− n `seq` Stream stepOuter (Just state)
492+groupsOf n (Fold fstep initial extract) (Stream step state) =
493+ n `seq` Stream step' (GroupStart state)
524494
525495 where
526496
527− {-# INLINE_LATE stepOuter #-}
528− stepOuter gst (Just st) = do
497+ {-# INLINE_LATE step' #-}
498+ step' _ (GroupStart st) = do
529499 -- XXX shall we use the Natural type instead? Need to check performance
530500 -- implications.
531501 when (n <= 0) $
532502 -- XXX we can pass the module string from the higher level API
533503 error $ "Streamly.Streams.StreamD.Type.groupsOf: the size of "
534504 ++ "groups [" ++ show n ++ "] must be a natural number"
505+ -- fs = fold state
506+ fs <- initial
507+ return $ Skip (GroupBuffer st fs 0)
535508
536− -- We retrieve the first element of the stream before we start to fold
537− -- a chunk so that we do not return an empty chunk in case the stream
538− -- is empty.
539− res <- step (adaptState gst) st
540− case res of
509+ step' gst (GroupBuffer st fs i) = do
510+ r <- step (adaptState gst) st
511+ case r of
541512 Yield x s -> do
542− -- XXX how to make sure that stepInner and f get fused
543− -- This problem seems to be similar to the concatMap problem
544− (r, s1) <- foldOneGroup n f x (adaptState gst) step s
545− return $ Yield r s1
546− Skip s -> return $ Skip $ Just s
547− Stop -> return Stop
548−
549− stepOuter _ Nothing = return Stop
513+ fs' <- fstep fs x
514+ let i' = i + 1
515+ return $
516+ if i' >= n
517+ then Skip (GroupYield fs' (GroupStart s))
518+ else Skip (GroupBuffer s fs' i')
519+ Skip s -> return $ Skip (GroupBuffer s fs i)
520+ Stop -> return $ Skip (GroupYield fs GroupFinish)
521+
522+ step' _ (GroupYield fs next) = do
523+ r <- extract fs
524+ return $ Yield r next
525+
526+ step' _ GroupFinish = return Stop