Reimplement Stream in "yoneda" style for efficiency

'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. It is well-known for free monads that this
representation can improve performance, and the test results
demonstrate this for GHC as well.

The operation mapAccumL is not used in the compiler and can't be
implemented efficiently because it requires destroying and rebuilding
the stream.

I removed one use of mapAccumL_ which has similar problems but the other
use was difficult to remove. In the future it may be worth exploring
whether the 'Stream' encoding could be modified further to capture the
mapAccumL pattern, and likewise defer the passing of accumulation
parameter until the stream is finally consumed.

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"

Metric Decrease:
    ManyConstructors

1 files changed, 87 insertions, 78 deletions

diff --git a/compiler/GHC/Data/Stream.hs b/compiler/GHC/Data/Stream.hs
index 7996ee7343..4e2bee4311 100644
--- a/compiler/GHC/Data/Stream.hs
+++ b/compiler/GHC/Data/Stream.hs
@@ -1,3 +1,6 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE TupleSections #-}
 -- -----------------------------------------------------------------------------
 --
 -- (c) The University of Glasgow 2012
@@ -6,25 +9,28 @@
 
 -- | Monadic streams
 module GHC.Data.Stream (
-    Stream(..), yield, liftIO,
-    collect, collect_, consume, fromList,
-    map, mapM, mapAccumL, mapAccumL_
+    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 @runStream@
+-- 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: @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.
+--  * 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
@@ -38,57 +44,73 @@ import Control.Monad hiding (mapM)
 -- 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)) }
+-- '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 { 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 st = runStreamInternal st pure Done
+
+data StreamS m a b = Yield a (StreamS m a b)
+                   | Done b
+                   | Effect (m (StreamS m a b))
 
-instance Monad f => Functor (Stream f a) where
+instance Monad m => Functor (StreamS m a) where
   fmap = liftM
 
-instance Monad m => Applicative (Stream m a) where
-  pure a = Stream (return (Left a))
+instance Monad m => Applicative (StreamS m a) where
+  pure = Done
   (<*>) = ap
 
-instance Monad m => Monad (Stream m a) where
+instance Monad m => Monad (StreamS m a) where
+  a >>= k = case a of
+                      Done r -> k r
+                      Yield a s -> Yield a (s >>= k)
+                      Effect m -> Effect (fmap (>>= k) m)
 
-  Stream m >>= k = Stream $ do
-                r <- m
-                case r of
-                  Left b        -> runStream (k b)
-                  Right (a,str) -> return (Right (a, str >>= k))
+instance Functor (Stream f a) where
+  fmap = liftM
 
-yield :: Monad m => a -> Stream m a ()
-yield a = Stream (return (Right (a, return ())))
+instance Applicative (Stream m a) where
+  pure a = Stream $ \_f g -> g a
+  (<*>) = ap
 
-liftIO :: IO a -> Stream IO b a
-liftIO io = Stream $ io >>= return . Left
+instance Monad (Stream m a) where
+  Stream m >>= k = Stream $ \f h -> m f (\a -> runStreamInternal (k a) f h)
+
+instance MonadIO m => MonadIO (Stream m b) where
+  liftIO io = Stream $ \_f g -> Effect (g <$> liftIO io)
+
+yield :: Monad m => a -> Stream m a ()
+yield a = Stream $ \f rest -> Effect (flip Yield (rest ())  <$> f a)
 
 -- | Turn a Stream into an ordinary list, by demanding all the elements.
 collect :: Monad m => Stream m a () -> m [a]
-collect str = go str []
+collect str = go [] (runStream str)
  where
-  go str acc = do
-    r <- runStream str
-    case r of
-      Left () -> return (reverse acc)
-      Right (a, str') -> go str' (a:acc)
+  go acc (Done ()) = return (reverse acc)
+  go acc (Effect m) = m >>= go acc
+  go acc (Yield a k) = go (a:acc) k
 
--- | 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
+consume :: (Monad m, Monad n) => Stream m a b -> (forall a . m a -> n a) -> (a -> n ()) -> n b
+consume str l f = go (runStream str)
+  where
+    go (Done r) = return r
+    go (Yield a p) = f a >> go p
+    go (Effect m)  = l m >>= go
 
 -- | Turn a list into a 'Stream', by yielding each element in turn.
 fromList :: Monad m => [a] -> Stream m a ()
@@ -96,40 +118,27 @@ 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'))
+map f str = Stream $ \g h -> runStreamInternal str (g . f) 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 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
+mapM f str = Stream $ \g h -> runStreamInternal str (g <=< f) 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_ 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'))
+mapAccumL_ f c str = Stream $ \f h -> go c f h (runStream str)
+
+  where
+    go :: c
+             -> (b -> m r')
+             -> ((c, r) -> StreamS m r' r1)
+             -> StreamS m a r
+             -> StreamS m r' r1
+    go c _f1 h1 (Done r) = h1 (c, r)
+    go c f1 h1 (Yield a p) = Effect (f c a >>= (\(c', b) -> f1 b
+                                           >>= \r' -> return $ Yield r' (go c' f1 h1 p)))
+    go c f1 h1 (Effect m) = Effect (go c f1 h1 <$> m)