@@ -475,43 +475,13 @@ take n (Stream step state) = n `seq` Stream step' (state, 0)
475-- Grouping/Splitting
475-- Grouping/Splitting
476------------------------------------------------------------------------------
476------------------------------------------------------------------------------
478−{-# INLINE_LATE foldOneGroup #-}
478+-- s = stream state, fs = fold state
479−foldOneGroup
479+data GroupState s fs
480− :: Monad m
480+ = GroupStart s
481− => Int
481+ | GroupBuffer s fs Int
482− -> Fold m a b
482+ | GroupYield fs (GroupState s fs)
483− -> a
483+ | GroupFinish
484− -> State K.Stream m a
485− -> (State K.Stream m a -> s -> m (Step s a))
488−foldOneGroup n (Fold fstep begin done) x gst step state = do
495− -- XXX is it strict enough?
496− go !_ st !acc i | i < n = do
501− go SPEC s acc1 (i + 1)
502− Skip s -> go SPEC s acc i
505− return (res, Nothing)
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.
515{-# INLINE_NORMAL groupsOf #-}
485{-# INLINE_NORMAL groupsOf #-}
517 :: Monad m
487 :: Monad m
@@ -519,31 +489,38 @@ groupsOf
519 -> Fold m a b
489 -> Fold m a b
520 -> Stream m a
490 -> Stream m a
521 -> Stream m b
491 -> Stream m b
522−groupsOf n f (Stream step state) =
492+groupsOf n (Fold fstep initial extract) (Stream step state) =
523− n `seq` Stream stepOuter (Just state)
493+ n `seq` Stream step' (GroupStart state)
527− {-# INLINE_LATE stepOuter #-}
497+ {-# INLINE_LATE step' #-}
528− stepOuter gst (Just st) = do
498+ step' _ (GroupStart st) = do
529 -- XXX shall we use the Natural type instead? Need to check performance
499 -- XXX shall we use the Natural type instead? Need to check performance
530 -- implications.
500 -- implications.
531 when (n <= 0) $
501 when (n <= 0) $
532 -- XXX we can pass the module string from the higher level API
502 -- XXX we can pass the module string from the higher level API
533 error $ "Streamly.Streams.StreamD.Type.groupsOf: the size of "
503 error $ "Streamly.Streams.StreamD.Type.groupsOf: the size of "
534 ++ "groups [" ++ show n ++ "] must be a natural number"
504 ++ "groups [" ++ show n ++ "] must be a natural number"
507+ return $ Skip (GroupBuffer st fs 0)
536− -- We retrieve the first element of the stream before we start to fold
509+ step' gst (GroupBuffer st fs i) = do
537− -- a chunk so that we do not return an empty chunk in case the stream
539− res <- step (adaptState gst) st
510+ r <- step (adaptState gst) st
540− case res of
511+ case r of
541 Yield x s -> do
512 Yield x s -> do
542− -- XXX how to make sure that stepInner and f get fused
513+ fs' <- fstep fs x
543− -- This problem seems to be similar to the concatMap problem
514+ let i' = i + 1
544− (r, s1) <- foldOneGroup n f x (adaptState gst) step s
515+ return $
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)
522+ step' _ (GroupYield fs next) = do
545− return $ Yield r s1
524+ return $ Yield r next
546− Skip s -> return $ Skip $ Just s
525+
549− stepOuter _ Nothing = return Stop
526+ step' _ GroupFinish = return Stop