never executed always true always false
1 {-# LANGUAGE FlexibleInstances #-}
2 {-# LANGUAGE RankNTypes #-}
3 -- -----------------------------------------------------------------------------
4 --
5 -- (c) The University of Glasgow 2012
6 --
7 -- -----------------------------------------------------------------------------
8
9 -- | Monadic streams
10 module GHC.Data.Stream (
11 Stream(..), StreamS(..), runStream, yield, liftIO,
12 collect, consume, fromList,
13 map, mapM, mapAccumL_
14 ) where
15
16 import GHC.Prelude hiding (map,mapM)
17
18 import Control.Monad hiding (mapM)
19 import Control.Monad.IO.Class
20
21 -- |
22 -- @Stream m a b@ is a computation in some Monad @m@ that delivers a sequence
23 -- of elements of type @a@ followed by a result of type @b@.
24 --
25 -- More concretely, a value of type @Stream m a b@ can be run using @runStreamInternal@
26 -- in the Monad @m@, and it delivers either
27 --
28 -- * the final result: @Done b@, or
29 -- * @Yield a str@ where @a@ is the next element in the stream, and @str@
30 -- is the rest of the stream
31 -- * @Effect mstr@ where @mstr@ is some action running in @m@ which
32 -- generates the rest of the stream.
33 --
34 -- Stream is itself a Monad, and provides an operation 'yield' that
35 -- produces a new element of the stream. This makes it convenient to turn
36 -- existing monadic computations into streams.
37 --
38 -- The idea is that Stream is useful for making a monadic computation
39 -- that produces values from time to time. This can be used for
40 -- knitting together two complex monadic operations, so that the
41 -- producer does not have to produce all its values before the
42 -- consumer starts consuming them. We make the producer into a
43 -- Stream, and the consumer pulls on the stream each time it wants a
44 -- new value.
45 --
46 -- 'Stream' is implemented in the "yoneda" style for efficiency. By
47 -- representing a stream in this manner 'fmap' and '>>=' operations are
48 -- accumulated in the function parameters before being applied once when
49 -- the stream is destroyed. In the old implementation each usage of 'mapM'
50 -- and '>>=' would traverse the entire stream in order to apply the
51 -- substitution at the leaves.
52 --
53 -- The >>= operation for 'Stream' was a hot-spot in the ticky profile for
54 -- the "ManyConstructors" test which called the 'cg' function many times in
55 -- @StgToCmm.hs@
56 --
57 newtype Stream m a b =
58 Stream { runStreamInternal :: forall r' r .
59 (a -> m r') -- For fusing calls to `map` and `mapM`
60 -> (b -> StreamS m r' r) -- For fusing `>>=`
61 -> StreamS m r' r }
62
63 runStream :: Applicative m => Stream m r' r -> StreamS m r' r
64 runStream st = runStreamInternal st pure Done
65
66 data StreamS m a b = Yield a (StreamS m a b)
67 | Done b
68 | Effect (m (StreamS m a b))
69
70 instance Monad m => Functor (StreamS m a) where
71 fmap = liftM
72
73 instance Monad m => Applicative (StreamS m a) where
74 pure = Done
75 (<*>) = ap
76
77 instance Monad m => Monad (StreamS m a) where
78 a >>= k = case a of
79 Done r -> k r
80 Yield a s -> Yield a (s >>= k)
81 Effect m -> Effect (fmap (>>= k) m)
82
83 instance Functor (Stream f a) where
84 fmap = liftM
85
86 instance Applicative (Stream m a) where
87 pure a = Stream $ \_f g -> g a
88 (<*>) = ap
89
90 instance Monad (Stream m a) where
91 Stream m >>= k = Stream $ \f h -> m f (\a -> runStreamInternal (k a) f h)
92
93 instance MonadIO m => MonadIO (Stream m b) where
94 liftIO io = Stream $ \_f g -> Effect (g <$> liftIO io)
95
96 yield :: Monad m => a -> Stream m a ()
97 yield a = Stream $ \f rest -> Effect (flip Yield (rest ()) <$> f a)
98
99 -- | Turn a Stream into an ordinary list, by demanding all the elements.
100 collect :: Monad m => Stream m a () -> m [a]
101 collect str = go [] (runStream str)
102 where
103 go acc (Done ()) = return (reverse acc)
104 go acc (Effect m) = m >>= go acc
105 go acc (Yield a k) = go (a:acc) k
106
107 consume :: (Monad m, Monad n) => Stream m a b -> (forall a . m a -> n a) -> (a -> n ()) -> n b
108 consume str l f = go (runStream str)
109 where
110 go (Done r) = return r
111 go (Yield a p) = f a >> go p
112 go (Effect m) = l m >>= go
113
114 -- | Turn a list into a 'Stream', by yielding each element in turn.
115 fromList :: Monad m => [a] -> Stream m a ()
116 fromList = mapM_ yield
117
118 -- | Apply a function to each element of a 'Stream', lazily
119 map :: Monad m => (a -> b) -> Stream m a x -> Stream m b x
120 map f str = Stream $ \g h -> runStreamInternal str (g . f) h
121
122 -- | Apply a monadic operation to each element of a 'Stream', lazily
123 mapM :: Monad m => (a -> m b) -> Stream m a x -> Stream m b x
124 mapM f str = Stream $ \g h -> runStreamInternal str (g <=< f) h
125
126 -- | Note this is not very efficient because it traverses the whole stream
127 -- before rebuilding it, avoid using it if you can. mapAccumL used to
128 -- implemented but it wasn't used anywhere in the compiler and has similar
129 -- effiency problems.
130 mapAccumL_ :: forall m a b c r . Monad m => (c -> a -> m (c,b)) -> c -> Stream m a r
131 -> Stream m b (c, r)
132 mapAccumL_ f c str = Stream $ \f h -> go c f h (runStream str)
133
134 where
135 go :: c
136 -> (b -> m r')
137 -> ((c, r) -> StreamS m r' r1)
138 -> StreamS m a r
139 -> StreamS m r' r1
140 go c _f1 h1 (Done r) = h1 (c, r)
141 go c f1 h1 (Yield a p) = Effect (f c a >>= (\(c', b) -> f1 b
142 >>= \r' -> return $ Yield r' (go c' f1 h1 p)))
143 go c f1 h1 (Effect m) = Effect (go c f1 h1 <$> m)