{-# 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)