on master help sign in
Add an unfold type to compose streams with better fusion
3b0f6dfHarendra Kumar
benchmark/BaseStreams.hs+1−1
@@ -103,7 +103,7 @@ main =
103 , benchIO "concatMapRepl4xN" D.concatMapRepl4xN
103 , benchIO "concatMapRepl4xN" D.concatMapRepl4xN
104 (D.sourceUnfoldrMN (D.value `div` 4))
104 (D.sourceUnfoldrMN (D.value `div` 4))
105 , benchIO "concatMapPureNxN" (D.concatMap 1) (D.sourceUnfoldrN D.value2)
105 , benchIO "concatMapPureNxN" (D.concatMap 1) (D.sourceUnfoldrN D.value2)
106− , benchIO "flattenRepl4xN" D.flattenRepl4xN
106+ , benchIO "concatMapURepl4xN" D.concatMapURepl4xN
107 (D.sourceUnfoldrMN (D.value `div` 4))
107 (D.sourceUnfoldrMN (D.value `div` 4))
108 , benchIO "intersperse" (D.intersperse 1) (D.sourceUnfoldrMN D.value2)
108 , benchIO "intersperse" (D.intersperse 1) (D.sourceUnfoldrMN D.value2)
109 , benchIO "interspersePure" (D.intersperse 1) (D.sourceUnfoldrN D.value2)
109 , benchIO "interspersePure" (D.intersperse 1) (D.sourceUnfoldrN D.value2)
benchmark/StreamDOps.hs+4−4
@@ -19,6 +19,7 @@ import Prelude
19import qualified Prelude as P
19import qualified Prelude as P
20
20
21import qualified Streamly.Streams.StreamD as S
21import qualified Streamly.Streams.StreamD as S
22+import qualified Streamly.Unfold as UF
22
23
23-- We try to keep the total number of iterations same irrespective of nesting
24-- We try to keep the total number of iterations same irrespective of nesting
24-- of the loops so that the overhead is easy to compare.
25-- of the loops so that the overhead is easy to compare.
@@ -270,10 +271,9 @@ zip src = transform $ S.zipWith (,) src src
270concatMapRepl4xN :: Monad m => Stream m Int -> m ()
271concatMapRepl4xN :: Monad m => Stream m Int -> m ()
271concatMapRepl4xN src = transform $ (S.concatMap (S.replicate 4) src)
272concatMapRepl4xN src = transform $ (S.concatMap (S.replicate 4) src)
272
273
273−{-# INLINE flattenRepl4xN #-}
274+{-# INLINE concatMapURepl4xN #-}
274−flattenRepl4xN :: Monad m => Stream m Int -> m ()
275+concatMapURepl4xN :: Monad m => Stream m Int -> m ()
275−flattenRepl4xN src = transform $
276+concatMapURepl4xN src = transform $ S.concatMapU (UF.replicateM 4) src
276− S.flatten S.replicateGen (\x -> return (x,4)) src
277
277
278-------------------------------------------------------------------------------
278-------------------------------------------------------------------------------
279-- Mixed Composition
279-- Mixed Composition
src/Streamly/Streams/StreamD.hs+36−10
@@ -63,8 +63,6 @@ module Streamly.Streams.StreamD
63 -- | Generate a monadic stream from a seed.
63 -- | Generate a monadic stream from a seed.
64 , repeat
64 , repeat
65 , replicate
65 , replicate
66− , replicateGen
67− , flatten
68 , replicateM
66 , replicateM
69 , fromIndices
67 , fromIndices
70 , fromIndicesM
68 , fromIndicesM
@@ -137,6 +135,7 @@ module Streamly.Streams.StreamD
137 -- ** Flattening nested streams
135 -- ** Flattening nested streams
138 , concatMapM
136 , concatMapM
139 , concatMap
137 , concatMap
138+ , concatMapU
140
139
141 -- ** Grouping
140 -- ** Grouping
142 , groupsOf
141 , groupsOf
@@ -284,6 +283,7 @@ import Streamly.Memory.Array.Types (Array(..))
284import Streamly.Fold.Types (Fold(..))
283import Streamly.Fold.Types (Fold(..))
285import Streamly.Pipe.Types (Pipe(..), PipeState(..))
284import Streamly.Pipe.Types (Pipe(..), PipeState(..))
286import Streamly.SVar (MonadAsync, defState, adaptState)
285import Streamly.SVar (MonadAsync, defState, adaptState)
286+import Streamly.Unfold.Types (Unfold(..))
287
287
288import Streamly.Streams.StreamD.Type
288import Streamly.Streams.StreamD.Type
289
289
@@ -386,14 +386,6 @@ replicateM n p = Stream step n
386 x <- p
386 x <- p
387 return $ Yield x (i - 1)
387 return $ Yield x (i - 1)
388
388
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−
397{-# INLINE_NORMAL replicate #-}
389{-# INLINE_NORMAL replicate #-}
398replicate :: Monad m => Int -> a -> Stream m a
390replicate :: Monad m => Int -> a -> Stream m a
399replicate n x = replicateM n (return x)
391replicate n x = replicateM n (return x)
@@ -1702,6 +1694,40 @@ stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1702mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
1694mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
1703mapM_ m = drain . mapM m
1695mapM_ m = drain . mapM m
1704
1696
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+
1705------------------------------------------------------------------------------
1731------------------------------------------------------------------------------
1706-- Exceptions
1732-- Exceptions
1707------------------------------------------------------------------------------
1733------------------------------------------------------------------------------
src/Streamly/Streams/StreamD/Type.hs−33
@@ -42,7 +42,6 @@ module Streamly.Streams.StreamD.Type
42 , yieldM
42 , yieldM
43 , concatMap
43 , concatMap
44 , concatMapM
44 , concatMapM
45− , flatten
46
45
47 , foldrT
46 , foldrT
48 , foldrM
47 , foldrM
@@ -222,38 +221,6 @@ concatMap f = concatMapM (return . f)
222-- {-# RULES "concatMap Array.toStreamD"
221-- {-# RULES "concatMap Array.toStreamD"
223-- concatMap Array.toStreamD = Array.flattenArray #-}
222-- concatMap Array.toStreamD = Array.flattenArray #-}
224
223
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−
257-- | Create a singleton 'Stream' from a pure value.
224-- | Create a singleton 'Stream' from a pure value.
258{-# INLINE_NORMAL yield #-}
225{-# INLINE_NORMAL yield #-}
259yield :: Monad m => a -> Stream m a
226yield :: 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 #-}
8+
9+#include "inline.hs"
10+
11+-- |
12+-- Module : Streamly.Unfold
13+-- Copyright : (c) 2019 Composewell Technologies
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 #-}
4+
5+-- |
6+-- Module : Streamly.Unfold.Types
7+-- Copyright : (c) 2019 Composewell Technologies
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
207 -- Pure streams
207 -- Pure streams
208 , Streamly.List
208 , Streamly.List
209
209
210+ -- Unfolds
211+ , Streamly.Unfold.Types
212+ , Streamly.Unfold
213+
210 -- Folds
214 -- Folds
211 , Streamly.Fold.Types
215 , Streamly.Fold.Types
212 , Streamly.Sink.Types
216 , Streamly.Sink.Types