Account for RULES that destroy ok-for-speculationwip/T18677

This patch addresses #18677.

I'll write a proper commit message in due course.

5 files changed, 64 insertions, 29 deletions

diff --git a/compiler/GHC/Core/Make.hs b/compiler/GHC/Core/Make.hs
index c242c776e6..8fc840fdec 100644
--- a/compiler/GHC/Core/Make.hs
+++ b/compiler/GHC/Core/Make.hs
@@ -6,7 +6,7 @@
 module GHC.Core.Make (
         -- * Constructing normal syntax
         mkCoreLet, mkCoreLets,
-        mkCoreApp, mkCoreApps, mkCoreConApps,
+        mkCoreApp, mkCoreApps, mkCoreAppTyped, mkCoreConApps,
         mkCoreLams, mkWildCase, mkIfThenElse,
         mkWildValBinder, mkWildEvBinder,
         mkSingleAltCase,
@@ -139,19 +139,24 @@ mkCoreConApps con args = mkCoreApps (Var (dataConWorkId con)) args
 --   See Note [Core let/app invariant] in "GHC.Core"
 mkCoreApps :: CoreExpr -> [CoreExpr] -> CoreExpr
 mkCoreApps fun args
-  = fst $
-    foldl' (mkCoreAppTyped doc_string) (fun, fun_ty) args
+  = fst $ foldl' mk_core_app (fun, fun_ty) args
   where
-    doc_string = ppr fun_ty $$ ppr fun $$ ppr args
     fun_ty = exprType fun
 
 -- | Construct an expression which represents the application of one expression
 -- to the other
 -- Respects the let/app invariant by building a case expression where necessary
 --   See Note [Core let/app invariant] in "GHC.Core"
-mkCoreApp :: SDoc -> CoreExpr -> CoreExpr -> CoreExpr
-mkCoreApp s fun arg
-  = fst $ mkCoreAppTyped s (fun, exprType fun) arg
+mkCoreApp :: HasDebugCallStack => CoreExpr -> CoreExpr -> CoreExpr
+mkCoreApp fun arg = mkCoreAppTyped fun (exprType fun) arg
+
+-- | Construct an expression which represents the application of one expression
+-- to the other.
+-- Precondition:  fun :: fun_ty
+-- Respects the let/app invariant by building a case expression where necessary
+--   See Note [Core let/app invariant] in "GHC.Core"
+mkCoreAppTyped :: HasDebugCallStack => CoreExpr -> Type -> CoreExpr -> CoreExpr
+mkCoreAppTyped fun fun_ty arg = fst $ mk_core_app (fun, fun_ty) arg
 
 -- | Construct an expression which represents the application of one expression
 -- paired with its type to an argument. The result is paired with its type. This
@@ -159,23 +164,24 @@ mkCoreApp s fun arg
 -- 'mkCoreApps'.
 -- Respects the let/app invariant by building a case expression where necessary
 --   See Note [Core let/app invariant] in "GHC.Core"
-mkCoreAppTyped :: SDoc -> (CoreExpr, Type) -> CoreExpr -> (CoreExpr, Type)
-mkCoreAppTyped _ (fun, fun_ty) (Type ty)
+mk_core_app :: HasDebugCallStack => (CoreExpr, Type) -> CoreExpr
+                                 -> (CoreExpr, Type)
+mk_core_app (fun, fun_ty) (Type ty)
   = (App fun (Type ty), piResultTy fun_ty ty)
-mkCoreAppTyped _ (fun, fun_ty) (Coercion co)
+mk_core_app (fun, fun_ty) (Coercion co)
   = (App fun (Coercion co), funResultTy fun_ty)
-mkCoreAppTyped d (fun, fun_ty) arg
-  = ASSERT2( isFunTy fun_ty, ppr fun $$ ppr arg $$ d )
-    (mkValApp fun arg (Scaled mult arg_ty) res_ty, res_ty)
+mk_core_app (fun, fun_ty) arg
+  = ASSERT2( isFunTy fun_ty, ppr fun $$ ppr arg )
+    (mkValApp fun arg arg_mult arg_ty res_ty, res_ty)
   where
-    (mult, arg_ty, res_ty) = splitFunTy fun_ty
+    (arg_mult, arg_ty, res_ty) = splitFunTy fun_ty
 
-mkValApp :: CoreExpr -> CoreExpr -> Scaled Type -> Type -> CoreExpr
+mkValApp :: CoreExpr -> CoreExpr -> Mult -> Type -> Type -> CoreExpr
 -- Build an application (e1 e2),
 -- or a strict binding  (case e2 of x -> e1 x)
 -- using the latter when necessary to respect the let/app invariant
 --   See Note [Core let/app invariant] in GHC.Core
-mkValApp fun arg (Scaled w arg_ty) res_ty
+mkValApp fun arg w arg_ty res_ty
   | not (needsCaseBinding arg_ty arg)
   = App fun arg                -- The vastly common case
   | otherwise
diff --git a/compiler/GHC/Core/Opt/Simplify.hs b/compiler/GHC/Core/Opt/Simplify.hs
index e219a0dba9..23e4063786 100644
--- a/compiler/GHC/Core/Opt/Simplify.hs
+++ b/compiler/GHC/Core/Opt/Simplify.hs
@@ -27,7 +27,7 @@ import GHC.Core.FamInstEnv ( FamInstEnv )
 import GHC.Types.Literal   ( litIsLifted ) --, mkLitInt ) -- temporalily commented out. See #8326
 import GHC.Types.Id
 import GHC.Types.Id.Make   ( seqId )
-import GHC.Core.Make       ( FloatBind, mkImpossibleExpr, castBottomExpr )
+import GHC.Core.Make       ( FloatBind, mkImpossibleExpr, castBottomExpr, mkCoreAppTyped )
 import qualified GHC.Core.Make
 import GHC.Types.Id.Info
 import GHC.Types.Name           ( mkSystemVarName, isExternalName, getOccFS )
@@ -1335,13 +1335,21 @@ rebuild env expr cont
               ; (floats2, expr') <- simplLam env' bs body cont
               ; return (floats1 `addFloats` floats2, expr') }
 
+      -- These next cases two don't happen much, because a call with
+      -- a variable at the head (f e1 ... en) is handled via rebuildCall,
+      -- which constructs ArgInfo, and with the final result being built
+      -- by argInfoExpr.  We only get here for non-variable heads, like
+      --     (case blah of alts) e1 e2
       ApplyToTy  { sc_arg_ty = ty, sc_cont = cont}
         -> rebuild env (App expr (Type ty)) cont
 
-      ApplyToVal { sc_arg = arg, sc_env = se, sc_dup = dup_flag, sc_cont = cont}
+      ApplyToVal { sc_arg = arg, sc_env = se, sc_dup = dup_flag
+                 , sc_cont = cont, sc_hole_ty = fun_ty }
         -- See Note [Avoid redundant simplification]
         -> do { (_, _, arg') <- simplArg env dup_flag se arg
-              ; rebuild env (App expr arg') cont }
+              ; rebuild env (mkCoreAppTyped expr fun_ty arg') cont }
+                 -- mkCoreAppTyped: see Note [RULEs can break let/app]
+                 -- in GHC.Core.Opt.Simplify.Env
 
 {-
 ************************************************************************
diff --git a/compiler/GHC/Core/Opt/Simplify/Env.hs b/compiler/GHC/Core/Opt/Simplify/Env.hs
index 4ceaf637ed..304c3e5b83 100644
--- a/compiler/GHC/Core/Opt/Simplify/Env.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Env.hs
@@ -615,9 +615,9 @@ mkRecFloats floats@(SimplFloats { sfLetFloats  = LetFloats bs ff
 wrapFloats :: SimplFloats -> OutExpr -> OutExpr
 -- Wrap the floats around the expression; they should all
 -- satisfy the let/app invariant, so mkLets should do the job just fine
-wrapFloats (SimplFloats { sfLetFloats  = LetFloats bs _
+wrapFloats (SimplFloats { sfLetFloats  = lbs
                         , sfJoinFloats = jbs }) body
-  = foldrOL Let (wrapJoinFloats jbs body) bs
+  = wrapLetFloats lbs $ wrapJoinFloats jbs body
      -- Note: Always safe to put the joins on the inside
      -- since the values can't refer to them
 
@@ -640,6 +640,20 @@ getTopFloatBinds (SimplFloats { sfLetFloats  = lbs
   = ASSERT( isNilOL jbs )  -- Can't be any top-level join bindings
     letFloatBinds lbs
 
+wrapLetFloats :: LetFloats -> OutExpr -> OutExpr
+wrapLetFloats (LetFloats bs _) body
+  = foldr wrap_bind body bs
+  where
+    wrap_bind bind body
+      | -- Horrid special case for a binding that doesn't satisfy
+        -- the let/app invariant; see Note [RULEs can break let/app]
+        NonRec bndr rhs <- bind
+      , isUnliftedType (idType bndr)
+      , not (exprOkForSpeculation rhs)
+      = Case rhs bndr (exprType rhs) [(DEFAULT, [], body)]
+      | otherwise
+      = Let bind body
+
 mapLetFloats :: LetFloats -> ((Id,CoreExpr) -> (Id,CoreExpr)) -> LetFloats
 mapLetFloats (LetFloats fs ff) fun
    = LetFloats (mapOL app fs) ff
@@ -647,8 +661,11 @@ mapLetFloats (LetFloats fs ff) fun
     app (NonRec b e) = case fun (b,e) of (b',e') -> NonRec b' e'
     app (Rec bs)     = Rec (map fun bs)
 
-{-
-************************************************************************
+{- Note [RULEs can break let/app]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+-}
+
+{- *********************************************************************
 *                                                                      *
                 Substitution of Vars
 *                                                                      *
diff --git a/compiler/GHC/Core/Opt/Simplify/Utils.hs b/compiler/GHC/Core/Opt/Simplify/Utils.hs
index 420d406eed..df03ffc4fd 100644
--- a/compiler/GHC/Core/Opt/Simplify/Utils.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Utils.hs
@@ -51,6 +51,7 @@ import GHC.Core.Ppr
 import GHC.Core.TyCo.Ppr ( pprParendType )
 import GHC.Core.FVs
 import GHC.Core.Utils
+import GHC.Core.Make     ( mkCoreAppTyped )
 import GHC.Core.Opt.Arity
 import GHC.Core.Unfold
 import GHC.Core.Unfold.Make
@@ -124,7 +125,7 @@ data SimplCont
 
   | ApplyToVal         -- (ApplyToVal arg K)[e] = K[ e arg ]
       { sc_dup     :: DupFlag   -- See Note [DupFlag invariants]
-      , sc_hole_ty :: OutType   -- Type of the function, presumably (forall a. blah)
+      , sc_hole_ty :: OutType   -- Type of the function, presumably (t1 -> t2)
                                 -- See Note [The hole type in ApplyToTy/Val]
       , sc_arg  :: InExpr       -- The argument,
       , sc_env  :: StaticEnv    -- see Note [StaticEnv invariant]
@@ -358,10 +359,13 @@ argInfoExpr :: OutId -> [ArgSpec] -> OutExpr
 argInfoExpr fun rev_args
   = go rev_args
   where
-    go []                              = Var fun
-    go (ValArg { as_arg = arg }  : as) = go as `App` arg
-    go (TyArg { as_arg_ty = ty } : as) = go as `App` Type ty
-    go (CastBy co                : as) = mkCast (go as) co
+    go []                                     = Var fun
+    go (ValArg { as_arg = arg
+               , as_hole_ty = fun_ty } : as) = mkCoreAppTyped (go as) fun_ty arg
+       -- mkCoreAppTyped: see Note [RULEs can break let/app]
+       -- in GHC.Core.Opt.Simplify.Env
+    go (TyArg { as_arg_ty = ty }       : as) = go as `App` Type ty
+    go (CastBy co                      : as) = mkCast (go as) co
 
 
 type FunRules = Maybe (Int, [CoreRule]) -- Remaining rules for this function
diff --git a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
index 8cc0eaa503..d74ba30895 100644
--- a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
@@ -469,7 +469,7 @@ mkWWargs subst fun_ty demands
     apply_or_bind_then k arg (Lam bndr body)
       = mkCoreLet (NonRec bndr arg) (k body)    -- Important that arg is fresh!
     apply_or_bind_then k arg fun
-      = k $ mkCoreApp (text "mkWWargs") fun arg
+      = k $ mkCoreApp fun arg
 applyToVars :: [Var] -> CoreExpr -> CoreExpr
 applyToVars vars fn = mkVarApps fn vars