Implement late lambda lift

Summary:
This implements a selective lambda-lifting pass late in the STG
pipeline.

Lambda lifting has the effect of avoiding closure allocation at the cost
of having to make former free vars available at call sites, possibly
enlarging closures surrounding call sites in turn.

We identify beneficial cases by means of an analysis that estimates
closure growth.

There's a Wiki page at
https://ghc.haskell.org/trac/ghc/wiki/LateLamLift.

Reviewers: simonpj, bgamari, simonmar

Reviewed By: simonpj

Subscribers: rwbarton, carter

GHC Trac Issues: #9476

Differential Revision: https://phabricator.haskell.org/D5224

5 files changed, 203 insertions, 75 deletions

diff --git a/compiler/stgSyn/CoreToStg.hs b/compiler/stgSyn/CoreToStg.hs
index 74bb7b6014..573db78a06 100644
--- a/compiler/stgSyn/CoreToStg.hs
+++ b/compiler/stgSyn/CoreToStg.hs
@@ -631,8 +631,8 @@ coreToStgLet bind body = do
 
         -- Compute the new let-expression
     let
-        new_let | isJoinBind bind = StgLetNoEscape bind2 body2
-                | otherwise       = StgLet bind2 body2
+        new_let | isJoinBind bind = StgLetNoEscape noExtSilent bind2 body2
+                | otherwise       = StgLet noExtSilent bind2 body2
 
     return new_let
   where
diff --git a/compiler/stgSyn/StgFVs.hs b/compiler/stgSyn/StgFVs.hs
index 80ce33ff7a..edfc94ed2d 100644
--- a/compiler/stgSyn/StgFVs.hs
+++ b/compiler/stgSyn/StgFVs.hs
@@ -1,6 +1,7 @@
 -- | Free variable analysis on STG terms.
 module StgFVs (
-    annTopBindingsFreeVars
+    annTopBindingsFreeVars,
+    annBindingFreeVars
   ) where
 
 import GhcPrelude
@@ -26,13 +27,17 @@ addLocals :: [Id] -> Env -> Env
 addLocals bndrs env
   = env { locals = extendVarSetList (locals env) bndrs }
 
--- | Annotates a top-level STG binding with its free variables.
+-- | Annotates a top-level STG binding group with its free variables.
 annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
 annTopBindingsFreeVars = map go
   where
     go (StgTopStringLit id bs) = StgTopStringLit id bs
     go (StgTopLifted bind)
-      = StgTopLifted (fst (binding emptyEnv emptyVarSet bind))
+      = StgTopLifted (annBindingFreeVars bind)
+
+-- | Annotates an STG binding with its free variables.
+annBindingFreeVars :: StgBinding -> CgStgBinding
+annBindingFreeVars = fst . binding emptyEnv emptyDVarSet
 
 boundIds :: StgBinding -> [Id]
 boundIds (StgNonRec b _) = [b]
@@ -53,35 +58,35 @@ boundIds (StgRec pairs)  = map fst pairs
 --      knot-tying.
 
 -- | This makes sure that only local, non-global free vars make it into the set.
-mkFreeVarSet :: Env -> [Id] -> IdSet
-mkFreeVarSet env = mkVarSet . filter (`elemVarSet` locals env)
+mkFreeVarSet :: Env -> [Id] -> DIdSet
+mkFreeVarSet env = mkDVarSet . filter (`elemVarSet` locals env)
 
-args :: Env -> [StgArg] -> IdSet
+args :: Env -> [StgArg] -> DIdSet
 args env = mkFreeVarSet env . mapMaybe f
   where
     f (StgVarArg occ) = Just occ
     f _               = Nothing
 
-binding :: Env -> IdSet -> StgBinding -> (CgStgBinding, IdSet)
+binding :: Env -> DIdSet -> StgBinding -> (CgStgBinding, DIdSet)
 binding env body_fv (StgNonRec bndr r) = (StgNonRec bndr r', fvs)
   where
     -- See Note [Tacking local binders]
     (r', rhs_fvs) = rhs env r
-    fvs = delVarSet body_fv bndr `unionVarSet` rhs_fvs
+    fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_fvs
 binding env body_fv (StgRec pairs) = (StgRec pairs', fvs)
   where
     -- See Note [Tacking local binders]
     bndrs = map fst pairs
     (rhss, rhs_fvss) = mapAndUnzip (rhs env . snd) pairs
     pairs' = zip bndrs rhss
-    fvs = delVarSetList (unionVarSets (body_fv:rhs_fvss)) bndrs
+    fvs = delDVarSetList (unionDVarSets (body_fv:rhs_fvss)) bndrs
 
-expr :: Env -> StgExpr -> (CgStgExpr, IdSet)
+expr :: Env -> StgExpr -> (CgStgExpr, DIdSet)
 expr env = go
   where
     go (StgApp occ as)
-      = (StgApp occ as, unionVarSet (args env as) (mkFreeVarSet env [occ]))
-    go (StgLit lit) = (StgLit lit, emptyVarSet)
+      = (StgApp occ as, unionDVarSet (args env as) (mkFreeVarSet env [occ]))
+    go (StgLit lit) = (StgLit lit, emptyDVarSet)
     go (StgConApp dc as tys) = (StgConApp dc as tys, args env as)
     go (StgOpApp op as ty) = (StgOpApp op as ty, args env as)
     go StgLam{} = pprPanic "StgFVs: StgLam" empty
@@ -90,16 +95,16 @@ expr env = go
         (scrut', scrut_fvs) = go scrut
         -- See Note [Tacking local binders]
         (alts', alt_fvss) = mapAndUnzip (alt (addLocals [bndr] env)) alts
-        alt_fvs = unionVarSets alt_fvss
-        fvs = delVarSet (unionVarSet scrut_fvs alt_fvs) bndr
-    go (StgLet bind body) = go_bind StgLet bind body
-    go (StgLetNoEscape bind body) = go_bind StgLetNoEscape bind body
+        alt_fvs = unionDVarSets alt_fvss
+        fvs = delDVarSet (unionDVarSet scrut_fvs alt_fvs) bndr
+    go (StgLet ext bind body) = go_bind (StgLet ext) bind body
+    go (StgLetNoEscape ext bind body) = go_bind (StgLetNoEscape ext) bind body
     go (StgTick tick e) = (StgTick tick e', fvs')
       where
         (e', fvs) = go e
-        fvs' = unionVarSet (tickish tick) fvs
-        tickish (Breakpoint _ ids) = mkVarSet ids
-        tickish _                  = emptyVarSet
+        fvs' = unionDVarSet (tickish tick) fvs
+        tickish (Breakpoint _ ids) = mkDVarSet ids
+        tickish _                  = emptyDVarSet
 
     go_bind dc bind body = (dc bind' body', fvs)
       where
@@ -108,18 +113,18 @@ expr env = go
         (body', body_fvs) = expr env' body
         (bind', fvs) = binding env' body_fvs bind
 
-rhs :: Env -> StgRhs -> (CgStgRhs, IdSet)
+rhs :: Env -> StgRhs -> (CgStgRhs, DIdSet)
 rhs env (StgRhsClosure _ ccs uf bndrs body)
   = (StgRhsClosure fvs ccs uf bndrs body', fvs)
   where
     -- See Note [Tacking local binders]
     (body', body_fvs) = expr (addLocals bndrs env) body
-    fvs = delVarSetList body_fvs bndrs
+    fvs = delDVarSetList body_fvs bndrs
 rhs env (StgRhsCon ccs dc as) = (StgRhsCon ccs dc as, args env as)
 
-alt :: Env -> StgAlt -> (CgStgAlt, IdSet)
+alt :: Env -> StgAlt -> (CgStgAlt, DIdSet)
 alt env (con, bndrs, e) = ((con, bndrs, e'), fvs)
   where
     -- See Note [Tacking local binders]
     (e', rhs_fvs) = expr (addLocals bndrs env) e
-    fvs = delVarSetList rhs_fvs bndrs
+    fvs = delDVarSetList rhs_fvs bndrs
diff --git a/compiler/stgSyn/StgLint.hs b/compiler/stgSyn/StgLint.hs
index 35a498f368..383b016f08 100644
--- a/compiler/stgSyn/StgLint.hs
+++ b/compiler/stgSyn/StgLint.hs
@@ -40,6 +40,8 @@ import StgSyn
 
 import DynFlags
 import Bag              ( Bag, emptyBag, isEmptyBag, snocBag, bagToList )
+import BasicTypes       ( TopLevelFlag(..), isTopLevel )
+import CostCentre       ( isCurrentCCS )
 import Id               ( Id, idType, isLocalId, isJoinId )
 import VarSet
 import DataCon
@@ -84,7 +86,7 @@ lintStgTopBindings dflags unarised whodunnit binds
         addInScopeVars binders $
             lint_binds binds
 
-    lint_bind (StgTopLifted bind) = lintStgBinds bind
+    lint_bind (StgTopLifted bind) = lintStgBinds TopLevel bind
     lint_bind (StgTopStringLit v _) = return [v]
 
 lintStgArg :: StgArg -> LintM ()
@@ -94,26 +96,39 @@ lintStgArg (StgVarArg v) = lintStgVar v
 lintStgVar :: Id -> LintM ()
 lintStgVar id = checkInScope id
 
-lintStgBinds :: StgBinding -> LintM [Id] -- Returns the binders
-lintStgBinds (StgNonRec binder rhs) = do
-    lint_binds_help (binder,rhs)
+lintStgBinds :: TopLevelFlag -> StgBinding -> LintM [Id] -- Returns the binders
+lintStgBinds top_lvl (StgNonRec binder rhs) = do
+    lint_binds_help top_lvl (binder,rhs)
     return [binder]
 
-lintStgBinds (StgRec pairs)
+lintStgBinds top_lvl (StgRec pairs)
   = addInScopeVars binders $ do
-        mapM_ lint_binds_help pairs
+        mapM_ (lint_binds_help top_lvl) pairs
         return binders
   where
     binders = [b | (b,_) <- pairs]
 
-lint_binds_help :: (Id, StgRhs) -> LintM ()
-lint_binds_help (binder, rhs)
+lint_binds_help :: TopLevelFlag -> (Id, StgRhs) -> LintM ()
+lint_binds_help top_lvl (binder, rhs)
   = addLoc (RhsOf binder) $ do
+        when (isTopLevel top_lvl) (checkNoCurrentCCS rhs)
         lintStgRhs rhs
         -- Check binder doesn't have unlifted type or it's a join point
         checkL (isJoinId binder || not (isUnliftedType (idType binder)))
                (mkUnliftedTyMsg binder rhs)
 
+-- | Top-level bindings can't inherit the cost centre stack from their
+-- (static) allocation site.
+checkNoCurrentCCS :: StgRhs -> LintM ()
+checkNoCurrentCCS (StgRhsClosure _ ccs _ _ _)
+  | isCurrentCCS ccs
+  = addErrL (text "Top-level StgRhsClosure with CurrentCCS")
+checkNoCurrentCCS (StgRhsCon ccs _ _)
+  | isCurrentCCS ccs
+  = addErrL (text "Top-level StgRhsCon with CurrentCCS")
+checkNoCurrentCCS _
+  = return ()
+
 lintStgRhs :: StgRhs -> LintM ()
 
 lintStgRhs (StgRhsClosure _ _ _ [] expr)
@@ -154,14 +169,14 @@ lintStgExpr (StgOpApp _ args _) =
 lintStgExpr lam@(StgLam _ _) =
     addErrL (text "Unexpected StgLam" <+> ppr lam)
 
-lintStgExpr (StgLet binds body) = do
-    binders <- lintStgBinds binds
+lintStgExpr (StgLet _ binds body) = do
+    binders <- lintStgBinds NotTopLevel binds
     addLoc (BodyOfLetRec binders) $
       addInScopeVars binders $
         lintStgExpr body
 
-lintStgExpr (StgLetNoEscape binds body) = do
-    binders <- lintStgBinds binds
+lintStgExpr (StgLetNoEscape _ binds body) = do
+    binders <- lintStgBinds NotTopLevel binds
     addLoc (BodyOfLetRec binders) $
       addInScopeVars binders $
         lintStgExpr body
diff --git a/compiler/stgSyn/StgSubst.hs b/compiler/stgSyn/StgSubst.hs
new file mode 100644
index 0000000000..72fbe418d1
--- /dev/null
+++ b/compiler/stgSyn/StgSubst.hs
@@ -0,0 +1,80 @@
+{-# LANGUAGE CPP #-}
+
+module StgSubst where
+
+#include "HsVersions.h"
+
+import GhcPrelude
+
+import Id
+import VarEnv
+import Control.Monad.Trans.State.Strict
+import Outputable
+import Util
+
+-- | A renaming substitution from 'Id's to 'Id's. Like 'RnEnv2', but not
+-- maintaining pairs of substitutions. Like @"CoreSubst".'CoreSubst.Subst'@, but
+-- with the domain being 'Id's instead of entire 'CoreExpr'.
+data Subst = Subst InScopeSet IdSubstEnv
+
+type IdSubstEnv = IdEnv Id
+
+-- | @emptySubst = 'mkEmptySubst' 'emptyInScopeSet'@
+emptySubst :: Subst
+emptySubst = mkEmptySubst emptyInScopeSet
+
+-- | Constructs a new 'Subst' assuming the variables in the given 'InScopeSet'
+-- are in scope.
+mkEmptySubst :: InScopeSet -> Subst
+mkEmptySubst in_scope = Subst in_scope emptyVarEnv
+
+-- | Substitutes an 'Id' for another one according to the 'Subst' given in a way
+-- that avoids shadowing the 'InScopeSet', returning the result and an updated
+-- 'Subst' that should be used by subsequent substitutions.
+substBndr :: Id -> Subst -> (Id, Subst)
+substBndr id (Subst in_scope env)
+  = (new_id, Subst new_in_scope new_env)
+  where
+    new_id = uniqAway in_scope id
+    no_change = new_id == id -- in case nothing shadowed
+    new_in_scope = in_scope `extendInScopeSet` new_id
+    new_env
+      | no_change = delVarEnv env id
+      | otherwise = extendVarEnv env id new_id
+
+-- | @substBndrs = runState . traverse (state . substBndr)@
+substBndrs :: Traversable f => f Id -> Subst -> (f Id, Subst)
+substBndrs = runState . traverse (state . substBndr)
+
+-- | Substitutes an occurrence of an identifier for its counterpart recorded
+-- in the 'Subst'.
+lookupIdSubst :: HasCallStack => Id -> Subst -> Id
+lookupIdSubst id (Subst in_scope env)
+  | not (isLocalId id) = id
+  | Just id' <- lookupVarEnv env id = id'
+  | Just id' <- lookupInScope in_scope id = id'
+  | otherwise = WARN( True, text "StgSubst.lookupIdSubst" <+> ppr id $$ ppr in_scope)
+                id
+
+-- | Substitutes an occurrence of an identifier for its counterpart recorded
+-- in the 'Subst'. Does not generate a debug warning if the identifier to
+-- to substitute wasn't in scope.
+noWarnLookupIdSubst :: HasCallStack => Id -> Subst -> Id
+noWarnLookupIdSubst id (Subst in_scope env)
+  | not (isLocalId id) = id
+  | Just id' <- lookupVarEnv env id = id'
+  | Just id' <- lookupInScope in_scope id = id'
+  | otherwise = id
+
+-- | Add the 'Id' to the in-scope set and remove any existing substitutions for
+-- it.
+extendInScope :: Id -> Subst -> Subst
+extendInScope id (Subst in_scope env) = Subst (in_scope `extendInScopeSet` id) env
+
+-- | Add a substitution for an 'Id' to the 'Subst': you must ensure that the
+-- in-scope set is such that TyCORep Note [The substitution invariant]
+-- holds after extending the substitution like this.
+extendSubst :: Id -> Id -> Subst -> Subst
+extendSubst id new_id (Subst in_scope env)
+  = ASSERT2( new_id `elemInScopeSet` in_scope, ppr id <+> ppr new_id $$ ppr in_scope )
+    Subst in_scope (extendVarEnv env id new_id)
diff --git a/compiler/stgSyn/StgSyn.hs b/compiler/stgSyn/StgSyn.hs
index 145c001046..5ba63e458c 100644
--- a/compiler/stgSyn/StgSyn.hs
+++ b/compiler/stgSyn/StgSyn.hs
@@ -16,6 +16,7 @@ generation.
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE ConstraintKinds #-}
 
 module StgSyn (
         StgArg(..),
@@ -23,7 +24,8 @@ module StgSyn (
         GenStgTopBinding(..), GenStgBinding(..), GenStgExpr(..), GenStgRhs(..),
         GenStgAlt, AltType(..),
 
-        StgPass(..), XRhsClosure, NoExtSilent, noExtSilent,
+        StgPass(..), BinderP, XRhsClosure, XLet, XLetNoEscape,
+        NoExtSilent, noExtSilent,
 
         UpdateFlag(..), isUpdatable,
 
@@ -33,6 +35,9 @@ module StgSyn (
         -- a set of synonyms for the code gen parameterisation
         CgStgTopBinding, CgStgBinding, CgStgExpr, CgStgRhs, CgStgAlt,
 
+        -- a set of synonyms for the lambda lifting parameterisation
+        LlStgTopBinding, LlStgBinding, LlStgExpr, LlStgRhs, LlStgAlt,
+
         -- a set of synonyms to distinguish in- and out variants
         InStgArg,  InStgTopBinding,  InStgBinding,  InStgExpr,  InStgRhs,  InStgAlt,
         OutStgArg, OutStgTopBinding, OutStgBinding, OutStgExpr, OutStgRhs, OutStgAlt,
@@ -101,8 +106,8 @@ data GenStgTopBinding pass
   | StgTopStringLit Id ByteString
 
 data GenStgBinding pass
-  = StgNonRec Id (GenStgRhs pass)
-  | StgRec    [(Id, GenStgRhs pass)]
+  = StgNonRec (BinderP pass) (GenStgRhs pass)
+  | StgRec    [(BinderP pass, GenStgRhs pass)]
 
 {-
 ************************************************************************
@@ -245,7 +250,7 @@ TODO: Encode this via an extension to GenStgExpr à la TTG.
 -}
 
   | StgLam
-        (NonEmpty Id)
+        (NonEmpty (BinderP pass))
         StgExpr    -- Body of lambda
 
 {-
@@ -259,13 +264,9 @@ This has the same boxed/unboxed business as Core case expressions.
 -}
 
   | StgCase
-        (GenStgExpr pass)
-                    -- the thing to examine
-
-        Id          -- binds the result of evaluating the scrutinee
-
+        (GenStgExpr pass) -- the thing to examine
+        (BinderP pass) -- binds the result of evaluating the scrutinee
         AltType
-
         [GenStgAlt pass]
                     -- The DEFAULT case is always *first*
                     -- if it is there at all
@@ -365,10 +366,12 @@ And so the code for let(rec)-things:
 -}
 
   | StgLet
+        (XLet pass)
         (GenStgBinding pass)    -- right hand sides (see below)
         (GenStgExpr pass)       -- body
 
   | StgLetNoEscape
+        (XLetNoEscape pass)
         (GenStgBinding pass)    -- right hand sides (see below)
         (GenStgExpr pass)       -- body
 
@@ -405,7 +408,7 @@ data GenStgRhs pass
                            --   list just before 'CodeGen'.
         CostCentreStack    -- ^ CCS to be attached (default is CurrentCCS)
         !UpdateFlag        -- ^ 'ReEntrant' | 'Updatable' | 'SingleEntry'
-        [Id]               -- ^ arguments; if empty, then not a function;
+        [BinderP pass]     -- ^ arguments; if empty, then not a function;
                            --   as above, order is important.
         (GenStgExpr pass)  -- ^ body
 
@@ -437,8 +440,9 @@ The second flavour of right-hand-side is for constructors (simple but important)
 
 -- | Used as a data type index for the stgSyn AST
 data StgPass
-  = CodeGen
-  | Vanilla
+  = Vanilla
+  | LiftLams
+  | CodeGen
 
 -- | Like 'HsExpression.NoExt', but with an 'Outputable' instance that returns
 -- 'empty'.
@@ -455,9 +459,24 @@ noExtSilent = NoExtSilent
 -- TODO: Maybe move this to HsExtensions? I'm not sure about the implications
 -- on build time...
 
-type family XRhsClosure (pass :: StgPass) where
-  XRhsClosure 'CodeGen = IdSet -- code gen needs to track non-global free vars
-  XRhsClosure 'Vanilla = NoExtSilent
+-- TODO: Do we really want to the extension point type families to have a closed
+-- domain?
+type family BinderP (pass :: StgPass)
+type instance BinderP 'Vanilla = Id
+type instance BinderP 'CodeGen = Id
+
+type family XRhsClosure (pass :: StgPass)
+type instance XRhsClosure 'Vanilla = NoExtSilent
+-- | Code gen needs to track non-global free vars
+type instance XRhsClosure 'CodeGen = DIdSet
+
+type family XLet (pass :: StgPass)
+type instance XLet 'Vanilla = NoExtSilent
+type instance XLet 'CodeGen = NoExtSilent
+
+type family XLetNoEscape (pass :: StgPass)
+type instance XLetNoEscape 'Vanilla = NoExtSilent
+type instance XLetNoEscape 'CodeGen = NoExtSilent
 
 stgRhsArity :: StgRhs -> Int
 stgRhsArity (StgRhsClosure _ _ _ bndrs _)
@@ -506,9 +525,9 @@ exprHasCafRefs (StgLam _ body)
   = exprHasCafRefs body
 exprHasCafRefs (StgCase scrt _ _ alts)
   = exprHasCafRefs scrt || any altHasCafRefs alts
-exprHasCafRefs (StgLet bind body)
+exprHasCafRefs (StgLet _ bind body)
   = bindHasCafRefs bind || exprHasCafRefs body
-exprHasCafRefs (StgLetNoEscape bind body)
+exprHasCafRefs (StgLetNoEscape _ bind body)
   = bindHasCafRefs bind || exprHasCafRefs body
 exprHasCafRefs (StgTick _ expr)
   = exprHasCafRefs expr
@@ -562,7 +581,7 @@ rather than from the scrutinee type.
 
 type GenStgAlt pass
   = (AltCon,          -- alts: data constructor,
-     [Id],            -- constructor's parameters,
+     [BinderP pass],  -- constructor's parameters,
      GenStgExpr pass) -- ...right-hand side.
 
 data AltType
@@ -589,6 +608,12 @@ type StgExpr       = GenStgExpr       'Vanilla
 type StgRhs        = GenStgRhs        'Vanilla
 type StgAlt        = GenStgAlt        'Vanilla
 
+type LlStgTopBinding = GenStgTopBinding 'LiftLams
+type LlStgBinding    = GenStgBinding    'LiftLams
+type LlStgExpr       = GenStgExpr       'LiftLams
+type LlStgRhs        = GenStgRhs        'LiftLams
+type LlStgAlt        = GenStgAlt        'LiftLams
+
 type CgStgTopBinding = GenStgTopBinding 'CodeGen
 type CgStgBinding    = GenStgBinding    'CodeGen
 type CgStgExpr       = GenStgExpr       'CodeGen
@@ -676,8 +701,15 @@ Robin Popplestone asked for semi-colon separators on STG binds; here's
 hoping he likes terminators instead...  Ditto for case alternatives.
 -}
 
+type OutputablePass pass =
+  ( Outputable (XLet pass)
+  , Outputable (XLetNoEscape pass)
+  , Outputable (XRhsClosure pass)
+  , OutputableBndr (BinderP pass)
+  )
+
 pprGenStgTopBinding
-  :: Outputable (XRhsClosure pass) => GenStgTopBinding pass -> SDoc
+  :: OutputablePass pass => GenStgTopBinding pass -> SDoc
 pprGenStgTopBinding (StgTopStringLit bndr str)
   = hang (hsep [pprBndr LetBind bndr, equals])
         4 (pprHsBytes str <> semi)
@@ -685,7 +717,7 @@ pprGenStgTopBinding (StgTopLifted bind)
   = pprGenStgBinding bind
 
 pprGenStgBinding
-  :: (Outputable (XRhsClosure pass)) => GenStgBinding pass -> SDoc
+  :: OutputablePass pass => GenStgBinding pass -> SDoc
 
 pprGenStgBinding (StgNonRec bndr rhs)
   = hang (hsep [pprBndr LetBind bndr, equals])
@@ -709,27 +741,23 @@ pprStgTopBindings binds
 instance Outputable StgArg where
     ppr = pprStgArg
 
-instance (Outputable (XRhsClosure pass))
-                => Outputable (GenStgTopBinding pass) where
+instance OutputablePass pass => Outputable (GenStgTopBinding pass) where
     ppr = pprGenStgTopBinding
 
-instance (Outputable (XRhsClosure pass))
-                => Outputable (GenStgBinding pass) where
+instance OutputablePass pass => Outputable (GenStgBinding pass) where
     ppr = pprGenStgBinding
 
-instance (Outputable (XRhsClosure pass))
-                => Outputable (GenStgExpr pass) where
+instance OutputablePass pass => Outputable (GenStgExpr pass) where
     ppr = pprStgExpr
 
-instance (Outputable (XRhsClosure pass))
-                => Outputable (GenStgRhs pass) where
+instance OutputablePass pass => Outputable (GenStgRhs pass) where
     ppr rhs = pprStgRhs rhs
 
 pprStgArg :: StgArg -> SDoc
 pprStgArg (StgVarArg var) = ppr var
 pprStgArg (StgLitArg con) = ppr con
 
-pprStgExpr :: (Outputable (XRhsClosure pass)) => GenStgExpr pass -> SDoc
+pprStgExpr :: OutputablePass pass => GenStgExpr pass -> SDoc
 -- special case
 pprStgExpr (StgLit lit)     = ppr lit
 
@@ -773,19 +801,19 @@ pprStgExpr (StgLet srt (StgNonRec bndr (StgRhsClosure cc bi free_vars upd_flag a
 
 -- special case: let ... in let ...
 
-pprStgExpr (StgLet bind expr@(StgLet _ _))
+pprStgExpr (StgLet ext bind expr@StgLet{})
   = ($$)
-      (sep [hang (text "let {")
+      (sep [hang (text "let" <+> ppr ext <+> text "{")
                 2 (hsep [pprGenStgBinding bind, text "} in"])])
       (ppr expr)
 
 -- general case
-pprStgExpr (StgLet bind expr)
-  = sep [hang (text "let {") 2 (pprGenStgBinding bind),
+pprStgExpr (StgLet ext bind expr)
+  = sep [hang (text "let" <+> ppr ext <+> text "{") 2 (pprGenStgBinding bind),
            hang (text "} in ") 2 (ppr expr)]
 
-pprStgExpr (StgLetNoEscape bind expr)
-  = sep [hang (text "let-no-escape {")
+pprStgExpr (StgLetNoEscape ext bind expr)
+  = sep [hang (text "let-no-escape" <+> ppr ext <+> text "{")
                 2 (pprGenStgBinding bind),
            hang (text "} in ")
                 2 (ppr expr)]
@@ -805,7 +833,7 @@ pprStgExpr (StgCase expr bndr alt_type alts)
            nest 2 (vcat (map pprStgAlt alts)),
            char '}']
 
-pprStgAlt :: (Outputable (XRhsClosure pass)) => GenStgAlt pass -> SDoc
+pprStgAlt :: OutputablePass pass => GenStgAlt pass -> SDoc
 pprStgAlt (con, params, expr)
   = hang (hsep [ppr con, sep (map (pprBndr CasePatBind) params), text "->"])
          4 (ppr expr <> semi)
@@ -821,7 +849,7 @@ instance Outputable AltType where
   ppr (AlgAlt tc)     = text "Alg"    <+> ppr tc
   ppr (PrimAlt tc)    = text "Prim"   <+> ppr tc
 
-pprStgRhs :: (Outputable (XRhsClosure pass)) => GenStgRhs pass -> SDoc
+pprStgRhs :: OutputablePass pass => GenStgRhs pass -> SDoc
 
 -- special case
 pprStgRhs (StgRhsClosure ext cc upd_flag [{-no args-}] (StgApp func []))