-- -----------------------------------------------------------------------------
--
-- (c) The University of Glasgow 2012
--
-- Monadic streams
--
-- -----------------------------------------------------------------------------
module Stream (
    Stream(..), yield, liftIO,
    collect, collect_, consume, fromList,
    Stream.map, Stream.mapM, Stream.mapAccumL, Stream.mapAccumL_
  ) where

import GhcPrelude

import Control.Monad

-- |
-- @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 { Stream m a b -> m (Either b (a, Stream m a b))
runStream :: m (Either b (a, Stream m a b)) }

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

instance Monad m => Applicative (Stream m a) where
  pure :: a -> Stream m a a
pure a
a = m (Either a (a, Stream m a a)) -> Stream m a a
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (Either a (a, Stream m a a) -> m (Either a (a, Stream m a a))
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Either a (a, Stream m a a)
forall a b. a -> Either a b
Left a
a))
  <*> :: Stream m a (a -> b) -> Stream m a a -> Stream m 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 m => Monad (Stream m a) where

  Stream m (Either a (a, Stream m a a))
m >>= :: Stream m a a -> (a -> Stream m a b) -> Stream m a b
>>= a -> Stream m a b
k = m (Either b (a, Stream m a b)) -> Stream m a b
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (m (Either b (a, Stream m a b)) -> Stream m a b)
-> m (Either b (a, Stream m a b)) -> Stream m a b
forall a b. (a -> b) -> a -> b
$ do
                Either a (a, Stream m a a)
r <- m (Either a (a, Stream m a a))
m
                case Either a (a, Stream m a a)
r of
                  Left a
b        -> Stream m a b -> m (Either b (a, Stream m a b))
forall (m :: * -> *) a b.
Stream m a b -> m (Either b (a, Stream m a b))
runStream (a -> Stream m a b
k a
b)
                  Right (a
a,Stream m a a
str) -> Either b (a, Stream m a b) -> m (Either b (a, Stream m a b))
forall (m :: * -> *) a. Monad m => a -> m a
return ((a, Stream m a b) -> Either b (a, Stream m a b)
forall a b. b -> Either a b
Right (a
a, Stream m a a
str Stream m a a -> (a -> Stream m a b) -> Stream m a b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Stream m a b
k))

yield :: Monad m => a -> Stream m a ()
yield :: a -> Stream m a ()
yield a
a = m (Either () (a, Stream m a ())) -> Stream m a ()
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (Either () (a, Stream m a ()) -> m (Either () (a, Stream m a ()))
forall (m :: * -> *) a. Monad m => a -> m a
return ((a, Stream m a ()) -> Either () (a, Stream m a ())
forall a b. b -> Either a b
Right (a
a, () -> Stream m a ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())))

liftIO :: IO a -> Stream IO b a
liftIO :: IO a -> Stream IO b a
liftIO IO a
io = IO (Either a (b, Stream IO b a)) -> Stream IO b a
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (IO (Either a (b, Stream IO b a)) -> Stream IO b a)
-> IO (Either a (b, Stream IO b a)) -> Stream IO b a
forall a b. (a -> b) -> a -> b
$ IO a
io IO a
-> (a -> IO (Either a (b, Stream IO b a)))
-> IO (Either a (b, Stream IO b a))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Either a (b, Stream IO b a) -> IO (Either a (b, Stream IO b a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either a (b, Stream IO b a) -> IO (Either a (b, Stream IO b a)))
-> (a -> Either a (b, Stream IO b a))
-> a
-> IO (Either a (b, Stream IO b a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either a (b, Stream IO b a)
forall a b. a -> Either a b
Left

-- | Turn a Stream into an ordinary list, by demanding all the elements.
collect :: Monad m => Stream m a () -> m [a]
collect :: Stream m a () -> m [a]
collect Stream m a ()
str = Stream m a () -> [a] -> m [a]
forall (m :: * -> *) a. Monad m => Stream m a () -> [a] -> m [a]
go Stream m a ()
str []
 where
  go :: Stream m a () -> [a] -> m [a]
go Stream m a ()
str [a]
acc = do
    Either () (a, Stream m a ())
r <- Stream m a () -> m (Either () (a, Stream m a ()))
forall (m :: * -> *) a b.
Stream m a b -> m (Either b (a, Stream m a b))
runStream Stream m a ()
str
    case Either () (a, Stream m a ())
r of
      Left () -> [a] -> m [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
acc)
      Right (a
a, Stream m a ()
str') -> Stream m a () -> [a] -> m [a]
go Stream m a ()
str' (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[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_ :: Stream m a r -> m ([a], r)
collect_ Stream m a r
str = Stream m a r -> [a] -> m ([a], r)
forall (m :: * -> *) a b.
Monad m =>
Stream m a b -> [a] -> m ([a], b)
go Stream m a r
str []
 where
  go :: Stream m a b -> [a] -> m ([a], b)
go Stream m a b
str [a]
acc = do
    Either b (a, Stream m a b)
r <- Stream m a b -> m (Either b (a, Stream m a b))
forall (m :: * -> *) a b.
Stream m a b -> m (Either b (a, Stream m a b))
runStream Stream m a b
str
    case Either b (a, Stream m a b)
r of
      Left b
r -> ([a], b) -> m ([a], b)
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
acc, b
r)
      Right (a
a, Stream m a b
str') -> Stream m a b -> [a] -> m ([a], b)
go Stream m a b
str' (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
acc)

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

-- | Turn a list into a 'Stream', by yielding each element in turn.
fromList :: Monad m => [a] -> Stream m a ()
fromList :: [a] -> Stream m a ()
fromList = (a -> Stream m a ()) -> [a] -> Stream m a ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ a -> Stream m a ()
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 :: (a -> b) -> Stream m a x -> Stream m b x
map a -> b
f Stream m a x
str = m (Either x (b, Stream m b x)) -> Stream m b x
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (m (Either x (b, Stream m b x)) -> Stream m b x)
-> m (Either x (b, Stream m b x)) -> Stream m b x
forall a b. (a -> b) -> a -> b
$ do
   Either x (a, Stream m a x)
r <- Stream m a x -> m (Either x (a, Stream m a x))
forall (m :: * -> *) a b.
Stream m a b -> m (Either b (a, Stream m a b))
runStream Stream m a x
str
   case Either x (a, Stream m a x)
r of
     Left x
x -> Either x (b, Stream m b x) -> m (Either x (b, Stream m b x))
forall (m :: * -> *) a. Monad m => a -> m a
return (x -> Either x (b, Stream m b x)
forall a b. a -> Either a b
Left x
x)
     Right (a
a, Stream m a x
str') -> Either x (b, Stream m b x) -> m (Either x (b, Stream m b x))
forall (m :: * -> *) a. Monad m => a -> m a
return ((b, Stream m b x) -> Either x (b, Stream m b x)
forall a b. b -> Either a b
Right (a -> b
f a
a, (a -> b) -> Stream m a x -> Stream m b x
forall (m :: * -> *) a b x.
Monad m =>
(a -> b) -> Stream m a x -> Stream m b x
Stream.map a -> b
f Stream m a x
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 :: (a -> m b) -> Stream m a x -> Stream m b x
mapM a -> m b
f Stream m a x
str = m (Either x (b, Stream m b x)) -> Stream m b x
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (m (Either x (b, Stream m b x)) -> Stream m b x)
-> m (Either x (b, Stream m b x)) -> Stream m b x
forall a b. (a -> b) -> a -> b
$ do
   Either x (a, Stream m a x)
r <- Stream m a x -> m (Either x (a, Stream m a x))
forall (m :: * -> *) a b.
Stream m a b -> m (Either b (a, Stream m a b))
runStream Stream m a x
str
   case Either x (a, Stream m a x)
r of
     Left x
x -> Either x (b, Stream m b x) -> m (Either x (b, Stream m b x))
forall (m :: * -> *) a. Monad m => a -> m a
return (x -> Either x (b, Stream m b x)
forall a b. a -> Either a b
Left x
x)
     Right (a
a, Stream m a x
str') -> do
        b
b <- a -> m b
f a
a
        Either x (b, Stream m b x) -> m (Either x (b, Stream m b x))
forall (m :: * -> *) a. Monad m => a -> m a
return ((b, Stream m b x) -> Either x (b, Stream m b x)
forall a b. b -> Either a b
Right (b
b, (a -> m b) -> Stream m a x -> Stream m b x
forall (m :: * -> *) a b x.
Monad m =>
(a -> m b) -> Stream m a x -> Stream m b x
Stream.mapM a -> m b
f Stream m a x
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 :: (c -> a -> m (c, b)) -> c -> Stream m a () -> Stream m b c
mapAccumL c -> a -> m (c, b)
f c
c Stream m a ()
str = m (Either c (b, Stream m b c)) -> Stream m b c
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (m (Either c (b, Stream m b c)) -> Stream m b c)
-> m (Either c (b, Stream m b c)) -> Stream m b c
forall a b. (a -> b) -> a -> b
$ do
  Either () (a, Stream m a ())
r <- Stream m a () -> m (Either () (a, Stream m a ()))
forall (m :: * -> *) a b.
Stream m a b -> m (Either b (a, Stream m a b))
runStream Stream m a ()
str
  case Either () (a, Stream m a ())
r of
    Left  () -> Either c (b, Stream m b c) -> m (Either c (b, Stream m b c))
forall (m :: * -> *) a. Monad m => a -> m a
return (c -> Either c (b, Stream m b c)
forall a b. a -> Either a b
Left c
c)
    Right (a
a, Stream m a ()
str') -> do
      (c
c',b
b) <- c -> a -> m (c, b)
f c
c a
a
      Either c (b, Stream m b c) -> m (Either c (b, Stream m b c))
forall (m :: * -> *) a. Monad m => a -> m a
return ((b, Stream m b c) -> Either c (b, Stream m b c)
forall a b. b -> Either a b
Right (b
b, (c -> a -> m (c, b)) -> c -> Stream m a () -> Stream m b c
forall (m :: * -> *) c a b.
Monad m =>
(c -> a -> m (c, b)) -> c -> Stream m a () -> Stream m b c
mapAccumL c -> a -> m (c, b)
f c
c' Stream m a ()
str'))

mapAccumL_ :: Monad m => (c -> a -> m (c,b)) -> c -> Stream m a r
           -> Stream m b (c, r)
mapAccumL_ :: (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 = m (Either (c, r) (b, Stream m b (c, r))) -> Stream m b (c, r)
forall (m :: * -> *) a b.
m (Either b (a, Stream m a b)) -> Stream m a b
Stream (m (Either (c, r) (b, Stream m b (c, r))) -> Stream m b (c, r))
-> m (Either (c, r) (b, Stream m b (c, r))) -> Stream m b (c, r)
forall a b. (a -> b) -> a -> b
$ do
  Either r (a, Stream m a r)
r <- Stream m a r -> m (Either r (a, Stream m a r))
forall (m :: * -> *) a b.
Stream m a b -> m (Either b (a, Stream m a b))
runStream Stream m a r
str
  case Either r (a, Stream m a r)
r of
    Left  r
r -> Either (c, r) (b, Stream m b (c, r))
-> m (Either (c, r) (b, Stream m b (c, r)))
forall (m :: * -> *) a. Monad m => a -> m a
return ((c, r) -> Either (c, r) (b, Stream m b (c, r))
forall a b. a -> Either a b
Left (c
c, r
r))
    Right (a
a, Stream m a r
str') -> do
      (c
c',b
b) <- c -> a -> m (c, b)
f c
c a
a
      Either (c, r) (b, Stream m b (c, r))
-> m (Either (c, r) (b, Stream m b (c, r)))
forall (m :: * -> *) a. Monad m => a -> m a
return ((b, Stream m b (c, r)) -> Either (c, r) (b, Stream m b (c, r))
forall a b. b -> Either a b
Right (b
b, (c -> a -> m (c, b)) -> c -> Stream m a r -> Stream m b (c, r)
forall (m :: * -> *) c a b 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'))