Improve performance of eta expansion

Eta expansion was taking ages on T18223.  This patch

* Aggressively squash reflexive casts in etaInfoApp.
  See Note [Check for reflexive casts in eta expansion]
  These changes decreased compile-time allocation by 80%!

* Passes the Simplifier's in-scope set to etaExpandAT, so we
  don't need to recompute it.  (This alone saved 10% of compile
  time.)

  Annoyingly several functions in the Simplifier (namely
  makeTrivialBinding and friends) need to get SimplEnv, rather
  than SimplMode, but that is no big deal.

Lots of small changes in compile-time allocation, less than 1%
and in both directions.  A couple of bigger changes, including
the rather delicate T18223

T12425(optasm) ghc/alloc    98448216.0    97121224.0  -1.3% GOOD
T18223(normal) ghc/alloc  5454689676.0  1138238008.0 -79.1% GOOD

Metric Decrease:
    T12425
    T18223

4 files changed, 152 insertions, 73 deletions

diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index ba3da23d93..aa0cf29754 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -73,7 +73,7 @@ module GHC.Core.Coercion (
 
         coToMCo, mkTransMCo, mkTransMCoL, mkTransMCoR, mkCastTyMCo, mkSymMCo,
         mkHomoForAllMCo, mkFunResMCo, mkPiMCos,
-        isReflMCo,
+        isReflMCo, checkReflexiveMCo,
 
         -- ** Coercion variables
         mkCoVar, isCoVar, coVarName, setCoVarName, setCoVarUnique,
@@ -311,6 +311,11 @@ coToMCo :: Coercion -> MCoercion
 coToMCo co | isReflCo co = MRefl
            | otherwise   = MCo co
 
+checkReflexiveMCo :: MCoercion -> MCoercion
+checkReflexiveMCo MRefl                       = MRefl
+checkReflexiveMCo (MCo co) | isReflexiveCo co = MRefl
+                           | otherwise        = MCo co
+
 -- | Tests if this MCoercion is obviously generalized reflexive
 -- Guaranteed to work very quickly.
 isGReflMCo :: MCoercion -> Bool
@@ -420,6 +425,10 @@ decomposeFunCo :: HasDebugCallStack
 -- Expects co :: (s1 -> t1) ~ (s2 -> t2)
 -- Returns (co1 :: s1~s2, co2 :: t1~t2)
 -- See Note [Function coercions] for the "3" and "4"
+
+decomposeFunCo _ (FunCo _ w co1 co2) = (w, co1, co2)
+   -- Short-circuits the calls to mkNthCo
+
 decomposeFunCo r co = assertPpr all_ok (ppr co)
                       (mkNthCo Nominal 0 co, mkNthCo r 3 co, mkNthCo r 4 co)
   where
diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index 53f2c28213..fcc35d76b5 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -592,7 +592,7 @@ arityTypeArity (AT oss _) = length oss
 
 -- | True <=> eta-expansion will add at least one lambda
 expandableArityType :: ArityType -> Bool
-expandableArityType at = arityTypeArity at /= 0
+expandableArityType at = arityTypeArity at > 0
 
 -- | See Note [Dead ends] in "GHC.Types.Demand".
 -- Bottom implies a dead end.
@@ -667,6 +667,7 @@ findRhsArity dflags bndr rhs old_arity
       where
         env = extendSigEnv (findRhsArityEnv dflags) bndr at
 
+
 {-
 Note [Arity analysis]
 ~~~~~~~~~~~~~~~~~~~~~
@@ -1293,12 +1294,22 @@ see Note [The one-shot state monad trick] in GHC.Utils.Monad.
 -- We should have that:
 --
 -- > ty = exprType e = exprType e'
-etaExpand   :: Arity     -> CoreExpr -> CoreExpr
-etaExpandAT :: ArityType -> CoreExpr -> CoreExpr
 
-etaExpand   n          orig_expr = eta_expand (replicate n NoOneShotInfo) orig_expr
-etaExpandAT (AT oss _) orig_expr = eta_expand oss                         orig_expr
-                           -- See Note [Eta expansion with ArityType]
+etaExpand :: Arity -> CoreExpr -> CoreExpr
+etaExpand n orig_expr
+  = eta_expand in_scope (replicate n NoOneShotInfo) orig_expr
+  where
+    in_scope = {-#SCC "eta_expand:in-scopeX" #-}
+               mkInScopeSet (exprFreeVars orig_expr)
+
+etaExpandAT :: InScopeSet -> ArityType -> CoreExpr -> CoreExpr
+-- See Note [Eta expansion with ArityType]
+--
+-- We pass in the InScopeSet from the simplifier to avoid recomputing
+-- it here, which can be jolly expensive if the casts are big
+-- In #18223 it took 10% of compile time just to do the exprFreeVars!
+etaExpandAT in_scope (AT oss _) orig_expr
+  = eta_expand in_scope oss orig_expr
 
 -- etaExpand arity e = res
 -- Then 'res' has at least 'arity' lambdas at the top
@@ -1311,12 +1322,12 @@ etaExpandAT (AT oss _) orig_expr = eta_expand oss                         orig_e
 -- would return
 --      (/\b. \y::a -> E b y)
 
-eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
-eta_expand one_shots (Cast expr co)
-  = mkCast (eta_expand one_shots expr) co
+eta_expand :: InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
+eta_expand in_scope one_shots (Cast expr co)
+  = Cast (eta_expand in_scope one_shots expr) co
 
-eta_expand one_shots orig_expr
-  = go one_shots [] orig_expr
+eta_expand in_scope one_shots orig_expr
+  = go in_scope one_shots [] orig_expr
   where
       -- Strip off existing lambdas and casts before handing off to mkEtaWW
       -- This is mainly to avoid spending time cloning binders and substituting
@@ -1324,20 +1335,20 @@ eta_expand one_shots orig_expr
       -- with casts here, apart from the topmost one, and they are rare, so
       -- if we find one we just hand off to mkEtaWW anyway
       -- Note [Eta expansion and SCCs]
-    go [] _ _ = orig_expr  -- Already has the specified arity; no-op
+    go _ [] _ _ = orig_expr  -- Already has the specified arity; no-op
 
-    go oss@(_:oss1) vs (Lam v body)
-      | isTyVar v = go oss  (v:vs) body
-      | otherwise = go oss1 (v:vs) body
+    go in_scope oss@(_:oss1) vs (Lam v body)
+      | isTyVar v = go (in_scope `extendInScopeSet` v) oss  (v:vs) body
+      | otherwise = go (in_scope `extendInScopeSet` v) oss1 (v:vs) body
 
-    go oss rev_vs expr
-      = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, pprEtaInfos etas]) $
-        retick $ etaInfoAbs top_eis $
-                 etaInfoApp in_scope' sexpr eis
+    go in_scope oss rev_vs expr
+      = -- pprTrace "ee" (vcat [ppr in_scope', ppr top_bndrs, ppr eis]) $
+        retick $
+        etaInfoAbs top_eis $
+        etaInfoApp in_scope' sexpr eis
       where
-          in_scope = mkInScopeSet (exprFreeVars expr)
           (in_scope', eis@(EI eta_bndrs mco))
-              = mkEtaWW oss (ppr orig_expr) in_scope (exprType expr)
+                    = mkEtaWW oss (ppr orig_expr) in_scope (exprType expr)
           top_bndrs = reverse rev_vs
           top_eis   = EI (top_bndrs ++ eta_bndrs) (mkPiMCos top_bndrs mco)
 
@@ -1380,23 +1391,29 @@ This what eta_expand does.  We do it in two steps:
 To a first approximation EtaInfo is just [Var].  But
 casts complicate the question.  If we have
    newtype N a = MkN (S -> a)
+     axN :: N a  ~  S -> a
 and
-   ty = N (N Int)
-then the eta-expansion must look like
-        (\x (\y. ((e |> co1) x) |> co2) y)
-           |> sym co2)
-        |> sym co1
+   e :: N (N Int)
+then the eta-expansion should look like
+   (\(x::S) (y::S) -> e |> co x y) |> sym co
 where
-  co1 :: N (N Int) ~ S -> N Int
-  co2 :: N Int     ~ S -> Int
+  co :: N (N Int) ~ S -> S -> Int
+  co = axN @(N Int) ; (S -> axN @Int)
 
-Blimey!  Look at all those casts.  Moreover, if the type
-is very deeply nested (as happens in #18223), the repetition
+We want to get one cast, at the top, to account for all those
+nested newtypes. This is expressed by the EtaInfo type:
+
+   data EtaInfo = EI [Var] MCoercionR
+
+Note [Check for reflexive casts in eta expansion]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+It turns out that the casts created by teh above mechanism are often Refl.
+When casts are very deeply nested (as happens in #18223), the repetition
 of types can make the overall term very large.  So there is a big
 payoff in cancelling out casts aggressively wherever possible.
 (See also Note [No crap in eta-expanded code].)
 
-This matters a lot in etaEInfoApp, where we
+This matters particularly in etaInfoApp, where we
 * Do beta-reduction on the fly
 * Use getArg_maybe to get a cast out of the way,
   so that we can do beta reduction
@@ -1413,23 +1430,56 @@ away.
 #18223 was a dramatic example in which the intermediate term was
 grotesquely huge, even though the next Simplifier iteration squashed
 it.  Better to kill it at birth.
+
+The crucial spots in etaInfoApp are:
+* `checkReflexiveMCo` in the (Cast e co) case of `go`
+* `checkReflexiveMCo` in `pushCoArg`
+* Less important: checkReflexiveMCo in the final case of `go`
+Collectively these make a factor-of-5 difference to the total
+allocation of T18223, so take care if you change this stuff!
+
+Example:
+   newtype N = MkN (Y->Z)
+   f :: X -> N
+   f = \(x::X). ((\(y::Y). blah) |> fco)
+
+where fco :: (Y->Z) ~ N
+
+mkEtaWW makes an EtaInfo of (EI [(eta1:X), (eta2:Y)] eta_co
+  where
+    eta_co :: (X->N) ~ (X->Y->Z)
+    eta_co =  (<X> -> nco)
+    nco :: N ~ (Y->Z)  -- Comes from topNormaliseNewType_maybe
+
+Now, when we push that eta_co inward in etaInfoApp:
+* In the (Cast e co) case, the 'fco' and 'nco' will meet, and
+  should cancel.
+* When we meet the (\y.e) we want no cast on the y.
+
 -}
 
 --------------
 data EtaInfo = EI [Var] MCoercionR
 
--- EI bs co
--- Abstraction:  (\b1 b2 .. bn. []) |> sym co
--- Application:  ([] |> co) b1 b2 .. bn
+-- (EI bs co) describes a particular eta-expansion, as follows:
+--  Abstraction:  (\b1 b2 .. bn. []) |> sym co
+--  Application:  ([] |> co) b1 b2 .. bn
 --
 --    e :: T    co :: T ~ (t1 -> t2 -> .. -> tn -> tr)
 --    e = (\b1 b2 ... bn. (e |> co) b1 b2 .. bn) |> sym co
 
+instance Outputable EtaInfo where
+  ppr (EI vs mco) = text "EI" <+> ppr vs <+> parens (ppr mco)
+
 
 etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
--- (etaInfoApp s e eis) returns something equivalent to
---             (substExpr s e `appliedto` eis)
+-- (etaInfoApp s e (EI bs mco) returns something equivalent to
+--             ((substExpr s e) |> mco b1 .. bn)
 -- See Note [The EtaInfo mechanism]
+--
+-- NB: With very deeply nested casts, this function can be expensive
+--     In T18223, this function alone costs 15% of allocation, all
+--     spent in the calls to substExprSC and substBindSC
 
 etaInfoApp in_scope expr eis
   = go (mkEmptySubst in_scope) expr eis
@@ -1442,7 +1492,10 @@ etaInfoApp in_scope expr eis
       = Tick (substTickish subst t) (go subst e eis)
 
     go subst (Cast e co) (EI bs mco)
-      = go subst e (EI bs (Core.substCo subst co `mkTransMCoR` mco))
+      = go subst e (EI bs mco')
+      where
+        mco' = checkReflexiveMCo (Core.substCo subst co `mkTransMCoR` mco)
+               -- See Note [Check for reflexive casts in eta-expansion]
 
     go subst (Case e b ty alts) eis
       = Case (Core.substExprSC subst e) b1 ty' alts'
@@ -1467,8 +1520,9 @@ etaInfoApp in_scope expr eis
       = go (Core.extendSubst subst v arg) e (EI bs mco')
 
     -- Stop pushing down; just wrap the expression up
+    -- See Note [Check for reflexive casts in eta expansion]
     go subst e (EI bs mco) = Core.substExprSC subst e
-                             `mkCastMCo` mco
+                             `mkCastMCo` checkReflexiveMCo mco
                              `mkVarApps` bs
 
 --------------
@@ -1617,10 +1671,17 @@ pushCoArg :: CoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
 -- This may fail, e.g. if (fun :: N) where N is a newtype
 -- C.f. simplCast in GHC.Core.Opt.Simplify
 -- 'co' is always Representational
-pushCoArg co (Type ty) = do { (ty', m_co') <- pushCoTyArg co ty
-                            ; return (Type ty', m_co') }
-pushCoArg co val_arg   = do { (arg_co, m_co') <- pushCoValArg co
-                            ; return (val_arg `mkCastMCo` arg_co, m_co') }
+pushCoArg co arg
+  | Type ty <- arg
+  = do { (ty', m_co') <- pushCoTyArg co ty
+       ; return (Type ty', m_co') }
+  | otherwise
+  = do { (arg_mco, m_co') <- pushCoValArg co
+       ; let arg_mco' = checkReflexiveMCo arg_mco
+             -- checkReflexiveMCo: see Note [Check for reflexive casts in eta expansion]
+             -- The coercion is very often (arg_co -> res_co), but without
+             -- the argument coercion actually being ReflCo
+       ; return (arg `mkCastMCo` arg_mco', m_co') }
 
 pushCoTyArg :: CoercionR -> Type -> Maybe (Type, MCoercionR)
 -- We have (fun |> co) @ty
diff --git a/compiler/GHC/Core/Opt/Simplify.hs b/compiler/GHC/Core/Opt/Simplify.hs
index 406c6ed1a0..ac55c68ded 100644
--- a/compiler/GHC/Core/Opt/Simplify.hs
+++ b/compiler/GHC/Core/Opt/Simplify.hs
@@ -362,7 +362,7 @@ simplLazyBind env top_lvl is_rec bndr bndr1 rhs rhs_se
         -- ANF-ise a constructor or PAP rhs
         -- We get at most one float per argument here
         ; (let_floats, body2) <- {-#SCC "prepareBinding" #-}
-                                 prepareBinding env top_lvl bndr1 body1
+                                 prepareBinding body_env top_lvl bndr1 body1
         ; let body_floats2 = body_floats1 `addLetFloats` let_floats
 
         ; (rhs_floats, rhs')
@@ -608,7 +608,7 @@ tryCastWorkerWrapper env top_lvl old_bndr occ_info bndr (Cast rhs co)
   , not (hasInlineUnfolding info)  -- Not INLINE things: Wrinkle 4
   , not (isUnliftedType rhs_ty)    -- Not if rhs has an unlifted type;
                                    --     see Note [Cast w/w: unlifted]
-  = do  { (rhs_floats, work_rhs) <- prepareRhs mode top_lvl occ_fs rhs
+  = do  { (rhs_floats, work_rhs) <- prepareRhs env top_lvl occ_fs rhs
         ; uniq <- getUniqueM
         ; let work_name = mkSystemVarName uniq occ_fs
               work_id   = mkLocalIdWithInfo work_name Many rhs_ty worker_info
@@ -691,7 +691,7 @@ prepareBinding :: SimplEnv -> TopLevelFlag
                -> OutId -> OutExpr
                -> SimplM (LetFloats, OutExpr)
 prepareBinding env top_lvl bndr rhs
-  = prepareRhs (getMode env) top_lvl (getOccFS bndr) rhs
+  = prepareRhs env top_lvl (getOccFS bndr) rhs
 
 {- Note [prepareRhs]
 ~~~~~~~~~~~~~~~~~~~~
@@ -711,7 +711,7 @@ Here we want to make e1,e2 trivial and get
 That's what the 'go' loop in prepareRhs does
 -}
 
-prepareRhs :: SimplMode -> TopLevelFlag
+prepareRhs :: SimplEnv -> TopLevelFlag
            -> FastString    -- Base for any new variables
            -> OutExpr
            -> SimplM (LetFloats, OutExpr)
@@ -721,7 +721,7 @@ prepareRhs :: SimplMode -> TopLevelFlag
 -- becomes    a = e
 --            x = Just a
 -- See Note [prepareRhs]
-prepareRhs mode top_lvl occ rhs0
+prepareRhs env top_lvl occ rhs0
   = do  { (_is_exp, floats, rhs1) <- go 0 rhs0
         ; return (floats, rhs1) }
   where
@@ -736,7 +736,7 @@ prepareRhs mode top_lvl occ rhs0
         = do { (is_exp, floats1, fun') <- go (n_val_args+1) fun
              ; case is_exp of
                 False -> return (False, emptyLetFloats, App fun arg)
-                True  -> do { (floats2, arg') <- makeTrivial mode top_lvl topDmd occ arg
+                True  -> do { (floats2, arg') <- makeTrivial env top_lvl topDmd occ arg
                             ; return (True, floats1 `addLetFlts` floats2, App fun' arg') } }
     go n_val_args (Var fun)
         = return (is_exp, emptyLetFloats, Var fun)
@@ -765,58 +765,60 @@ prepareRhs mode top_lvl occ rhs0
     go _ other
         = return (False, emptyLetFloats, other)
 
-makeTrivialArg :: SimplMode -> ArgSpec -> SimplM (LetFloats, ArgSpec)
-makeTrivialArg mode arg@(ValArg { as_arg = e, as_dmd = dmd })
-  = do { (floats, e') <- makeTrivial mode NotTopLevel dmd (fsLit "arg") e
+makeTrivialArg :: SimplEnv -> ArgSpec -> SimplM (LetFloats, ArgSpec)
+makeTrivialArg env arg@(ValArg { as_arg = e, as_dmd = dmd })
+  = do { (floats, e') <- makeTrivial env NotTopLevel dmd (fsLit "arg") e
        ; return (floats, arg { as_arg = e' }) }
 makeTrivialArg _ arg
   = return (emptyLetFloats, arg)  -- CastBy, TyArg
 
-makeTrivial :: SimplMode -> TopLevelFlag -> Demand
+makeTrivial :: SimplEnv -> TopLevelFlag -> Demand
             -> FastString  -- ^ A "friendly name" to build the new binder from
             -> OutExpr     -- ^ This expression satisfies the let/app invariant
             -> SimplM (LetFloats, OutExpr)
 -- Binds the expression to a variable, if it's not trivial, returning the variable
 -- For the Demand argument, see Note [Keeping demand info in StrictArg Plan A]
-makeTrivial mode top_lvl dmd occ_fs expr
+makeTrivial env top_lvl dmd occ_fs expr
   | exprIsTrivial expr                          -- Already trivial
   || not (bindingOk top_lvl expr expr_ty)       -- Cannot trivialise
                                                 --   See Note [Cannot trivialise]
   = return (emptyLetFloats, expr)
 
   | Cast expr' co <- expr
-  = do { (floats, triv_expr) <- makeTrivial mode top_lvl dmd occ_fs expr'
+  = do { (floats, triv_expr) <- makeTrivial env top_lvl dmd occ_fs expr'
        ; return (floats, Cast triv_expr co) }
 
   | otherwise
-  = do { (floats, new_id) <- makeTrivialBinding mode top_lvl occ_fs
+  = do { (floats, new_id) <- makeTrivialBinding env top_lvl occ_fs
                                                 id_info expr expr_ty
        ; return (floats, Var new_id) }
   where
     id_info = vanillaIdInfo `setDemandInfo` dmd
     expr_ty = exprType expr
 
-makeTrivialBinding :: SimplMode -> TopLevelFlag
+makeTrivialBinding :: SimplEnv -> TopLevelFlag
                    -> FastString  -- ^ a "friendly name" to build the new binder from
                    -> IdInfo
                    -> OutExpr     -- ^ This expression satisfies the let/app invariant
                    -> OutType     -- Type of the expression
                    -> SimplM (LetFloats, OutId)
-makeTrivialBinding mode top_lvl occ_fs info expr expr_ty
-  = do  { (floats, expr1) <- prepareRhs mode top_lvl occ_fs expr
+makeTrivialBinding env top_lvl occ_fs info expr expr_ty
+  = do  { (floats, expr1) <- prepareRhs env top_lvl occ_fs expr
         ; uniq <- getUniqueM
         ; let name = mkSystemVarName uniq occ_fs
               var  = mkLocalIdWithInfo name Many expr_ty info
 
         -- Now something very like completeBind,
         -- but without the postInlineUnconditionally part
-        ; (arity_type, expr2) <- tryEtaExpandRhs mode var expr1
+        ; (arity_type, expr2) <- tryEtaExpandRhs env var expr1
         ; unf <- mkLetUnfolding (sm_uf_opts mode) top_lvl InlineRhs var expr2
 
         ; let final_id = addLetBndrInfo var arity_type unf
               bind     = NonRec final_id expr2
 
         ; return ( floats `addLetFlts` unitLetFloat bind, final_id ) }
+  where
+    mode = getMode env
 
 bindingOk :: TopLevelFlag -> CoreExpr -> Type -> Bool
 -- True iff we can have a binding of this expression at this level
@@ -900,11 +902,10 @@ completeBind env top_lvl mb_cont old_bndr new_bndr new_rhs
    do { let old_info = idInfo old_bndr
             old_unf  = realUnfoldingInfo old_info
             occ_info = occInfo old_info
-            mode     = getMode env
 
          -- Do eta-expansion on the RHS of the binding
          -- See Note [Eta-expanding at let bindings] in GHC.Core.Opt.Simplify.Utils
-      ; (new_arity, eta_rhs) <- tryEtaExpandRhs mode new_bndr new_rhs
+      ; (new_arity, eta_rhs) <- tryEtaExpandRhs env new_bndr new_rhs
 
         -- Simplify the unfolding
       ; new_unfolding <- simplLetUnfolding env top_lvl mb_cont old_bndr
@@ -1651,8 +1652,8 @@ simplLam env bndrs body (TickIt tickish cont)
         -- Not enough args, so there are real lambdas left to put in the result
 simplLam env bndrs body cont
   = do  { (env', bndrs') <- simplLamBndrs env bndrs
-        ; body' <- simplExpr env' body
-        ; new_lam <- mkLam env bndrs' body' cont
+        ; body'   <- simplExpr env' body
+        ; new_lam <- mkLam env' bndrs' body' cont
         ; rebuild env' new_lam cont }
 
 -------------
@@ -3478,7 +3479,7 @@ mkDupableContWithDmds env _
        ; (floats1, cont')  <- mkDupableContWithDmds env dmds cont
                               -- Use the demands from the function to add the right
                               -- demand info on any bindings we make for further args
-       ; (floats_s, args') <- mapAndUnzipM (makeTrivialArg (getMode env))
+       ; (floats_s, args') <- mapAndUnzipM (makeTrivialArg env)
                                            (ai_args fun)
        ; return ( foldl' addLetFloats floats1 floats_s
                 , StrictArg { sc_fun = fun { ai_args = args' }
@@ -3524,7 +3525,7 @@ mkDupableContWithDmds env dmds
         ; (floats1, cont') <- mkDupableContWithDmds env dmds cont
         ; let env' = env `setInScopeFromF` floats1
         ; (_, se', arg') <- simplArg env' dup se arg
-        ; (let_floats2, arg'') <- makeTrivial (getMode env) NotTopLevel dmd (fsLit "karg") arg'
+        ; (let_floats2, arg'') <- makeTrivial env NotTopLevel dmd (fsLit "karg") arg'
         ; let all_floats = floats1 `addLetFloats` let_floats2
         ; return ( all_floats
                  , ApplyToVal { sc_arg = arg''
@@ -4109,7 +4110,7 @@ simplStableUnfolding env top_lvl mb_cont id rhs_ty id_arity unf
     eta_expand expr
       | not eta_on         = expr
       | exprIsTrivial expr = expr
-      | otherwise          = etaExpandAT id_arity expr
+      | otherwise          = etaExpandAT (getInScope env) id_arity expr
     eta_on = sm_eta_expand (getMode env)
 
 {- Note [Eta-expand stable unfoldings]
diff --git a/compiler/GHC/Core/Opt/Simplify/Utils.hs b/compiler/GHC/Core/Opt/Simplify/Utils.hs
index 39f62d8744..5c3114e76b 100644
--- a/compiler/GHC/Core/Opt/Simplify/Utils.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Utils.hs
@@ -409,6 +409,7 @@ contIsRhsOrArg _               = False
 
 contIsRhs :: SimplCont -> Bool
 contIsRhs (Stop _ RhsCtxt) = True
+contIsRhs (CastIt _ k)     = contIsRhs k   -- For f = e |> co, treat e as Rhs context
 contIsRhs _                = False
 
 -------------------
@@ -1557,11 +1558,13 @@ mkLam :: SimplEnv -> [OutBndr] -> OutExpr -> SimplCont -> SimplM OutExpr
 -- mkLam tries three things
 --      a) eta reduction, if that gives a trivial expression
 --      b) eta expansion [only if there are some value lambdas]
-
+--
+-- NB: the SimplEnv already includes the [OutBndr] in its in-scope set
 mkLam _env [] body _cont
   = return body
 mkLam env bndrs body cont
-  = do { dflags <- getDynFlags
+  = {-#SCC "mkLam" #-}
+    do { dflags <- getDynFlags
        ; mkLam' dflags bndrs body }
   where
     mkLam' :: DynFlags -> [OutBndr] -> OutExpr -> SimplM OutExpr
@@ -1595,13 +1598,16 @@ mkLam env bndrs body cont
       , let body_arity = exprEtaExpandArity dflags body
       , expandableArityType body_arity
       = do { tick (EtaExpansion (head bndrs))
-           ; let res = mkLams bndrs (etaExpandAT body_arity body)
+           ; let res = mkLams bndrs $
+                       etaExpandAT in_scope body_arity body
            ; traceSmpl "eta expand" (vcat [text "before" <+> ppr (mkLams bndrs body)
                                           , text "after" <+> ppr res])
            ; return res }
 
       | otherwise
       = return (mkLams bndrs body)
+      where
+        in_scope = getInScope env  -- Includes 'bndrs'
 
 {-
 Note [Eta expanding lambdas]
@@ -1664,13 +1670,13 @@ because the latter is not well-kinded.
 ************************************************************************
 -}
 
-tryEtaExpandRhs :: SimplMode -> OutId -> OutExpr
+tryEtaExpandRhs :: SimplEnv -> OutId -> OutExpr
                 -> SimplM (ArityType, OutExpr)
 -- See Note [Eta-expanding at let bindings]
 -- If tryEtaExpandRhs rhs = (n, is_bot, rhs') then
 --   (a) rhs' has manifest arity n
 --   (b) if is_bot is True then rhs' applied to n args is guaranteed bottom
-tryEtaExpandRhs mode bndr rhs
+tryEtaExpandRhs env bndr rhs
   | Just join_arity <- isJoinId_maybe bndr
   = do { let (join_bndrs, join_body) = collectNBinders join_arity rhs
              oss   = [idOneShotInfo id | id <- join_bndrs, isId id]
@@ -1686,12 +1692,14 @@ tryEtaExpandRhs mode bndr rhs
   , new_arity > old_arity   -- And the current manifest arity isn't enough
   , want_eta rhs
   = do { tick (EtaExpansion bndr)
-       ; return (arity_type, etaExpandAT arity_type rhs) }
+       ; return (arity_type, etaExpandAT in_scope arity_type rhs) }
 
   | otherwise
   = return (arity_type, rhs)
 
   where
+    mode      = getMode env
+    in_scope  = getInScope env
     dflags    = sm_dflags mode
     old_arity = exprArity rhs