summaryrefslogtreecommitdiff
path: root/compiler/GHC/Data/Stream.hs
blob: 7996ee7343dc106bd105a474e48ab2587855630d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
-- -----------------------------------------------------------------------------
--
-- (c) The University of Glasgow 2012
--
-- -----------------------------------------------------------------------------

-- | Monadic streams
module GHC.Data.Stream (
    Stream(..), yield, liftIO,
    collect, collect_, consume, fromList,
    map, mapM, mapAccumL, mapAccumL_
  ) where

import GHC.Prelude hiding (map,mapM)

import Control.Monad hiding (mapM)

-- |
-- @Stream m a b@ is a computation in some Monad @m@ that delivers a sequence
-- of elements of type @a@ followed by a result of type @b@.
--
-- More concretely, a value of type @Stream m a b@ can be run using @runStream@
-- in the Monad @m@, and it delivers either
--
--  * the final result: @Left b@, or
--  * @Right (a,str)@, where @a@ is the next element in the stream, and @str@
--    is a computation to get the rest of the stream.
--
-- Stream is itself a Monad, and provides an operation 'yield' that
-- produces a new element of the stream.  This makes it convenient to turn
-- existing monadic computations into streams.
--
-- The idea is that Stream is useful for making a monadic computation
-- that produces values from time to time.  This can be used for
-- knitting together two complex monadic operations, so that the
-- producer does not have to produce all its values before the
-- consumer starts consuming them.  We make the producer into a
-- Stream, and the consumer pulls on the stream each time it wants a
-- new value.
--
newtype Stream m a b = Stream { runStream :: m (Either b (a, Stream m a b)) }

instance Monad f => Functor (Stream f a) where
  fmap = liftM

instance Monad m => Applicative (Stream m a) where
  pure a = Stream (return (Left a))
  (<*>) = ap

instance Monad m => Monad (Stream m a) where

  Stream m >>= k = Stream $ do
                r <- m
                case r of
                  Left b        -> runStream (k b)
                  Right (a,str) -> return (Right (a, str >>= k))

yield :: Monad m => a -> Stream m a ()
yield a = Stream (return (Right (a, return ())))

liftIO :: IO a -> Stream IO b a
liftIO io = Stream $ io >>= return . Left

-- | Turn a Stream into an ordinary list, by demanding all the elements.
collect :: Monad m => Stream m a () -> m [a]
collect str = go str []
 where
  go str acc = do
    r <- runStream str
    case r of
      Left () -> return (reverse acc)
      Right (a, str') -> go str' (a:acc)

-- | Turn a Stream into an ordinary list, by demanding all the elements.
collect_ :: Monad m => Stream m a r -> m ([a], r)
collect_ str = go str []
 where
  go str acc = do
    r <- runStream str
    case r of
      Left r -> return (reverse acc, r)
      Right (a, str') -> go str' (a:acc)

consume :: Monad m => Stream m a b -> (a -> m ()) -> m b
consume str f = do
    r <- runStream str
    case r of
      Left ret -> return ret
      Right (a, str') -> do
        f a
        consume str' f

-- | Turn a list into a 'Stream', by yielding each element in turn.
fromList :: Monad m => [a] -> Stream m a ()
fromList = mapM_ yield

-- | Apply a function to each element of a 'Stream', lazily
map :: Monad m => (a -> b) -> Stream m a x -> Stream m b x
map f str = Stream $ do
   r <- runStream str
   case r of
     Left x -> return (Left x)
     Right (a, str') -> return (Right (f a, map f str'))

-- | Apply a monadic operation to each element of a 'Stream', lazily
mapM :: Monad m => (a -> m b) -> Stream m a x -> Stream m b x
mapM f str = Stream $ do
   r <- runStream str
   case r of
     Left x -> return (Left x)
     Right (a, str') -> do
        b <- f a
        return (Right (b, mapM f str'))

-- | analog of the list-based 'mapAccumL' on Streams.  This is a simple
-- way to map over a Stream while carrying some state around.
mapAccumL :: Monad m => (c -> a -> m (c,b)) -> c -> Stream m a ()
          -> Stream m b c
mapAccumL f c str = Stream $ do
  r <- runStream str
  case r of
    Left  () -> return (Left c)
    Right (a, str') -> do
      (c',b) <- f c a
      return (Right (b, mapAccumL f c' str'))

mapAccumL_ :: Monad m => (c -> a -> m (c,b)) -> c -> Stream m a r
           -> Stream m b (c, r)
mapAccumL_ f c str = Stream $ do
  r <- runStream str
  case r of
    Left  r -> return (Left (c, r))
    Right (a, str') -> do
      (c',b) <- f c a
      return (Right (b, mapAccumL_ f c' str'))