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.

1 files changed, 41 insertions, 20 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