Remove StgLam

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.* codebase 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. Instead, we use a new intermediate
type in CoreToStg, PreStgRhs, to represent the RHS expression of a
binding.

11 files changed, 65 insertions, 67 deletions

diff --git a/compiler/GHC/CoreToStg.hs b/compiler/GHC/CoreToStg.hs
index bc890ea6cb..ee885eaacf 100644
--- a/compiler/GHC/CoreToStg.hs
+++ b/compiler/GHC/CoreToStg.hs
@@ -52,7 +52,6 @@ import GHC.Types.SrcLoc    ( mkGeneralSrcSpan )
 import GHC.Builtin.Names   ( unsafeEqualityProofName )
 
 import Control.Monad (ap)
-import Data.List.NonEmpty (nonEmpty, toList)
 import Data.Maybe (fromMaybe)
 import Data.Tuple (swap)
 import qualified Data.Set as Set
@@ -326,7 +325,7 @@ coreToTopStgRhs
         -> CtsM (StgRhs, CollectedCCs)
 
 coreToTopStgRhs dflags ccs this_mod (bndr, rhs)
-  = do { new_rhs <- coreToStgExpr rhs
+  = do { new_rhs <- coreToPreStgRhs rhs
 
        ; let (stg_rhs, ccs') =
                mkTopStgRhs dflags this_mod ccs bndr new_rhs
@@ -359,6 +358,10 @@ coreToTopStgRhs dflags ccs this_mod (bndr, rhs)
 -- Expressions
 -- ---------------------------------------------------------------------------
 
+-- coreToStgExpr panics if the input expression is a value lambda. CorePrep
+-- ensures that value lambdas only exist as the RHS of bindings, which we
+-- handle with the function coreToPreStgRhs.
+
 coreToStgExpr
         :: CoreExpr
         -> CtsM StgExpr
@@ -392,16 +395,13 @@ coreToStgExpr expr@(App _ _)
 coreToStgExpr expr@(Lam _ _)
   = let
         (args, body) = myCollectBinders expr
-        args'        = filterStgBinders args
     in
-    extendVarEnvCts [ (a, LambdaBound) | a <- args' ] $ do
-    body' <- coreToStgExpr body
-    let
-        result_expr = case nonEmpty args' of
-          Nothing     -> body'
-          Just args'' -> StgLam args'' body'
+    case filterStgBinders args of
 
-    return result_expr
+      [] -> coreToStgExpr body
+
+      _ -> pprPanic "coretoStgExpr" $
+        text "Unexpected value lambda:" $$ ppr expr
 
 coreToStgExpr (Tick tick expr)
   = do case tick of
@@ -674,23 +674,42 @@ coreToStgRhs :: (Id,CoreExpr)
              -> CtsM StgRhs
 
 coreToStgRhs (bndr, rhs) = do
-    new_rhs <- coreToStgExpr rhs
+    new_rhs <- coreToPreStgRhs rhs
     return (mkStgRhs bndr new_rhs)
 
+-- Represents the RHS of a binding for use with mk(Top)StgRhs.
+data PreStgRhs = PreStgRhs [Id] StgExpr -- The [Id] is empty for thunks
+
+-- Convert the RHS of a binding from Core to STG. This is a wrapper around
+-- coreToStgExpr that can handle value lambdas.
+coreToPreStgRhs :: CoreExpr -> CtsM PreStgRhs
+coreToPreStgRhs (Cast expr _) = coreToPreStgRhs expr
+coreToPreStgRhs expr@(Lam _ _) =
+    let
+        (args, body) = myCollectBinders expr
+        args'        = filterStgBinders args
+    in
+        extendVarEnvCts [ (a, LambdaBound) | a <- args' ] $ do
+          body' <- coreToStgExpr body
+          return (PreStgRhs args' body')
+coreToPreStgRhs expr = PreStgRhs [] <$> coreToStgExpr expr
+
 -- 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 -> PreStgRhs -> (StgRhs, CollectedCCs)
 
-mkTopStgRhs dflags this_mod ccs bndr rhs
-  | StgLam bndrs body <- rhs
-  = -- StgLam can't have empty arguments, so not CAF
+mkTopStgRhs dflags this_mod ccs bndr (PreStgRhs bndrs rhs)
+  | not (null bndrs)
+  = -- The list of arguments is non-empty, so not CAF
     ( StgRhsClosure noExtFieldSilent
                     dontCareCCS
                     ReEntrant
-                    (toList bndrs) body
+                    bndrs rhs
     , ccs )
 
+  -- After this point we know that `bndrs` is empty,
+  -- so this is not a function binding
   | StgConApp con args _ <- unticked_rhs
   , -- Dynamic StgConApps are updatable
     not (isDllConApp dflags this_mod con args)
@@ -732,14 +751,16 @@ 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 bndr rhs
-  | StgLam bndrs body <- rhs
+mkStgRhs :: Id -> PreStgRhs -> StgRhs
+mkStgRhs bndr (PreStgRhs bndrs rhs)
+  | not (null bndrs)
   = StgRhsClosure noExtFieldSilent
                   currentCCS
                   ReEntrant
-                  (toList bndrs) body
+                  bndrs rhs
 
+  -- After this point we know that `bndrs` is empty,
+  -- so this is not a function binding
   | isJoinId bndr -- must be a nullary join point
   = ASSERT(idJoinArity bndr == 0)
     StgRhsClosure noExtFieldSilent
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..742b29ee71 100644
--- a/compiler/GHC/Stg/Lint.hs
+++ b/compiler/GHC/Stg/Lint.hs
@@ -192,10 +192,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) $
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..25d01079df 100644
--- a/compiler/GHC/Stg/Syntax.hs
+++ b/compiler/GHC/Stg/Syntax.hs
@@ -89,8 +89,6 @@ import GHC.Types.RepType ( typePrimRep1 )
 import GHC.Utils.Misc
 import GHC.Utils.Panic
 
-import Data.List.NonEmpty ( NonEmpty, toList )
-
 {-
 ************************************************************************
 *                                                                      *
@@ -256,22 +254,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
 *                                                                      *
 ************************************************************************
@@ -436,6 +418,30 @@ important):
                         -- are not allocated.
         [StgArg]        -- Args
 
+{-
+Note Stg Passes
+~~~~~~~~~~~~~~~
+Here is a short summary of the STG pipeline and where we use the different
+StgPass data type indexes:
+
+  1. CoreToStg.Prep performs several transformations that prepare the desugared
+     and simplified core to be converted to STG. One of these transformations is
+     making it so that value lambdas only exist as the RHS of a binding.
+
+  2. CoreToStg converts the prepared core to STG, specifically GenStg*
+     parameterised by 'Vanilla.
+
+  3. Stg.Pipeline does a number of passes on the generated STG. One of these is
+     the lambda-lifting pass, which internally uses the 'LiftLams
+     parameterisation to store information for deciding whether or not to lift
+     each binding.
+
+  4. Stg.FVs annotates closures with their free variables. To store these
+     annotations we use the 'CodeGen parameterisation.
+
+  5. Stg.StgToCmm generates Cmm from the annotated STG.
+-}
+
 -- | Used as a data type index for the stgSyn AST
 data StgPass
   = Vanilla
@@ -709,11 +715,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
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 0048a4c9a2..7427547bf4 100644
--- a/compiler/GHC/StgToCmm/Expr.hs
+++ b/compiler/GHC/StgToCmm/Expr.hs
@@ -129,8 +129,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
 ------------------------------------------------------------------------