    {-# LANGUAGE FlexibleInstances #-}
    {-# LANGUAGE RankNTypes #-}
    -- -----------------------------------------------------------------------------
    --
    -- (c) The University of Glasgow 2012
    --
    -- -----------------------------------------------------------------------------

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

    import GHC.Prelude hiding (map,mapM)

    import Control.Monad hiding (mapM)
    import Control.Monad.IO.Class

    -- |
    -- @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 @runStreamInternal@
    -- in the Monad @m@, and it delivers either
    --
    --  * the final result: @Done b@, or
    --  * @Yield a str@ where @a@ is the next element in the stream, and @str@
    --     is the rest of the stream
    --  * @Effect mstr@ where @mstr@ is some action running in @m@ which
    --  generates 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.
    --
    -- 'Stream' is implemented in the "yoneda" style for efficiency. By
    -- representing a stream in this manner 'fmap' and '>>=' operations are
    -- accumulated in the function parameters before being applied once when
    -- the stream is destroyed. In the old implementation each usage of 'mapM'
    -- and '>>=' would traverse the entire stream in order to apply the
    -- substitution at the leaves.
    --
    -- The >>= operation for 'Stream' was a hot-spot in the ticky profile for
    -- the "ManyConstructors" test which called the 'cg' function many times in
    -- @StgToCmm.hs@
    --
    newtype Stream m a b =
              Stream { forall (m :: * -> *) a b.
    Stream m a b
    -> forall r' r.
       (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r
    runStreamInternal :: forall r' r .
                                            (a -> m r') -- For fusing calls to `map` and `mapM`
                                         -> (b -> StreamS m r' r)  -- For fusing `>>=`
                                         -> StreamS m r' r }

    runStream :: Applicative m => Stream m r' r -> StreamS m r' r
    runStream :: forall (m :: * -> *) r' r.
    Applicative m =>
    Stream m r' r -> StreamS m r' r
    runStream Stream m r' r
    st = forall (m :: * -> *) a b.
    Stream m a b
    -> forall r' r.
       (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r
    runStreamInternal Stream m r' r
    st forall (f :: * -> *) a. Applicative f => a -> f a
    pure forall (m :: * -> *) a b. b -> StreamS m a b
    Done

    data StreamS m a b = Yield a (StreamS m a b)
                       | Done b
                       | Effect (m (StreamS m a b))

    instance Monad m => Functor (StreamS m a) where
      fmap :: forall a b. (a -> b) -> StreamS m a a -> StreamS m a b
    fmap = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
    liftM

    instance Monad m => Applicative (StreamS m a) where
      pure :: forall a. a -> StreamS m a a
    pure = forall (m :: * -> *) a b. b -> StreamS m a b
    Done
      <*> :: forall a b. StreamS m a (a -> b) -> StreamS m a a -> StreamS m a b
    (<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
    ap

    instance Monad m => Monad (StreamS m a) where
      StreamS m a a
    a >>= :: forall a b. StreamS m a a -> (a -> StreamS m a b) -> StreamS m a b
    >>= a -> StreamS m a b
    k = case StreamS m a a
    a of
                          Done a
    r -> a -> StreamS m a b
    k a
    r
                          Yield a
    a StreamS m a a
    s -> forall (m :: * -> *) a b. a -> StreamS m a b -> StreamS m a b
    Yield a
    a (StreamS m a a
    s forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= a -> StreamS m a b
    k)
                          Effect m (StreamS m a a)
    m -> forall (m :: * -> *) a b. m (StreamS m a b) -> StreamS m a b
    Effect (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= a -> StreamS m a b
    k) m (StreamS m a a)
    m)

    instance Functor (Stream f a) where
      fmap :: forall a b. (a -> b) -> Stream f a a -> Stream f a b
    fmap = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
    liftM

    instance Applicative (Stream m a) where
      pure :: forall a. a -> Stream m a a
    pure a
    a = forall (m :: * -> *) a b.
    (forall r' r.
     (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r)
    -> Stream m a b
    Stream forall a b. (a -> b) -> a -> b
    $ \a -> m r'
    _f a -> StreamS m r' r
    g -> a -> StreamS m r' r
    g a
    a
      <*> :: forall a b. Stream m a (a -> b) -> Stream m a a -> Stream m a b
    (<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
    ap

    instance Monad (Stream m a) where
      Stream forall r' r. (a -> m r') -> (a -> StreamS m r' r) -> StreamS m r' r
    m >>= :: forall a b. Stream m a a -> (a -> Stream m a b) -> Stream m a b
    >>= a -> Stream m a b
    k = forall (m :: * -> *) a b.
    (forall r' r.
     (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r)
    -> Stream m a b
    Stream forall a b. (a -> b) -> a -> b
    $ \a -> m r'
    f b -> StreamS m r' r
    h -> forall r' r. (a -> m r') -> (a -> StreamS m r' r) -> StreamS m r' r
    m a -> m r'
    f (\a
    a -> forall (m :: * -> *) a b.
    Stream m a b
    -> forall r' r.
       (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r
    runStreamInternal (a -> Stream m a b
    k a
    a) a -> m r'
    f b -> StreamS m r' r
    h)

    instance MonadIO m => MonadIO (Stream m b) where
      liftIO :: forall a. IO a -> Stream m b a
    liftIO IO a
    io = forall (m :: * -> *) a b.
    (forall r' r.
     (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r)
    -> Stream m a b
    Stream forall a b. (a -> b) -> a -> b
    $ \b -> m r'
    _f a -> StreamS m r' r
    g -> forall (m :: * -> *) a b. m (StreamS m a b) -> StreamS m a b
    Effect (a -> StreamS m r' r
    g forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO IO a
    io)

    yield :: Monad m => a -> Stream m a ()
    yield :: forall (m :: * -> *) a. Monad m => a -> Stream m a ()
    yield a
    a = forall (m :: * -> *) a b.
    (forall r' r.
     (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r)
    -> Stream m a b
    Stream forall a b. (a -> b) -> a -> b
    $ \a -> m r'
    f () -> StreamS m r' r
    rest -> forall (m :: * -> *) a b. m (StreamS m a b) -> StreamS m a b
    Effect (forall a b c. (a -> b -> c) -> b -> a -> c
    flip forall (m :: * -> *) a b. a -> StreamS m a b -> StreamS m a b
    Yield (() -> StreamS m r' r
    rest ())  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> a -> m r'
    f a
    a)

    -- | Turn a Stream into an ordinary list, by demanding all the elements.
    collect :: Monad m => Stream m a () -> m [a]
    collect :: forall (m :: * -> *) a. Monad m => Stream m a () -> m [a]
    collect Stream m a ()
    str = forall {m :: * -> *} {a}. Monad m => [a] -> StreamS m a () -> m [a]
    go [] (forall (m :: * -> *) r' r.
    Applicative m =>
    Stream m r' r -> StreamS m r' r
    runStream Stream m a ()
    str)
     where
      go :: [a] -> StreamS m a () -> m [a]
    go [a]
    acc (Done ()) = forall (m :: * -> *) a. Monad m => a -> m a
    return (forall a. [a] -> [a]
    reverse [a]
    acc)
      go [a]
    acc (Effect m (StreamS m a ())
    m) = m (StreamS m a ())
    m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= [a] -> StreamS m a () -> m [a]
    go [a]
    acc
      go [a]
    acc (Yield a
    a StreamS m a ()
    k) = [a] -> StreamS m a () -> m [a]
    go (a
    aforall a. a -> [a] -> [a]
    :[a]
    acc) StreamS m a ()
    k

    consume :: (Monad m, Monad n) => Stream m a b -> (forall a . m a -> n a) -> (a -> n ()) -> n b
    consume :: forall (m :: * -> *) (n :: * -> *) a b.
    (Monad m, Monad n) =>
    Stream m a b -> (forall a. m a -> n a) -> (a -> n ()) -> n b
    consume Stream m a b
    str forall a. m a -> n a
    l a -> n ()
    f = StreamS m a b -> n b
    go (forall (m :: * -> *) r' r.
    Applicative m =>
    Stream m r' r -> StreamS m r' r
    runStream Stream m a b
    str)
      where
        go :: StreamS m a b -> n b
    go (Done b
    r) = forall (m :: * -> *) a. Monad m => a -> m a
    return b
    r
        go (Yield a
    a StreamS m a b
    p) = a -> n ()
    f a
    a forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> StreamS m a b -> n b
    go StreamS m a b
    p
        go (Effect m (StreamS m a b)
    m)  = forall a. m a -> n a
    l m (StreamS m a b)
    m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= StreamS m a b -> n b
    go

    -- | Turn a list into a 'Stream', by yielding each element in turn.
    fromList :: Monad m => [a] -> Stream m a ()
    fromList :: forall (m :: * -> *) a. Monad m => [a] -> Stream m a ()
    fromList = forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ forall (m :: * -> *) a. Monad m => a -> Stream m a ()
    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 :: forall (m :: * -> *) a b x.
    Monad m =>
    (a -> b) -> Stream m a x -> Stream m b x
    map a -> b
    f Stream m a x
    str = forall (m :: * -> *) a b.
    (forall r' r.
     (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r)
    -> Stream m a b
    Stream forall a b. (a -> b) -> a -> b
    $ \b -> m r'
    g x -> StreamS m r' r
    h -> forall (m :: * -> *) a b.
    Stream m a b
    -> forall r' r.
       (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r
    runStreamInternal Stream m a x
    str (b -> m r'
    g forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> b
    f) x -> StreamS m r' r
    h

    -- | 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 :: forall (m :: * -> *) a b x.
    Monad m =>
    (a -> m b) -> Stream m a x -> Stream m b x
    mapM a -> m b
    f Stream m a x
    str = forall (m :: * -> *) a b.
    (forall r' r.
     (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r)
    -> Stream m a b
    Stream forall a b. (a -> b) -> a -> b
    $ \b -> m r'
    g x -> StreamS m r' r
    h -> forall (m :: * -> *) a b.
    Stream m a b
    -> forall r' r.
       (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r
    runStreamInternal Stream m a x
    str (b -> m r'
    g forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< a -> m b
    f) x -> StreamS m r' r
    h

    -- | Note this is not very efficient because it traverses the whole stream
    -- before rebuilding it, avoid using it if you can. mapAccumL used to
    -- implemented but it wasn't used anywhere in the compiler and has similar
    -- effiency problems.
    mapAccumL_ :: forall m a b c r . Monad m => (c -> a -> m (c,b)) -> c -> Stream m a r
               -> Stream m b (c, r)
    mapAccumL_ :: forall (m :: * -> *) a b c r.
    Monad m =>
    (c -> a -> m (c, b)) -> c -> Stream m a r -> Stream m b (c, r)
    mapAccumL_ c -> a -> m (c, b)
    f c
    c Stream m a r
    str = forall (m :: * -> *) a b.
    (forall r' r.
     (a -> m r') -> (b -> StreamS m r' r) -> StreamS m r' r)
    -> Stream m a b
    Stream forall a b. (a -> b) -> a -> b
    $ \b -> m r'
    f (c, r) -> StreamS m r' r
    h -> forall r' r1.
    c
    -> (b -> m r')
    -> ((c, r) -> StreamS m r' r1)
    -> StreamS m a r
    -> StreamS m r' r1
    go c
    c b -> m r'
    f (c, r) -> StreamS m r' r
    h (forall (m :: * -> *) r' r.
    Applicative m =>
    Stream m r' r -> StreamS m r' r
    runStream Stream m a r
    str)

      where
        go :: c
                 -> (b -> m r')
                 -> ((c, r) -> StreamS m r' r1)
                 -> StreamS m a r
                 -> StreamS m r' r1
        go :: forall r' r1.
    c
    -> (b -> m r')
    -> ((c, r) -> StreamS m r' r1)
    -> StreamS m a r
    -> StreamS m r' r1
    go c
    c b -> m r'
    _f1 (c, r) -> StreamS m r' r1
    h1 (Done r
    r) = (c, r) -> StreamS m r' r1
    h1 (c
    c, r
    r)
        go c
    c b -> m r'
    f1 (c, r) -> StreamS m r' r1
    h1 (Yield a
    a StreamS m a r
    p) = forall (m :: * -> *) a b. m (StreamS m a b) -> StreamS m a b
    Effect (c -> a -> m (c, b)
    f c
    c a
    a forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= (\(c
    c', b
    b) -> b -> m r'
    f1 b
    b
                                               forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \r'
    r' -> forall (m :: * -> *) a. Monad m => a -> m a
    return forall a b. (a -> b) -> a -> b
    $ forall (m :: * -> *) a b. a -> StreamS m a b -> StreamS m a b
    Yield r'
    r' (forall r' r1.
    c
    -> (b -> m r')
    -> ((c, r) -> StreamS m r' r1)
    -> StreamS m a r
    -> StreamS m r' r1
    go c
    c' b -> m r'
    f1 (c, r) -> StreamS m r' r1
    h1 StreamS m a r
    p)))
        go c
    c b -> m r'
    f1 (c, r) -> StreamS m r' r1
    h1 (Effect m (StreamS m a r)
    m) = forall (m :: * -> *) a b. m (StreamS m a b) -> StreamS m a b
    Effect (forall r' r1.
    c
    -> (b -> m r')
    -> ((c, r) -> StreamS m r' r1)
    -> StreamS m a r
    -> StreamS m r' r1
    go c
    c b -> m r'
    f1 (c, r) -> StreamS m r' r1
    h1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> m (StreamS m a r)
    m)