@@ -42,6 +42,7 @@ module Streamly.Streams.StreamD.Type
4242 , yieldM
4343 , concatMap
4444 , concatMapM
45+ , flatten
4546
4647 , foldrT
4748 , foldrM
@@ -221,6 +222,38 @@ concatMap f = concatMapM (return . f)
221222-- {-# RULES "concatMap Array.toStreamD"
222223-- concatMap Array.toStreamD = Array.flattenArray #-}
223224
225+-- | @flatten step inject stream@ uses the unfold @step@ function to generate
226+-- streams from the input stream and flatten them into a single output stream.
227+-- The @inject@ function is used to wrap the input element into the state of
228+-- unfold step.
229+--
230+-- This is like 'concatMap' but uses an unfold with an explicit state to
231+-- generate the stream instead of a 'Stream' type generator. This allows better
232+-- optimization via fusion. This can be many times more efficient than
233+-- 'concatMap'.
234+
235+{-# INLINE_NORMAL flatten #-}
236+flatten :: Monad m
237+ => (s -> m (Step s b)) -> (a -> m s) -> Stream m a -> Stream m b
238+flatten istep inject (Stream ostep u) = Stream step (Left u)
239+ where
240+ {-# INLINE_LATE step #-}
241+ step gst (Left t) = do
242+ r <- ostep (adaptState gst) t
243+ case r of
244+ Yield a t' -> do
245+ s <- inject a
246+ s `seq` return (Skip (Right (s, t')))
247+ Skip t' -> return $ Skip (Left t')
248+ Stop -> return $ Stop
249+
250+ step _ (Right (s, t)) = do
251+ r <- istep s
252+ return $ case r of
253+ Yield x s' -> Yield x (Right (s', t))
254+ Skip s' -> Skip (Right (s', t))
255+ Stop -> Skip (Left t)
256+
224257-- | Create a singleton 'Stream' from a pure value.
225258{-# INLINE_NORMAL yield #-}
226259yield :: Monad m => a -> Stream m a