@@ -58,6 +58,8 @@ module Streamly.Streams.StreamD
5858 -- | Direct style stream does not support @fromFoldable@.
5959 , yield
6060 , yieldM
61+ , generate
62+ , generateM
6163 , fromList
6264 , fromListM
6365 , fromStreamK
@@ -87,8 +89,11 @@ module Streamly.Streams.StreamD
8789 , minimumBy
8890 , findIndices
8991 , lookup
92+ , findM
9093 , find
9194 , (!!)
95+ , concatMapM
96+ , concatMap
9297
9398 -- ** Substreams
9499 , isPrefixOf
@@ -162,7 +167,7 @@ import Prelude
162167 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
163168 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
164169 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
165− (!!), scanl, scanl1)
170+ (!!), scanl, scanl1, concatMap)
166171
167172import Streamly.SVar (MonadAsync, State(..), defState, rstState)
168173
@@ -285,6 +290,20 @@ yieldM m = Stream step True
285290 step _ True = m >>= \x -> return $ Yield x False
286291 step _ False = return Stop
287292
293+{-# INLINE generate #-}
294+generate :: Monad m => Int -> (Int -> a) -> Stream m a
295+generate n gen = generateM n (return . gen)
296+
297+{-# INLINE_NORMAL generateM #-}
298+generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
299+generateM n gen = n `seq` Stream step 0
300+ where
301+ {-# INLINE_LATE step #-}
302+ step _ i | i < n = do
303+ x <- gen i
304+ return $ Yield x (i+1)
305+ | otherwise = return Stop
306+
288307-- XXX we need the MonadAsync constraint because of a rewrite rule.
289308-- | Convert a list of monadic actions to a 'Stream'
290309{-# INLINE_LATE fromListM #-}
@@ -562,17 +581,23 @@ lookup e (Stream step state) = go state
562581 Skip s -> go s
563582 Stop -> return Nothing
564583
565−{-# INLINE_NORMAL find #-}
566−find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
567−find p (Stream step state) = go state
584+{-# INLINE_NORMAL findM #-}
585+findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
586+findM p (Stream step state) = go SPEC state
568587 where
569− go st = do
588+ go !_ st = do
570589 r <- step defState st
571590 case r of
572− Yield x s -> if p x then return (Just x) else go s
573− Skip s -> go s
591+ Yield x s -> do
592+ b <- p x
593+ if b then return (Just x) else go SPEC s
594+ Skip s -> go SPEC s
574595 Stop -> return Nothing
575596
597+{-# INLINE find #-}
598+find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
599+find p = findM (return . p)
600+
576601{-# INLINE_NORMAL findIndices #-}
577602findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
578603findIndices p (Stream step state) = Stream step' (state, 0)
@@ -585,6 +610,32 @@ findIndices p (Stream step state) = Stream step' (state, 0)
585610 Skip s -> Skip (s, i+1)
586611 Stop -> Stop
587612
613+{-# INLINE concatMap #-}
614+concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
615+concatMap f = concatMapM (return . f)
616+
617+{-# INLINE_NORMAL concatMapM #-}
618+concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
619+concatMapM f (Stream step state) = Stream step' (Left state)
620+ where
621+ step' gst (Left st) = do
622+ r <- step (rstState gst) st
623+ case r of
624+ Yield a s -> do
625+ b_stream <- f a
626+ return $ Skip (Right (b_stream, s))
627+ Skip s -> return $ Skip (Left s)
628+ Stop -> return Stop
629+
630+ step' _ (Right (Stream inner_step inner_st, st)) = do
631+ r <- inner_step defState inner_st
632+ case r of
633+ Yield b inner_s ->
634+ return $ Yield b (Right (Stream inner_step inner_s, st))
635+ Skip inner_s ->
636+ return $ Skip (Right (Stream inner_step inner_s, st))
637+ Stop -> return $ Skip (Left st)
638+
588639------------------------------------------------------------------------------
589640-- Substreams
590641------------------------------------------------------------------------------