3b0f6df · streamly · git.iko.soy

git.iko.soy/streamly

Add an unfold type to compose streams with better fusion

commit3b0f6df5cf30deafde1c6697a391d120cc60d70d authorHarendra Kumar <harendra.kumar@gmail.com> dateSat 17 Aug 2019, 14:04 (7 years ago) parentfeb44d8

7 files changed+144−48

compare with master

unifiedsplit

benchmark/BaseStreams.hs+1−1 benchmark/StreamDOps.hs+4−4 src/Streamly/Streams/StreamD.hs+36−10 src/Streamly/Streams/StreamD/Type.hs−33 src/Streamly/Unfold.hs+68 src/Streamly/Unfold/Types.hs+31 streamly.cabal+4

benchmark/BaseStreams.hs+1−1

@@ -103,7 +103,7 @@ main =

103103 , benchIO "concatMapRepl4xN" D.concatMapRepl4xN

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

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

106− , benchIO "flattenRepl4xN" D.flattenRepl4xN

106+ , benchIO "concatMapURepl4xN" D.concatMapURepl4xN

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

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

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

benchmark/StreamDOps.hs+4−4

@@ -19,6 +19,7 @@ import Prelude

1919import qualified Prelude as P

2020

2121import qualified Streamly.Streams.StreamD as S

22+import qualified Streamly.Unfold as UF

2223

2324-- We try to keep the total number of iterations same irrespective of nesting

2425-- of the loops so that the overhead is easy to compare.

@@ -270,10 +271,9 @@ zip src = transform $ S.zipWith (,) src src

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

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

272273

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

274+{-# INLINE concatMapURepl4xN #-}

275+concatMapURepl4xN :: Monad m => Stream m Int -> m ()

276+concatMapURepl4xN src = transform $ S.concatMapU (UF.replicateM 4) src

277277

278278-------------------------------------------------------------------------------

279279-- Mixed Composition

src/Streamly/Streams/StreamD.hs+36−10

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

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

6464 , repeat

6565 , replicate

66− , replicateGen

67− , flatten

6866 , replicateM

6967 , fromIndices

7068 , fromIndicesM

@@ -137,6 +135,7 @@ module Streamly.Streams.StreamD

137135 -- ** Flattening nested streams

138136 , concatMapM

139137 , concatMap

138+ , concatMapU

140139

141140 -- ** Grouping

142141 , groupsOf

@@ -284,6 +283,7 @@ import Streamly.Memory.Array.Types (Array(..))

284283import Streamly.Fold.Types (Fold(..))

285284import Streamly.Pipe.Types (Pipe(..), PipeState(..))

286285import Streamly.SVar (MonadAsync, defState, adaptState)

286+import Streamly.Unfold.Types (Unfold(..))

287287

288288import Streamly.Streams.StreamD.Type

289289

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

386386 x <- p

387387 return $ Yield x (i - 1)

388388

389−-- | For use with flatten

390−{-# INLINE replicateGen #-}

391−replicateGen :: Monad m => (a, Int) -> m (Step (a, Int) a)

392−replicateGen (x, i) = return $

393− if i <= 0

394− then Stop

395− else Yield x (x, (i - 1))

396−

397389{-# INLINE_NORMAL replicate #-}

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

399391replicate n x = replicateM n (return x)

@@ -1702,6 +1694,40 @@ stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)

17021694mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()

17031695mapM_ m = drain . mapM m

17041696

1697+-------------------------------------------------------------------------------

1698+-- Stream transformations using Unfolds

1699+-------------------------------------------------------------------------------

1700+

1701+-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements

1702+-- to streams and then flattens the generated streams into a single output

1703+-- stream.

1704+

1705+-- This is like 'concatMap' but uses an unfold with an explicit state to

1706+-- generate the stream instead of a 'Stream' type generator. This allows better

1707+-- optimization via fusion. This can be many times more efficient than

1708+-- 'concatMap'.

1709+

1710+{-# INLINE_NORMAL concatMapU #-}

1711+concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b

1712+concatMapU (Unfold istep inject) (Stream ostep u) = Stream step (Left u)

1713+ where

1714+ {-# INLINE_LATE step #-}

1715+ step gst (Left t) = do

1716+ r <- ostep (adaptState gst) t

1717+ case r of

1718+ Yield a t' -> do

1719+ s <- inject a

1720+ s `seq` return (Skip (Right (s, t')))

1721+ Skip t' -> return $ Skip (Left t')

1722+ Stop -> return $ Stop

1723+

1724+ step _ (Right (s, t)) = do

1725+ r <- istep s

1726+ return $ case r of

1727+ Yield x s' -> Yield x (Right (s', t))

1728+ Skip s' -> Skip (Right (s', t))

1729+ Stop -> Skip (Left t)

1730+

17051731------------------------------------------------------------------------------

17061732-- Exceptions

17071733------------------------------------------------------------------------------

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

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

4242 , yieldM

4343 , concatMap

4444 , concatMapM

45− , flatten

4645

4746 , foldrT

4847 , foldrM

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

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

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

224223

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−

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

258225{-# INLINE_NORMAL yield #-}

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

src/Streamly/Unfold.hsnew file+68

@@ -0,0 +1,68 @@

1+{-# LANGUAGE BangPatterns #-}

2+{-# LANGUAGE CPP #-}

3+{-# LANGUAGE ExistentialQuantification #-}

4+{-# LANGUAGE FlexibleContexts #-}

5+{-# LANGUAGE RankNTypes #-}

6+{-# LANGUAGE RecordWildCards #-}

7+{-# LANGUAGE ScopedTypeVariables #-}

9+#include "inline.hs"

10+

11+-- |

12+-- Module : Streamly.Unfold

14+-- License : BSD3

15+-- Maintainer : harendra.kumar@gmail.com

16+-- Stability : experimental

17+-- Portability : GHC

18+--

19+

20+module Streamly.Unfold

21+ (

22+ -- * Unfold Type

23+ -- |

24+ -- A 'Fold' can be run over a seed using the 'unfold' combinator.

25+ --

26+ -- >>> unfold UF.replicateM (putStrLn "hello") 10

27+

28+ Unfold

29+ , unfold

30+

31+ -- ** Unfolds

32+ , replicateM

33+ )

34+where

35+

36+import Streamly.Streams.StreamD.Type (Stream(..), Step(..))

37+import Streamly.Unfold.Types (Unfold(..))

38+

39+-------------------------------------------------------------------------------

40+-- Running unfolds

41+-------------------------------------------------------------------------------

42+

43+{-# INLINE_NORMAL unfold #-}

44+unfold :: Monad m => a -> Unfold m a b -> Stream m b

45+unfold seed (Unfold ustep inject) = Stream step Nothing

46+ where

47+ {-# INLINE_LATE step #-}

48+ step _ Nothing = inject seed >>= return . Skip . Just

49+ step _ (Just st) = do

50+ r <- ustep st

51+ return $ case r of

52+ Yield x s -> Yield x (Just s)

53+ Skip s -> Skip (Just s)

54+ Stop -> Stop

55+

56+-------------------------------------------------------------------------------

57+-- Unfolds

58+-------------------------------------------------------------------------------

59+

60+{-# INLINE replicateM #-}

61+replicateM :: Monad m => Int -> Unfold m a a

62+replicateM n = Unfold step inject

63+ where

64+ inject x = return (x, n)

65+ step (x, i) = return $

66+ if i <= 0

67+ then Stop

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

src/Streamly/Unfold/Types.hsnew file+31

@@ -0,0 +1,31 @@

1+{-# LANGUAGE CPP #-}

2+{-# LANGUAGE ExistentialQuantification #-}

3+{-# LANGUAGE FlexibleContexts #-}

5+-- |

6+-- Module : Streamly.Unfold.Types

8+-- License : BSD3

9+-- Maintainer : harendra.kumar@gmail.com

10+-- Stability : experimental

11+-- Portability : GHC

12+

13+module Streamly.Unfold.Types

14+ ( Unfold (..)

15+ )

16+where

17+

18+import Streamly.Streams.StreamD.Type (Step(..))

19+

20+------------------------------------------------------------------------------

21+-- Monadic Unfolds

22+------------------------------------------------------------------------------

23+

24+-- | An @Unfold m a b@ is a generator of a stream of values of type @b@ from a

25+-- seed of type 'a' in 'Monad' @m@.

26+--

27+-- @since 0.7.0

28+

29+data Unfold m a b =

30+ -- | @Unfold step inject@

31+ forall s. Unfold (s -> m (Step s b)) (a -> m s)

streamly.cabal+4

@@ -207,6 +207,10 @@ library

207207 -- Pure streams

208208 , Streamly.List

209209

210+ -- Unfolds

211+ , Streamly.Unfold.Types

212+ , Streamly.Unfold

213+

210214 -- Folds

211215 , Streamly.Fold.Types

212216 , Streamly.Sink.Types