687fb9a · streamly · git.iko.soy

benchmark/BaseStreams.hs+2

compare with master

unifiedsplit

@@ -103,6 +103,8 @@ main =

103 , benchIO "concatMapRepl4xN" D.concatMapRepl4xN

104 (D.sourceUnfoldrMN (D.value `div` 4))

105 , benchIO "concatMapPureNxN" (D.concatMap 1) (D.sourceUnfoldrN D.value2)

106+ , benchIO "flattenRepl4xN" D.flattenRepl4xN

107+ (D.sourceUnfoldrMN (D.value `div` 4))

106 , benchIO "intersperse" (D.intersperse 1) (D.sourceUnfoldrMN D.value2)

108 , benchIO "intersperse" (D.intersperse 1) (D.sourceUnfoldrMN D.value2)

107 , benchIO "interspersePure" (D.intersperse 1) (D.sourceUnfoldrN D.value2)

109 , benchIO "interspersePure" (D.intersperse 1) (D.sourceUnfoldrN D.value2)

108 -- , benchIO "foldrS" (D.foldrS 1) D.sourceUnfoldrM

110 -- , benchIO "foldrS" (D.foldrS 1) D.sourceUnfoldrM

benchmark/StreamDOps.hs+5

@@ -270,6 +270,11 @@ zip src = transform $ S.zipWith (,) src src

270concatMapRepl4xN :: Monad m => Stream m Int -> m ()

271concatMapRepl4xN src = transform $ (S.concatMap (S.replicate 4) src)

272

273+{-# INLINE flattenRepl4xN #-}

274+flattenRepl4xN :: Monad m => Stream m Int -> m ()

275+flattenRepl4xN src = transform $

276+ S.flatten S.replicateGen (\x -> return (x,4)) src

277+

273-------------------------------------------------------------------------------

278-------------------------------------------------------------------------------

274-- Mixed Composition

279-- Mixed Composition

275-------------------------------------------------------------------------------

280-------------------------------------------------------------------------------

src/Streamly/Streams/StreamD.hs+10

@@ -58,6 +58,8 @@ module Streamly.Streams.StreamD

58 -- | Generate a monadic stream from a seed.

59 , repeat

60 , replicate

61+ , replicateGen

62+ , flatten

61 , replicateM

63 , replicateM

62 , fromIndices

64 , fromIndices

63 , fromIndicesM

65 , fromIndicesM

@@ -376,6 +378,14 @@ replicateM n p = Stream step n

376 x <- p

378 x <- p

377 return $ Yield x (i - 1)

379 return $ Yield x (i - 1)

378

380

381+-- | For use with flatten

382+{-# INLINE replicateGen #-}

383+replicateGen :: Monad m => (a, Int) -> m (Step (a, Int) a)

384+replicateGen (x, i) = return $

385+ if i <= 0

386+ then Stop

387+ else Yield x (x, (i - 1))

388+

379{-# INLINE_NORMAL replicate #-}

389{-# INLINE_NORMAL replicate #-}

380replicate :: Monad m => Int -> a -> Stream m a

390replicate :: Monad m => Int -> a -> Stream m a

381replicate n x = replicateM n (return x)

391replicate n x = replicateM n (return x)

src/Streamly/Streams/StreamD/Type.hs+33

@@ -42,6 +42,7 @@ module Streamly.Streams.StreamD.Type

42 , yieldM

43 , concatMap

44 , concatMapM

45+ , flatten

45

46

46 , foldrT

47 , foldrT

47 , foldrM

48 , foldrM

@@ -221,6 +222,38 @@ concatMap f = concatMapM (return . f)

221-- {-# RULES "concatMap Array.toStreamD"

222-- {-# RULES "concatMap Array.toStreamD"

222-- concatMap Array.toStreamD = Array.flattenArray #-}

223-- concatMap Array.toStreamD = Array.flattenArray #-}

223

224

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+

224-- | Create a singleton 'Stream' from a pure value.

257-- | Create a singleton 'Stream' from a pure value.

225{-# INLINE_NORMAL yield #-}

258{-# INLINE_NORMAL yield #-}

226yield :: Monad m => a -> Stream m a

259yield :: Monad m => a -> Stream m a