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)