git.iko.soy/streamly
on master help sign in
Implement Data.Fold.concatMap.
039eb7aPranay Sashank
src/Streamly/Internal/Data/Fold/Types.hs+46−8
compare with master
unifiedsplit
@@ -685,19 +685,57 @@ longest :: -- Monad m =>
685 Fold m x a -> Fold m x a -> Fold m x a
685 Fold m x a -> Fold m x a -> Fold m x a
686longest _f1 _f2 = undefined
686longest _f1 _f2 = undefined
687
687
688+data ConcatMapState m sa a c
689+ = B !sa
690+ | forall s. C (s -> a -> m (Step s c)) !s (s -> m c)
691+
688-- | Map a 'Fold' returning function on the result of a 'Fold'.
692-- | Map a 'Fold' returning function on the result of a 'Fold'.
689--
693--
690−-- Compare with 'Monad' instance method '>>='. This implementation allows
694+-- Compare with 'Monad' instance method '>>='.
691−-- stream fusion but has quadratic complexity. This can fuse with other
692−-- operations and can be much faster than 'Monad' instance for small number
693−-- (less than 8) of compositions.
694--
695--
695−-- /Unimplemented/
696+-- >>> Stream.fold (concatMap (flip Fold.takeLE Fold.sum) (Fold.rmapM (return . fromJust) Fold.head)) $ Stream.fromList [10,9..1]
697+-- 45
698+--
699+-- /Internal/
696--
700--
697{-# INLINE concatMap #-}
701{-# INLINE concatMap #-}
698−concatMap :: -- Monad m =>
699− (b -> Fold m a c) -> Fold m a b -> Fold m a c
702+concatMap :: Monad m => (b -> Fold m a c) -> Fold m a b -> Fold m a c
700−concatMap _func _fld = undefined
703+concatMap f (Fold stepa initiala extracta) = Fold stepc initialc extractc
704+ where
705+ initialc = do
706+ r <- initiala
707+ case r of
708+ Partial s -> return $ Partial (B s)
709+ Done b -> initInnerFold (f b)
710+
711+ stepc (B s) a = do
712+ r <- stepa s a
713+ case r of
714+ Partial s1 -> return $ Partial (B s1)
715+ Done b -> initInnerFold (f b)
716+
717+ stepc (C stepInner s extractInner) a = do
718+ r <- stepInner s a
719+ return $ case r of
720+ Partial sc -> Partial (C stepInner sc extractInner)
721+ Done c -> Done c
722+
723+ extractc (B s) = do
724+ r <- extracta s
725+ initExtract (f r)
726+ extractc (C _ sInner extractInner) = extractInner sInner
727+
728+ initInnerFold (Fold step i e) = do
729+ r <- i
730+ return $ case r of
731+ Partial s -> Partial (C step s e)
732+ Done c -> Done c
733+
734+ initExtract (Fold _ i e) = do
735+ r <- i
736+ case r of
737+ Partial s -> e s
738+ Done c -> return c
701
739
702-- | Combines the outputs of the folds (the type @b@) using their 'Semigroup'
740-- | Combines the outputs of the folds (the type @b@) using their 'Semigroup'
703-- instances.
741-- instances.