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

6 files changed, 119 insertions, 111 deletions

diff --git a/compiler/GHC/Cmm/Info.hs b/compiler/GHC/Cmm/Info.hs
index 9298df2544..eaa3c2a923 100644
--- a/compiler/GHC/Cmm/Info.hs
+++ b/compiler/GHC/Cmm/Info.hs
@@ -72,19 +72,20 @@ mkEmptyContInfoTable info_lbl
 cmmToRawCmm :: Logger -> DynFlags -> Stream IO CmmGroupSRTs a
             -> IO (Stream IO RawCmmGroup a)
 cmmToRawCmm logger dflags cmms
-  = do { uniqs <- mkSplitUniqSupply 'i'
-       ; let do_one :: UniqSupply -> [CmmDeclSRTs] -> IO (UniqSupply, [RawCmmDecl])
-             do_one uniqs cmm =
+  = do {
+       ; let do_one :: [CmmDeclSRTs] -> IO [RawCmmDecl]
+             do_one cmm = do
+               uniqs <- mkSplitUniqSupply 'i'
                -- NB. strictness fixes a space leak.  DO NOT REMOVE.
                withTimingSilent logger dflags (text "Cmm -> Raw Cmm")
-                                forceRes $
-                 case initUs uniqs $ concatMapM (mkInfoTable dflags) cmm of
-                   (b,uniqs') -> return (uniqs',b)
-       ; return (snd <$> Stream.mapAccumL_ do_one uniqs cmms)
+                          (\x -> seqList x ())
+                  -- TODO: It might be better to make `mkInfoTable` run in
+                  -- IO as well so we don't have to pass around
+                  -- a UniqSupply (see #16843)
+                 (return $ initUs_ uniqs $ concatMapM (mkInfoTable dflags) cmm)
+       ; return (Stream.mapM do_one cmms)
        }
 
-    where forceRes (uniqs, rawcmms) =
-            uniqs `seq` foldr (\decl r -> decl `seq` r) () rawcmms
 
 -- Make a concrete info table, represented as a list of CmmStatic
 -- (it can't be simply a list of Word, because the SRT field is
diff --git a/compiler/GHC/CmmToAsm.hs b/compiler/GHC/CmmToAsm.hs
index 8201b14ab9..5eda3f03a8 100644
--- a/compiler/GHC/CmmToAsm.hs
+++ b/compiler/GHC/CmmToAsm.hs
@@ -11,6 +11,7 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 {-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
 
@@ -296,7 +297,7 @@ finishNativeGen logger dflags config modLoc bufh@(BufHandle _ _ h) us ngs
                    Opt_D_dump_asm_stats "NCG stats"
                    FormatText
 
-cmmNativeGenStream :: (OutputableP Platform statics, Outputable jumpDest, Instruction instr)
+cmmNativeGenStream :: forall statics jumpDest instr a . (OutputableP Platform statics, Outputable jumpDest, Instruction instr)
               => Logger
               -> DynFlags
               -> NCGConfig
@@ -304,14 +305,21 @@ cmmNativeGenStream :: (OutputableP Platform statics, Outputable jumpDest, Instru
               -> NcgImpl statics instr jumpDest
               -> BufHandle
               -> UniqSupply
-              -> Stream IO RawCmmGroup a
+              -> Stream.Stream IO RawCmmGroup a
               -> NativeGenAcc statics instr
               -> IO (NativeGenAcc statics instr, UniqSupply, a)
 
 cmmNativeGenStream logger dflags config modLoc ncgImpl h us cmm_stream ngs
- = do r <- Stream.runStream cmm_stream
-      case r of
-        Left a ->
+ = loop us (Stream.runStream cmm_stream) ngs
+  where
+    ncglabel = text "NCG"
+    loop :: UniqSupply
+              -> Stream.StreamS IO RawCmmGroup a
+              -> NativeGenAcc statics instr
+              -> IO (NativeGenAcc statics instr, UniqSupply, a)
+    loop us s ngs =
+      case s of
+        Stream.Done a ->
           return (ngs { ngs_imports = reverse $ ngs_imports ngs
                       , ngs_natives = reverse $ ngs_natives ngs
                       , ngs_colorStats = reverse $ ngs_colorStats ngs
@@ -319,7 +327,8 @@ cmmNativeGenStream logger dflags config modLoc ncgImpl h us cmm_stream ngs
                       },
                   us,
                   a)
-        Right (cmms, cmm_stream') -> do
+        Stream.Effect m -> m >>= \cmm_stream' -> loop us cmm_stream' ngs
+        Stream.Yield cmms cmm_stream' -> do
           (us', ngs'') <-
             withTimingSilent logger
                 dflags
@@ -345,10 +354,8 @@ cmmNativeGenStream logger dflags config modLoc ncgImpl h us cmm_stream ngs
               let ngs'' = ngs' { ngs_debug = ngs_debug ngs' ++ ldbgs, ngs_labels = [] }
               return (us', ngs'')
 
-          cmmNativeGenStream logger dflags config modLoc ncgImpl h us'
-              cmm_stream' ngs''
+          loop us' cmm_stream' ngs''
 
-    where ncglabel = text "NCG"
 
 -- | Do native code generation on all these cmms.
 --
diff --git a/compiler/GHC/CmmToLlvm.hs b/compiler/GHC/CmmToLlvm.hs
index 3cf7b50ceb..21cfdf6dcd 100644
--- a/compiler/GHC/CmmToLlvm.hs
+++ b/compiler/GHC/CmmToLlvm.hs
@@ -77,13 +77,13 @@ llvmCodeGen logger dflags h cmm_stream
 
        -- run code generation
        a <- runLlvm logger dflags (fromMaybe supportedLlvmVersion mb_ver) bufh $
-         llvmCodeGen' dflags (liftStream cmm_stream)
+         llvmCodeGen' dflags cmm_stream
 
        bFlush bufh
 
        return a
 
-llvmCodeGen' :: DynFlags -> Stream.Stream LlvmM RawCmmGroup a -> LlvmM a
+llvmCodeGen' :: DynFlags -> Stream.Stream IO RawCmmGroup a -> LlvmM a
 llvmCodeGen' dflags cmm_stream
   = do  -- Preamble
         renderLlvm header
@@ -91,7 +91,7 @@ llvmCodeGen' dflags cmm_stream
         cmmMetaLlvmPrelude
 
         -- Procedures
-        a <- Stream.consume cmm_stream llvmGroupLlvmGens
+        a <- Stream.consume cmm_stream liftIO llvmGroupLlvmGens
 
         -- Declare aliases for forward references
         opts <- getLlvmOpts
diff --git a/compiler/GHC/CmmToLlvm/Base.hs b/compiler/GHC/CmmToLlvm/Base.hs
index 84c82ef873..a943bfcebb 100644
--- a/compiler/GHC/CmmToLlvm/Base.hs
+++ b/compiler/GHC/CmmToLlvm/Base.hs
@@ -19,14 +19,14 @@ module GHC.CmmToLlvm.Base (
         llvmVersionStr, llvmVersionList,
 
         LlvmM,
-        runLlvm, liftStream, withClearVars, varLookup, varInsert,
+        runLlvm, withClearVars, varLookup, varInsert,
         markStackReg, checkStackReg,
         funLookup, funInsert, getLlvmVer, getDynFlags,
         dumpIfSetLlvm, renderLlvm, markUsedVar, getUsedVars,
         ghcInternalFunctions, getPlatform, getLlvmOpts,
 
         getMetaUniqueId,
-        setUniqMeta, getUniqMeta,
+        setUniqMeta, getUniqMeta, liftIO,
 
         cmmToLlvmType, widthToLlvmFloat, widthToLlvmInt, llvmFunTy,
         llvmFunSig, llvmFunArgs, llvmStdFunAttrs, llvmFunAlign, llvmInfAlign,
@@ -62,7 +62,6 @@ import GHC.Utils.BufHandle   ( BufHandle )
 import GHC.Types.Unique.Set
 import GHC.Types.Unique.Supply
 import GHC.Utils.Logger
-import qualified GHC.Data.Stream as Stream
 
 import Data.Maybe (fromJust)
 import Control.Monad (ap)
@@ -387,14 +386,6 @@ getEnv f = LlvmM (\env -> return (f env, env))
 modifyEnv :: (LlvmEnv -> LlvmEnv) -> LlvmM ()
 modifyEnv f = LlvmM (\env -> return ((), f env))
 
--- | Lift a stream into the LlvmM monad
-liftStream :: Stream.Stream IO a x -> Stream.Stream LlvmM a x
-liftStream s = Stream.Stream $ do
-  r <- liftIO $ Stream.runStream s
-  case r of
-    Left b        -> return (Left b)
-    Right (a, r2) -> return (Right (a, liftStream r2))
-
 -- | Clear variables from the environment for a subcomputation
 withClearVars :: LlvmM a -> LlvmM a
 withClearVars m = LlvmM $ \env -> do
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)
diff --git a/compiler/GHC/Driver/CodeOutput.hs b/compiler/GHC/Driver/CodeOutput.hs
index fb6d04afbf..134ee2f960 100644
--- a/compiler/GHC/Driver/CodeOutput.hs
+++ b/compiler/GHC/Driver/CodeOutput.hs
@@ -149,7 +149,7 @@ outputC logger dflags filenm cmm_stream packages =
                           FormatC
                           doc
             printForC dflags h doc
-      Stream.consume cmm_stream writeC
+      Stream.consume cmm_stream id writeC
 
 {-
 ************************************************************************