TTG-ify StgLamwip/leif/stg-ttg

This makes the StgLam constructor of StgExpr an extension using Trees
that Grow.  StgLam is used exclusively in the work of CoreToStg, but
there's nothing in the type of StgExpr that indicates this, so we're
forced throughout the Stg.* code to handle cases like:

case expr of
  ...
  StgLam lam -> panic "Unexpected StgLam"
  ...

This patch removes the StgLam constructor from the base StgExpr so these
cases no longer need to be handled, except in one case. We still need to
convert from SgStgRhs (the type with StgLam) to StgRhs (the type without
StgLam) after CoreToStg has finished. This conversion is safe because
CorePrep guarantees that no value lambdas exist except as the body of a
let binding.

11 files changed, 136 insertions, 67 deletions

diff --git a/compiler/GHC/CoreToStg.hs b/compiler/GHC/CoreToStg.hs
index 8082023ae7..2193577a3d 100644
--- a/compiler/GHC/CoreToStg.hs
+++ b/compiler/GHC/CoreToStg.hs
@@ -1,4 +1,6 @@
 {-# LANGUAGE CPP, DeriveFunctor #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE DataKinds #-}
 
 --
 -- (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
@@ -330,10 +332,11 @@ coreToTopStgRhs dflags ccs this_mod (bndr, rhs)
 
        ; let (stg_rhs, ccs') =
                mkTopStgRhs dflags this_mod ccs bndr new_rhs
+             stg_rhs' = coerceStgRhs stg_rhs
              stg_arity =
-               stgRhsArity stg_rhs
+               stgRhsArity stg_rhs'
 
-       ; return (ASSERT2( arity_ok stg_arity, mk_arity_msg stg_arity) stg_rhs,
+       ; return (ASSERT2( arity_ok stg_arity, mk_arity_msg stg_arity) stg_rhs',
                  ccs') }
   where
         -- It's vital that the arity on a top-level Id matches
@@ -355,13 +358,43 @@ coreToTopStgRhs dflags ccs this_mod (bndr, rhs)
                 text "Id arity:" <+> ppr id_arity,
                 text "STG arity:" <+> ppr stg_arity]
 
+    -- Coerces the intermediate SgStgRhs that is used in this module to StgRhs,
+    -- panicking if any StgLams are found. mkStgRhs and mkTopStgRhs remove any
+    -- lambdas that appear in the body of a let binding, and CorePrep ensures
+    -- that no other lambdas exist, so this should never panic.
+    coerceStgRhs :: SgStgRhs -> StgRhs
+    coerceStgRhs (StgRhsCon cc con args) = StgRhsCon cc con args
+    coerceStgRhs (StgRhsClosure ext cc upd_flag args body)
+      = StgRhsClosure ext cc upd_flag args (coerceStgExpr body)
+      where
+        coerceStgExpr :: SgStgExpr -> StgExpr
+        coerceStgExpr (StgApp func args) = StgApp func args
+        coerceStgExpr (StgLit lit) = StgLit lit
+        coerceStgExpr (StgConApp con args tys) = StgConApp con args tys
+        coerceStgExpr (StgOpApp op args ty) = StgOpApp op args ty
+        coerceStgExpr (StgCase scrut bndr alt_type alts)
+          = StgCase (coerceStgExpr scrut) bndr alt_type ((fmap . mapLast) coerceStgExpr alts)
+        coerceStgExpr (StgLet ext bind expr)
+          = StgLet ext (coerceStgBinding bind) (coerceStgExpr expr)
+        coerceStgExpr (StgLetNoEscape ext bind expr)
+          = StgLetNoEscape ext (coerceStgBinding bind) (coerceStgExpr expr)
+        coerceStgExpr (StgTick tick expr) = StgTick tick (coerceStgExpr expr)
+        coerceStgExpr (XStgExpr _) = pprPanic "CoreToStg.coerceStgExpr" (text "StgLam")
+
+        coerceStgBinding :: SgStgBinding -> StgBinding
+        coerceStgBinding (StgNonRec bndr rhs) = StgNonRec bndr (coerceStgRhs rhs)
+        coerceStgBinding (StgRec binds) = StgRec $ (fmap . fmap) coerceStgRhs binds
+
+        mapLast :: (a -> b) -> (c, d, a) -> (c, d, b)
+        mapLast f (c, d, a) = (c, d, f a)
+
 -- ---------------------------------------------------------------------------
 -- Expressions
 -- ---------------------------------------------------------------------------
 
 coreToStgExpr
         :: CoreExpr
-        -> CtsM StgExpr
+        -> CtsM SgStgExpr
 
 -- The second and third components can be derived in a simple bottom up pass, not
 -- dependent on any decisions about which variables will be let-no-escaped or
@@ -399,7 +432,7 @@ coreToStgExpr expr@(Lam _ _)
     let
         result_expr = case nonEmpty args' of
           Nothing     -> body'
-          Just args'' -> StgLam args'' body'
+          Just args'' -> XStgExpr $ StgLam args'' body'
 
     return result_expr
 
@@ -448,7 +481,7 @@ coreToStgExpr e0@(Case scrut bndr _ alts) = do
               text "STG:" $$ pprStgExpr panicStgPprOpts stg
       _ -> return stg
   where
-    vars_alt :: (AltCon, [Var], CoreExpr) -> CtsM (AltCon, [Var], StgExpr)
+    vars_alt :: (AltCon, [Var], CoreExpr) -> CtsM (AltCon, [Var], SgStgExpr)
     vars_alt (con, binders, rhs)
       | DataAlt c <- con, c == unboxedUnitDataCon
       = -- This case is a bit smelly.
@@ -515,7 +548,7 @@ mkStgAltType bndr alts
 coreToStgApp :: Id            -- Function
              -> [CoreArg]     -- Arguments
              -> [Tickish Id]  -- Debug ticks
-             -> CtsM StgExpr
+             -> CtsM SgStgExpr
 coreToStgApp f args ticks = do
     (args', ticks') <- coreToStgArgs args
     how_bound <- lookupVarCts f
@@ -624,7 +657,7 @@ coreToStgArgs (arg : args) = do         -- Non-type argument
 coreToStgLet
          :: CoreBind     -- bindings
          -> CoreExpr     -- body
-         -> CtsM StgExpr -- new let
+         -> CtsM SgStgExpr -- new let
 
 coreToStgLet bind body = do
     (bind2, body2)
@@ -650,7 +683,7 @@ coreToStgLet bind body = do
         = (binder, LetBound NestedLet (manifestArity rhs))
 
     vars_bind :: CoreBind
-              -> CtsM (StgBinding,
+              -> CtsM (SgStgBinding,
                        [(Id, HowBound)])  -- extension to environment
 
     vars_bind (NonRec binder rhs) = do
@@ -671,7 +704,7 @@ coreToStgLet bind body = do
               return (StgRec (binders `zip` rhss2), env_ext)
 
 coreToStgRhs :: (Id,CoreExpr)
-             -> CtsM StgRhs
+             -> CtsM SgStgRhs
 
 coreToStgRhs (bndr, rhs) = do
     new_rhs <- coreToStgExpr rhs
@@ -680,10 +713,10 @@ coreToStgRhs (bndr, rhs) = do
 -- Generate a top-level RHS. Any new cost centres generated for CAFs will be
 -- appended to `CollectedCCs` argument.
 mkTopStgRhs :: DynFlags -> Module -> CollectedCCs
-            -> Id -> StgExpr -> (StgRhs, CollectedCCs)
+            -> Id -> SgStgExpr -> (SgStgRhs, CollectedCCs)
 
 mkTopStgRhs dflags this_mod ccs bndr rhs
-  | StgLam bndrs body <- rhs
+  | XStgExpr (StgLam bndrs body) <- rhs
   = -- StgLam can't have empty arguments, so not CAF
     ( StgRhsClosure noExtFieldSilent
                     dontCareCCS
@@ -732,9 +765,9 @@ mkTopStgRhs dflags this_mod ccs bndr rhs
 
 -- Generate a non-top-level RHS. Cost-centre is always currentCCS,
 -- see Note [Cost-centre initialization plan].
-mkStgRhs :: Id -> StgExpr -> StgRhs
+mkStgRhs :: Id -> SgStgExpr -> SgStgRhs
 mkStgRhs bndr rhs
-  | StgLam bndrs body <- rhs
+  | XStgExpr (StgLam bndrs body) <- rhs
   = StgRhsClosure noExtFieldSilent
                   currentCCS
                   ReEntrant
diff --git a/compiler/GHC/Stg/CSE.hs b/compiler/GHC/Stg/CSE.hs
index 61a7824188..b9e6782f77 100644
--- a/compiler/GHC/Stg/CSE.hs
+++ b/compiler/GHC/Stg/CSE.hs
@@ -95,8 +95,6 @@ import GHC.Prelude
 import GHC.Core.DataCon
 import GHC.Types.Id
 import GHC.Stg.Syntax
-import GHC.Utils.Outputable
-import GHC.Utils.Panic
 import GHC.Types.Basic (isWeakLoopBreaker)
 import GHC.Types.Var.Env
 import GHC.Core (AltCon(..))
@@ -312,8 +310,6 @@ stgCseExpr _ (StgLit lit)
 stgCseExpr env (StgOpApp op args tys)
     = StgOpApp op args' tys
   where args' = substArgs env args
-stgCseExpr _ (StgLam _ _)
-    = pprPanic "stgCseExp" (text "StgLam")
 stgCseExpr env (StgTick tick body)
     = let body' = stgCseExpr env body
       in StgTick tick body'
diff --git a/compiler/GHC/Stg/DepAnal.hs b/compiler/GHC/Stg/DepAnal.hs
index 223ab0c5bb..9bf4249f6f 100644
--- a/compiler/GHC/Stg/DepAnal.hs
+++ b/compiler/GHC/Stg/DepAnal.hs
@@ -91,8 +91,6 @@ annTopBindingsDeps this_mod bs = zip bs (map top_bind bs)
       args bounds as
     expr bounds (StgOpApp _ as _) =
       args bounds as
-    expr _ lam@StgLam{} =
-      pprPanic "annTopBindingsDeps" (text "Found lambda:" $$ pprStgExpr panicStgPprOpts lam)
     expr bounds (StgCase scrut scrut_bndr _ as) =
       expr bounds scrut `unionVarSet`
         alts (extendVarSet bounds scrut_bndr) as
diff --git a/compiler/GHC/Stg/FVs.hs b/compiler/GHC/Stg/FVs.hs
index 211a0cb315..bd699a1fe1 100644
--- a/compiler/GHC/Stg/FVs.hs
+++ b/compiler/GHC/Stg/FVs.hs
@@ -48,9 +48,7 @@ import GHC.Stg.Syntax
 import GHC.Types.Id
 import GHC.Types.Var.Set
 import GHC.Core    ( Tickish(Breakpoint) )
-import GHC.Utils.Outputable
 import GHC.Utils.Misc
-import GHC.Utils.Panic
 
 import Data.Maybe ( mapMaybe )
 
@@ -128,7 +126,6 @@ expr env = go
     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
     go (StgCase scrut bndr ty alts) = (StgCase scrut' bndr ty alts', fvs)
       where
         (scrut', scrut_fvs) = go scrut
diff --git a/compiler/GHC/Stg/Lift.hs b/compiler/GHC/Stg/Lift.hs
index 27e63f9313..8f2337120e 100644
--- a/compiler/GHC/Stg/Lift.hs
+++ b/compiler/GHC/Stg/Lift.hs
@@ -229,7 +229,6 @@ liftExpr (StgApp f args) = do
   pure (StgApp f' top_lvl_args)
 liftExpr (StgConApp con args tys) = StgConApp con <$> traverse liftArgs args <*> pure tys
 liftExpr (StgOpApp op args ty) = StgOpApp op <$> traverse liftArgs args <*> pure ty
-liftExpr (StgLam _ _) = pprPanic "stgLiftLams" (text "StgLam")
 liftExpr (StgCase scrut info ty alts) = do
   scrut' <- liftExpr scrut
   withSubstBndr (binderInfoBndr info) $ \bndr' -> do
diff --git a/compiler/GHC/Stg/Lift/Analysis.hs b/compiler/GHC/Stg/Lift/Analysis.hs
index 5aef95c008..314e010ead 100644
--- a/compiler/GHC/Stg/Lift/Analysis.hs
+++ b/compiler/GHC/Stg/Lift/Analysis.hs
@@ -34,7 +34,6 @@ import qualified GHC.StgToCmm.ArgRep  as StgToCmm.ArgRep
 import qualified GHC.StgToCmm.Closure as StgToCmm.Closure
 import qualified GHC.StgToCmm.Layout  as StgToCmm.Layout
 import GHC.Utils.Outputable
-import GHC.Utils.Panic
 import GHC.Utils.Misc
 import GHC.Types.Var.Set
 
@@ -223,7 +222,6 @@ tagSkeletonExpr (StgApp f args)
       -- argument occurrences, see "GHC.Stg.Lift.Analysis#arg_occs".
       | null args = unitVarSet f
       | otherwise = mkArgOccs args
-tagSkeletonExpr (StgLam _ _) = pprPanic "stgLiftLams" (text "StgLam")
 tagSkeletonExpr (StgCase scrut bndr ty alts)
   = (skel, arg_occs, StgCase scrut' bndr' ty alts')
   where
diff --git a/compiler/GHC/Stg/Lint.hs b/compiler/GHC/Stg/Lint.hs
index 1485a11458..8ba754971e 100644
--- a/compiler/GHC/Stg/Lint.hs
+++ b/compiler/GHC/Stg/Lint.hs
@@ -33,13 +33,17 @@ basic properties listed above.
 -}
 
 {-# LANGUAGE ScopedTypeVariables, FlexibleContexts, TypeFamilies,
-  DeriveFunctor #-}
+  DeriveFunctor, TypeApplications, CPP #-}
 
 module GHC.Stg.Lint ( lintStgTopBindings ) where
 
 import GHC.Prelude
 
 import GHC.Stg.Syntax
+import GHC.Hs.Extension ( NoExtCon )
+#if __GLASGOW_HASKELL__ < 811
+import GHC.Hs.Extension ( noExtCon )
+#endif
 
 import GHC.Driver.Session
 import GHC.Data.Bag         ( Bag, emptyBag, isEmptyBag, snocBag, bagToList )
@@ -60,7 +64,7 @@ import qualified GHC.Utils.Error as Err
 import Control.Applicative ((<|>))
 import Control.Monad
 
-lintStgTopBindings :: forall a . (OutputablePass a, BinderP a ~ Id)
+lintStgTopBindings :: forall a . (OutputablePass a, BinderP a ~ Id, XXStgExpr a ~ NoExtCon)
                    => DynFlags
                    -> Module -- ^ module being compiled
                    -> Bool   -- ^ have we run Unarise yet?
@@ -108,7 +112,7 @@ lintStgVar :: Id -> LintM ()
 lintStgVar id = checkInScope id
 
 lintStgBinds
-    :: (OutputablePass a, BinderP a ~ Id)
+    :: (OutputablePass a, BinderP a ~ Id, XXStgExpr a ~ NoExtCon)
     => TopLevelFlag -> GenStgBinding a -> LintM [Id] -- Returns the binders
 lintStgBinds top_lvl (StgNonRec binder rhs) = do
     lint_binds_help top_lvl (binder,rhs)
@@ -122,7 +126,7 @@ lintStgBinds top_lvl (StgRec pairs)
     binders = [b | (b,_) <- pairs]
 
 lint_binds_help
-    :: (OutputablePass a, BinderP a ~ Id)
+    :: (OutputablePass a, BinderP a ~ Id, XXStgExpr a ~ NoExtCon)
     => TopLevelFlag
     -> (Id, GenStgRhs a)
     -> LintM ()
@@ -153,7 +157,7 @@ checkNoCurrentCCS rhs = do
          -> addErrL (text "Top-level StgRhsCon with CurrentCCS" $$ rhs')
       _ -> return ()
 
-lintStgRhs :: (OutputablePass a, BinderP a ~ Id) => GenStgRhs a -> LintM ()
+lintStgRhs :: (OutputablePass a, BinderP a ~ Id, XXStgExpr a ~ NoExtCon) => GenStgRhs a -> LintM ()
 
 lintStgRhs (StgRhsClosure _ _ _ [] expr)
   = lintStgExpr expr
@@ -171,7 +175,7 @@ lintStgRhs rhs@(StgRhsCon _ con args) = do
     mapM_ lintStgArg args
     mapM_ checkPostUnariseConArg args
 
-lintStgExpr :: (OutputablePass a, BinderP a ~ Id) => GenStgExpr a -> LintM ()
+lintStgExpr :: (OutputablePass a, BinderP a ~ Id, XXStgExpr a ~ NoExtCon) => GenStgExpr a -> LintM ()
 
 lintStgExpr (StgLit _) = return ()
 
@@ -192,10 +196,6 @@ lintStgExpr app@(StgConApp con args _arg_tys) = do
 lintStgExpr (StgOpApp _ args _) =
     mapM_ lintStgArg args
 
-lintStgExpr lam@(StgLam _ _) = do
-    opts <- getStgPprOpts
-    addErrL (text "Unexpected StgLam" <+> pprStgExpr opts lam)
-
 lintStgExpr (StgLet _ binds body) = do
     binders <- lintStgBinds NotTopLevel binds
     addLoc (BodyOfLetRec binders) $
@@ -218,8 +218,12 @@ lintStgExpr (StgCase scrut bndr alts_type alts) = do
 
     addInScopeVars [bndr | in_scope] (mapM_ lintAlt alts)
 
+#if __GLASGOW_HASKELL__ < 811
+lintStgExpr (XStgExpr xStgExpr) = noExtCon xStgExpr
+#endif
+
 lintAlt
-    :: (OutputablePass a, BinderP a ~ Id)
+    :: (OutputablePass a, BinderP a ~ Id, XXStgExpr a ~ NoExtCon)
     => (AltCon, [Id], GenStgExpr a) -> LintM ()
 
 lintAlt (DEFAULT, _, rhs) =
diff --git a/compiler/GHC/Stg/Stats.hs b/compiler/GHC/Stg/Stats.hs
index 329f319a47..0f806a3175 100644
--- a/compiler/GHC/Stg/Stats.hs
+++ b/compiler/GHC/Stg/Stats.hs
@@ -32,7 +32,6 @@ import GHC.Prelude
 import GHC.Stg.Syntax
 
 import GHC.Types.Id (Id)
-import GHC.Utils.Panic
 
 import Data.Map (Map)
 import qualified Data.Map as Map
@@ -169,5 +168,3 @@ statExpr (StgCase expr _ _ alts)
   where
     stat_alts alts
         = combineSEs (map statExpr [ e | (_,_,e) <- alts ])
-
-statExpr (StgLam {}) = panic "statExpr StgLam"
diff --git a/compiler/GHC/Stg/Syntax.hs b/compiler/GHC/Stg/Syntax.hs
index b38c2f1ab0..37f834d2bb 100644
--- a/compiler/GHC/Stg/Syntax.hs
+++ b/compiler/GHC/Stg/Syntax.hs
@@ -3,7 +3,9 @@
 {-# LANGUAGE DataKinds            #-}
 {-# LANGUAGE DeriveDataTypeable   #-}
 {-# LANGUAGE FlexibleContexts     #-}
+{-# LANGUAGE GADTs                #-}
 {-# LANGUAGE LambdaCase           #-}
+{-# LANGUAGE TypeApplications     #-}
 {-# LANGUAGE TypeFamilies         #-}
 {-# LANGUAGE UndecidableInstances #-}
 
@@ -25,9 +27,11 @@ module GHC.Stg.Syntax (
         GenStgTopBinding(..), GenStgBinding(..), GenStgExpr(..), GenStgRhs(..),
         GenStgAlt, AltType(..),
 
-        StgPass(..), BinderP, XRhsClosure, XLet, XLetNoEscape,
+        StgPass(..), BinderP, XRhsClosure, XLet, XLetNoEscape, XXStgExpr,
+        StgPassWitness(..), IsStgPass,
         NoExtFieldSilent, noExtFieldSilent,
         OutputablePass,
+        StgLam(..),
 
         UpdateFlag(..), isUpdatable,
 
@@ -40,6 +44,9 @@ module GHC.Stg.Syntax (
         -- a set of synonyms for the lambda lifting parameterisation
         LlStgTopBinding, LlStgBinding, LlStgExpr, LlStgRhs, LlStgAlt,
 
+        -- a set of synonyms for the core to stg parameterisation
+        SgStgTopBinding, SgStgBinding, SgStgExpr, SgStgRhs, SgStgAlt,
+
         -- a set of synonyms to distinguish in- and out variants
         InStgArg,  InStgTopBinding,  InStgBinding,  InStgExpr,  InStgRhs,  InStgAlt,
         OutStgArg, OutStgTopBinding, OutStgBinding, OutStgExpr, OutStgRhs, OutStgAlt,
@@ -73,6 +80,10 @@ import Data.Data   ( Data )
 import Data.List   ( intersperse )
 import GHC.Core.DataCon
 import GHC.Driver.Session
+import GHC.Hs.Extension ( NoExtCon )
+#if __GLASGOW_HASKELL__ < 811
+import GHC.Hs.Extension ( noExtCon )
+#endif
 import GHC.Types.ForeignCall ( ForeignCall )
 import GHC.Types.Id
 import GHC.Types.Name        ( isDynLinkName )
@@ -256,22 +267,6 @@ literals.
 {-
 ************************************************************************
 *                                                                      *
-StgLam
-*                                                                      *
-************************************************************************
-
-StgLam is used *only* during CoreToStg's work. Before CoreToStg has finished it
-encodes (\x -> e) as (let f = \x -> e in f) TODO: Encode this via an extension
-to GenStgExpr à la TTG.
--}
-
-  | StgLam
-        (NonEmpty (BinderP pass))
-        StgExpr    -- Body of lambda
-
-{-
-************************************************************************
-*                                                                      *
 GenStgExpr: case-expressions
 *                                                                      *
 ************************************************************************
@@ -386,6 +381,8 @@ Finally for @hpc@ expressions we introduce a new STG construct.
     (Tickish Id)
     (GenStgExpr pass)       -- sub expression
 
+  | XStgExpr !(XXStgExpr pass)
+
 -- END of GenStgExpr
 
 {-
@@ -441,6 +438,28 @@ data StgPass
   = Vanilla
   | LiftLams
   | CodeGen
+  | StgGen
+
+data StgPassWitness (pass :: StgPass) where
+  VanillaWitness :: StgPassWitness 'Vanilla
+  LiftLamsWitness :: StgPassWitness 'LiftLams
+  CodeGenWitness :: StgPassWitness 'CodeGen
+  StgGenWitness :: StgPassWitness 'StgGen
+
+class IsStgPass (pass :: StgPass) where
+  stgPass :: StgPassWitness pass
+
+instance IsStgPass 'Vanilla where
+  stgPass = VanillaWitness
+
+instance IsStgPass 'LiftLams where
+  stgPass = LiftLamsWitness
+
+instance IsStgPass 'CodeGen where
+  stgPass = CodeGenWitness
+
+instance IsStgPass 'StgGen where
+  stgPass = StgGenWitness
 
 -- | Like 'GHC.Hs.Extension.NoExtField', but with an 'Outputable' instance that
 -- returns 'empty'.
@@ -462,19 +481,37 @@ noExtFieldSilent = NoExtFieldSilent
 type family BinderP (pass :: StgPass)
 type instance BinderP 'Vanilla = Id
 type instance BinderP 'CodeGen = Id
+type instance BinderP 'StgGen = Id
 
 type family XRhsClosure (pass :: StgPass)
 type instance XRhsClosure 'Vanilla = NoExtFieldSilent
 -- | Code gen needs to track non-global free vars
 type instance XRhsClosure 'CodeGen = DIdSet
+type instance XRhsClosure 'StgGen = NoExtFieldSilent
 
 type family XLet (pass :: StgPass)
 type instance XLet 'Vanilla = NoExtFieldSilent
 type instance XLet 'CodeGen = NoExtFieldSilent
+type instance XLet 'StgGen = NoExtFieldSilent
 
 type family XLetNoEscape (pass :: StgPass)
 type instance XLetNoEscape 'Vanilla = NoExtFieldSilent
 type instance XLetNoEscape 'CodeGen = NoExtFieldSilent
+type instance XLetNoEscape 'StgGen = NoExtFieldSilent
+
+type family XXStgExpr (pass :: StgPass)
+type instance XXStgExpr 'Vanilla = NoExtCon
+type instance XXStgExpr 'CodeGen = NoExtCon
+type instance XXStgExpr 'LiftLams = NoExtCon
+type instance XXStgExpr 'StgGen = StgLam
+
+-- StgLam is used *only* during CoreToStg's work. Before CoreToStg has finished
+-- it encodes (\x -> e) as (let f = \x -> e in f). We define this here rather
+-- than in CoreToStg because we need it for pretty-printing StgExpr and
+-- importing CoreToStg would create an import cycle.
+data StgLam = StgLam
+  (NonEmpty Id)
+  SgStgExpr
 
 stgRhsArity :: StgRhs -> Int
 stgRhsArity (StgRhsClosure _ _ _ bndrs _)
@@ -541,6 +578,12 @@ type CgStgExpr       = GenStgExpr       'CodeGen
 type CgStgRhs        = GenStgRhs        'CodeGen
 type CgStgAlt        = GenStgAlt        'CodeGen
 
+type SgStgTopBinding = GenStgTopBinding 'StgGen
+type SgStgBinding    = GenStgBinding    'StgGen
+type SgStgExpr       = GenStgExpr       'StgGen
+type SgStgRhs        = GenStgRhs        'StgGen
+type SgStgAlt        = GenStgAlt        'StgGen
+
 {- Many passes apply a substitution, and it's very handy to have type
    synonyms to remind us whether or not the substitution has been applied.
    See GHC.Core for precedence in Core land
@@ -645,6 +688,7 @@ type OutputablePass pass =
   , Outputable (XLetNoEscape pass)
   , Outputable (XRhsClosure pass)
   , OutputableBndr (BinderP pass)
+  , IsStgPass pass
   )
 
 -- | STG pretty-printing options
@@ -701,7 +745,7 @@ pprStgArg :: StgArg -> SDoc
 pprStgArg (StgVarArg var) = ppr var
 pprStgArg (StgLitArg con) = ppr con
 
-pprStgExpr :: OutputablePass pass => StgPprOpts -> GenStgExpr pass -> SDoc
+pprStgExpr :: forall pass. OutputablePass pass => StgPprOpts -> GenStgExpr pass -> SDoc
 pprStgExpr opts e = case e of
                            -- special case
    StgLit lit           -> ppr lit
@@ -709,11 +753,6 @@ pprStgExpr opts e = case e of
    StgApp func args     -> hang (ppr func) 4 (interppSP args)
    StgConApp con args _ -> hsep [ ppr con, brackets (interppSP args) ]
    StgOpApp op args _   -> hsep [ pprStgOp op, brackets (interppSP args)]
-   StgLam bndrs body    -> let ppr_list = brackets . fsep . punctuate comma
-                           in sep [ char '\\' <+> ppr_list (map (pprBndr LambdaBind) (toList bndrs))
-                                      <+> text "->"
-                                  , pprStgExpr opts body
-                                  ]
 
 -- special case: let v = <very specific thing>
 --               in
@@ -786,6 +825,19 @@ pprStgExpr opts e = case e of
              , char '}'
              ]
 
+   XStgExpr xStgExpr -> case stgPass @pass of
+#if __GLASGOW_HASKELL__ < 811
+     VanillaWitness -> noExtCon xStgExpr
+     LiftLamsWitness -> noExtCon xStgExpr
+     CodeGenWitness -> noExtCon xStgExpr
+#endif
+     StgGenWitness -> case xStgExpr of
+       StgLam bndrs body -> let ppr_list = brackets . fsep . punctuate comma
+         in sep [ char '\\' <+> ppr_list (map (pprBndr LambdaBind) (toList bndrs))
+                    <+> text "->"
+                , pprStgExpr opts body
+                ]
+
 
 pprStgAlt :: OutputablePass pass => StgPprOpts -> Bool -> GenStgAlt pass -> SDoc
 pprStgAlt opts indent (con, params, expr)
diff --git a/compiler/GHC/Stg/Unarise.hs b/compiler/GHC/Stg/Unarise.hs
index eb4c968f5b..40dff5f33b 100644
--- a/compiler/GHC/Stg/Unarise.hs
+++ b/compiler/GHC/Stg/Unarise.hs
@@ -336,9 +336,6 @@ unariseExpr rho (StgConApp dc args ty_args)
 unariseExpr rho (StgOpApp op args ty)
   = return (StgOpApp op (unariseFunArgs rho args) ty)
 
-unariseExpr _ e@StgLam{}
-  = pprPanic "unariseExpr: found lambda" (pprStgExpr panicStgPprOpts e)
-
 unariseExpr rho (StgCase scrut bndr alt_ty alts)
   -- tuple/sum binders in the scrutinee can always be eliminated
   | StgApp v [] <- scrut
diff --git a/compiler/GHC/StgToCmm/Expr.hs b/compiler/GHC/StgToCmm/Expr.hs
index eb56a6ad09..49451ae648 100644
--- a/compiler/GHC/StgToCmm/Expr.hs
+++ b/compiler/GHC/StgToCmm/Expr.hs
@@ -98,8 +98,6 @@ cgExpr (StgLetNoEscape _ binds expr) =
 cgExpr (StgCase expr bndr alt_type alts) =
   cgCase expr bndr alt_type alts
 
-cgExpr (StgLam {}) = panic "cgExpr: StgLam"
-
 ------------------------------------------------------------------------
 --              Let no escape
 ------------------------------------------------------------------------