CorePrep: Handle over-saturated primitiveswip/T22937

Fixes #22937. See the new wrinkle (W2) in
Note [Calling primitives with the right arity].

6 files changed, 121 insertions, 82 deletions

diff --git a/compiler/GHC/Builtin/PrimOps.hs b/compiler/GHC/Builtin/PrimOps.hs
index c9f0d56aaf..19a3b09fbc 100644
--- a/compiler/GHC/Builtin/PrimOps.hs
+++ b/compiler/GHC/Builtin/PrimOps.hs
@@ -700,7 +700,7 @@ convention for curried applications that can accommodate representation
 polymorphism.
 
 To ensure saturation, CorePrep eta expands all primop applications as
-described in Note [Eta expansion of hasNoBinding things in CorePrep] in
+described in Note [Calling primitives with the right arity] in
 GHC.Core.Prep.
 
 Historical Note:
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index c0b72cefed..927fc07872 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -6,7 +6,7 @@
 (c) The University of Glasgow, 1994-2006
 
 
-Core pass to saturate constructors and PrimOps
+Core pass to ANF-ise and saturate PrimOps and cbv-functions
 -}
 
 module GHC.CoreToStg.Prep
@@ -80,7 +80,7 @@ Note [CorePrep Overview]
 
 The goal of this pass is to prepare for code generation.
 
-1.  Saturate constructor and primop applications.
+1.  Saturate applications of primops and cbv functions.
 
 2.  Convert to A-normal form; that is, function arguments
     are always variables.
@@ -1101,12 +1101,14 @@ cpeApp top_env expr
       = do { v1 <- fiddleCCall v
            ; let e2 = lookupCorePrepEnv env v1
                  hd = getIdFromTrivialExpr_maybe e2
-                 -- Determine number of required arguments. See Note [Ticks and mandatory eta expansion]
-                 min_arity = case hd of
+                 -- Determine the number of required arguments.
+                 -- See Note [Calling primitives with the right arity]
+                 -- and Note [Ticks and mandatory eta expansion]
+                 exact_arity = case hd of
                    Just v_hd -> if hasNoBinding v_hd then Just $! (idArity v_hd) else Nothing
                    Nothing -> Nothing
           --  ; pprTraceM "cpe_app:stricts:" (ppr v <+> ppr args $$ ppr stricts $$ ppr (idCbvMarks_maybe v))
-           ; (app, floats, unsat_ticks) <- rebuild_app env args e2 emptyFloats stricts min_arity
+           ; (app, floats, unsat_ticks) <- rebuild_app env args e2 emptyFloats stricts exact_arity
            ; mb_saturate hd app floats unsat_ticks depth }
         where
           depth = val_args args
@@ -1134,9 +1136,12 @@ cpeApp top_env expr
                           -- If evalDmd says that it's sure to be evaluated,
                           -- we'll end up case-binding it
            ; (app, floats,unsat_ticks) <- rebuild_app env args fun' fun_floats [] Nothing
-           ; mb_saturate Nothing app floats unsat_ticks (val_args args) }
+           ; massert (null unsat_ticks)
+           ; return (floats, app) }
 
-    -- Count the number of value arguments *and* coercions (since we don't eliminate the later in STG)
+
+    -- Count the number of value arguments *including* coercions
+    -- (since we don't eliminate the latter in STG)
     val_args :: [ArgInfo] -> Int
     val_args args = go args 0
       where
@@ -1174,13 +1179,13 @@ cpeApp top_env expr
         -> CpeApp                     -- The function
         -> Floats
         -> [Demand]
-        -> Maybe Arity
+        -> Maybe Arity                -- (Just arity) when headed by a hasNoBinding Id
         -> UniqSM (CpeApp
                   ,Floats
                   ,[CoreTickish] -- Underscoped ticks. See Note [Ticks and mandatory eta expansion]
                   )
-    rebuild_app env args app floats ss req_depth =
-      rebuild_app' env args app floats ss [] (fromMaybe 0 req_depth)
+    rebuild_app env args app floats ss req_depth  =
+      rebuild_app' env args app floats ss [] (fromMaybe (0-1) req_depth)
 
     rebuild_app'
         :: CorePrepEnv
@@ -1189,33 +1194,37 @@ cpeApp top_env expr
         -> Floats
         -> [Demand]
         -> [CoreTickish]
-        -> Int -- Number of arguments required to satisfy minimal tick scopes.
+        -> Int -- Negative for normal functions;
+               -- number of remaining value arguments for hasNoBinding Ids;
+               -- see Note [Calling primitives with the right arity]
+               -- and Note [Ticks and mandatory eta expansion]
         -> UniqSM (CpeApp, Floats, [CoreTickish])
     rebuild_app' _ [] app floats ss rt_ticks !_req_depth
       = assertPpr (null ss) (ppr ss)-- make sure we used all the strictness info
         return (app, floats, rt_ticks)
 
     rebuild_app' env (a : as) fun' floats ss rt_ticks req_depth = case a of
-      -- See Note [Ticks and mandatory eta expansion]
-      _
-        | not (null rt_ticks)
-        , req_depth <= 0
-        ->
-            let tick_fun = foldr mkTick fun' rt_ticks
-            in rebuild_app' env (a : as) tick_fun floats ss rt_ticks req_depth
-
       CpeApp (Type arg_ty)
         -> rebuild_app' env as (App fun' (Type arg_ty')) floats ss rt_ticks req_depth
         where
           arg_ty' = cpSubstTy env arg_ty
 
-      CpeApp (Coercion co)
-        -> rebuild_app' env as (App fun' (Coercion co')) floats (drop 1 ss) rt_ticks req_depth
-        where
-            co' = cpSubstCo env co
-
-      CpeApp arg -> do
-        let (ss1, ss_rest)  -- See Note [lazyId magic] in GHC.Types.Id.Make
+      CpeApp arg
+        | req_depth == 0 -> do
+            -- See Note [Calling primitives with the right arity], wrinkle W2:
+            -- The primitive already has the right number of value arguments
+            -- we must case-bind before we can apply it to another argument.
+            -- We also apply any collected profiling ticks now; see
+            -- Note [Ticks and mandatory eta expansion]
+            v <- newVar (exprType fun')
+            let tick_fun = foldr mkTick fun' rt_ticks
+                float = mkFloat env evalDmd False v tick_fun
+            rebuild_app' env (a : as) (Var v) (addFloat floats float) ss [] (0-1)
+        | Coercion co <- arg
+        , let co' = cpSubstCo env co
+        -> rebuild_app' env as (App fun' (Coercion co')) floats (drop 1 ss) rt_ticks (req_depth-1)
+        | otherwise -> do
+        let (ss1, ss_rest)  -- See Note [lazyId magic] in GHC.Types.Id.Make, wrinkle W3
                = case (ss, isLazyExpr arg) of
                    (_   : ss_rest, True)  -> (topDmd, ss_rest)
                    (ss1 : ss_rest, False) -> (ss1,    ss_rest)
@@ -1227,10 +1236,11 @@ cpeApp top_env expr
         -> rebuild_app' env as (Cast fun' co') floats ss rt_ticks req_depth
         where
            co' = cpSubstCo env co
+
       -- See Note [Ticks and mandatory eta expansion]
       CpeTick tickish
         | tickishPlace tickish == PlaceRuntime
-        , req_depth > 0
+        , req_depth >= 0
         -> assert (isProfTick tickish) $
            rebuild_app' env as fun' floats ss (tickish:rt_ticks) req_depth
         | otherwise
@@ -1238,7 +1248,7 @@ cpeApp top_env expr
         -> rebuild_app' env as fun' (addFloat floats (FloatTick tickish)) ss rt_ticks req_depth
 
 isLazyExpr :: CoreExpr -> Bool
--- See Note [lazyId magic] in GHC.Types.Id.Make
+-- See Note [lazyId magic] in GHC.Types.Id.Make, wrinkle W3
 isLazyExpr (Cast e _)              = isLazyExpr e
 isLazyExpr (Tick _ e)              = isLazyExpr e
 isLazyExpr (Var f `App` _ `App` _) = f `hasKey` lazyIdKey
@@ -1445,6 +1455,11 @@ the continuation may not be a manifest lambda.
 Note [ANF-ising literal string arguments]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
+*** Is this Note still necessary? Yes, the example transformation to
+       foo = u\ [] case "turtle"# of s { __DEFAULT__ -> Foo s }
+    seems pretty bad.  But these days, we'd expect the simplifier to
+    have floated "turtle"# to top-level anyway.  Right?
+
 Consider a program like,
 
     data Foo = Foo Addr#
@@ -1516,18 +1531,42 @@ because that has different strictness.  Hence the use of 'allLazy'.
 -- Building the saturated syntax
 -- ---------------------------------------------------------------------------
 
-Note [Eta expansion of hasNoBinding things in CorePrep]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-maybeSaturate deals with eta expanding to saturate things that can't deal with
-unsaturated applications (identified by 'hasNoBinding', currently
-foreign calls, unboxed tuple/sum constructors, and representation-polymorphic
-primitives such as 'coerce' and 'unsafeCoerce#').
+Note [Calling primitives with the right arity]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+For several low-level things, the code generator can only handle
+saturated applications, i.e. applications with exactly the right
+number of arguments.  (These things are identified by 'hasNoBinding'.
+Currently, they are: foreign calls, unboxed tuple/sum constructors,
+and representation-polymorphic primitives such as 'coerce' and
+'unsafeCoerce#'.)
+
+W1: If an application has too few arguments, we must eta-expand. For
+    example, we transform `(+#) x` into `\y -> (+#) x y`.  This happens
+    in maybeSaturate.
+
+W2: Perhaps surprisingly, an application of a primitive can have too
+    many arguments!  This can make sense if the primitive returns a
+    function.  Here's an example, from #22937:
+
+      let  arg3 = \s' f -> unIO f s'
+           arg4 = putStrLn "test"
+       in  keepAlive# () s arg3 arg4
+
+    keepAlive# is a primop with arity 3, so we must apply it to its
+    first 3 arguments, and then apply the resulting function to the
+    remaining argument, as follows:
+
+      let  arg3 = \s' f -> unIO f s'
+           arg4 = putStrLn "test"
+       in  case  keepAlive# () s arg3  of  fun {
+             __DEFAULT -> fun arg4
+           };
+
+    We perform this transformation in rebuild_app.
+
 
 Historical Note: Note that eta expansion in CorePrep used to be very fragile
 due to the "prediction" of CAFfyness that we used to make during tidying.
-We previously saturated primop
-applications here as well but due to this fragility (see #16846) we now deal
-with this another way, as described in Note [Primop wrappers] in GHC.Builtin.PrimOps.
 -}
 
 maybeSaturate :: Id -> CpeApp -> Int -> [CoreTickish] -> UniqSM CpeRhs
@@ -1769,19 +1808,6 @@ mkFloat env dmd is_unlifted bndr rhs
     -- Otherwise we get  case (\x -> e) of ...!
 
   | is_unlifted = FloatCase rhs bndr DEFAULT [] True
-      -- we used to assertPpr ok_for_spec (ppr rhs) here, but it is now disabled
-      -- because exprOkForSpeculation isn't stable under ANF-ing. See for
-      -- example #19489 where the following unlifted expression:
-      --
-      --    GHC.Prim.(#|_#) @LiftedRep @LiftedRep @[a_ax0] @[a_ax0]
-      --                    (GHC.Types.: @a_ax0 a2_agq a3_agl)
-      --
-      -- is ok-for-spec but is ANF-ised into:
-      --
-      --    let sat = GHC.Types.: @a_ax0 a2_agq a3_agl
-      --    in GHC.Prim.(#|_#) @LiftedRep @LiftedRep @[a_ax0] @[a_ax0] sat
-      --
-      -- which isn't ok-for-spec because of the let-expression.
 
   | is_hnf      = FloatLet (NonRec bndr                       rhs)
   | otherwise   = FloatLet (NonRec (setIdDemandInfo bndr dmd) rhs)
diff --git a/compiler/GHC/Types/Id/Make.hs b/compiler/GHC/Types/Id/Make.hs
index bf579c0d36..0fd057553b 100644
--- a/compiler/GHC/Types/Id/Make.hs
+++ b/compiler/GHC/Types/Id/Make.hs
@@ -1904,43 +1904,43 @@ Note [lazyId magic]
 lazy :: forall a. a -> a
 
 'lazy' is used to make sure that a sub-expression, and its free variables,
-are truly used call-by-need, with no code motion.  Key examples:
+are truly used call-by-need, with no code motion.  Key example:
 
 * pseq:    pseq a b = a `seq` lazy b
   We want to make sure that the free vars of 'b' are not evaluated
   before 'a', even though the expression is plainly strict in 'b'.
-
-* catch:   catch a b = catch# (lazy a) b
-  Again, it's clear that 'a' will be evaluated strictly (and indeed
-  applied to a state token) but we want to make sure that any exceptions
-  arising from the evaluation of 'a' are caught by the catch (see
-  #11555).
+    *** This isn't especially robust. See #23233.
 
 Implementing 'lazy' is a bit tricky:
 
-* It must not have a strictness signature: by being a built-in Id,
-  all the info about lazyId comes from here, not from GHC.Magic.hi.
-  This is important, because the strictness analyser will spot it as
-  strict!
-
-* It must not have an unfolding: it gets "inlined" by a HACK in
-  CorePrep. It's very important to do this inlining *after* unfoldings
-  are exposed in the interface file.  Otherwise, the unfolding for
-  (say) pseq in the interface file will not mention 'lazy', so if we
-  inline 'pseq' we'll totally miss the very thing that 'lazy' was
-  there for in the first place. See #3259 for a real world
-  example.
-
-* Suppose CorePrep sees (catch# (lazy e) b).  At all costs we must
-  avoid using call by value here:
-     case e of r -> catch# r b
-  Avoiding that is the whole point of 'lazy'.  So in CorePrep (which
-  generate the 'case' expression for a call-by-value call) we must
-  spot the 'lazy' on the arg (in CorePrep.cpeApp), and build a 'let'
-  instead.
-
-* lazyId is defined in GHC.Base, so we don't *have* to inline it.  If it
-  appears un-applied, we'll end up just calling it.
+W1: It must not have a strictness signature: by being a built-in Id,
+    all the info about lazyId comes from here, not from GHC.Magic.hi.
+    This is important, because the strictness analyser will spot it as
+    strict!
+
+W2: It must not have an unfolding: it gets "inlined" by a HACK in
+    CorePrep. It's very important to do this inlining *after* unfoldings
+    are exposed in the interface file.  Otherwise, the unfolding for
+    (say) pseq in the interface file will not mention 'lazy', so if we
+    inline 'pseq' we'll totally miss the very thing that 'lazy' was
+    there for in the first place. See #3259 for a real world
+    example.
+
+W3: Suppose CorePrep sees (catch# (lazy e) b).  At all costs we must
+    avoid using call by value here:
+       case e of r -> catch# r b
+    Avoiding that is the whole point of 'lazy'.  So in CorePrep (which
+    generate the 'case' expression for a call-by-value call) we must
+    spot the 'lazy' on the arg (in CorePrep.cpeApp), and build a 'let'
+    instead.
+      *** We wouldn't use call-by-value in this example anyway, since
+          catch# is no longer considered strict.  (See primops.txt.pp
+          Note [Strictness for mask/unmask/catch]) This property of
+          lazy is news to me (clyring, Apr 2023). Is it documented
+          anywhere else? Is there any reason to keep it?
+
+W4: lazyId is defined in GHC.Magic, so we don't *have* to inline it.  If it
+    appears un-applied, we'll end up just calling it.
 
 Note [noinlineId magic]
 ~~~~~~~~~~~~~~~~~~~~~~~
@@ -1982,7 +1982,7 @@ Wrinkles
 (W1) Sometimes case (2) above needs to apply `noinline` to a type of kind
      Constraint; e.g.
                     noinline @(Eq Int) $dfEqInt
-     We don't have type-or-kind polymorphism, so we simply have two `inline`
+     We don't have type-or-kind polymorphism, so we simply have two `noinline`
      Ids, namely `noinlineId` and `noinlineConstraintId`.
 
 (W2) Note that noinline as currently implemented can hide some simplifications
diff --git a/testsuite/tests/core-to-stg/T22937.hs b/testsuite/tests/core-to-stg/T22937.hs
new file mode 100644
index 0000000000..fe9aa62746
--- /dev/null
+++ b/testsuite/tests/core-to-stg/T22937.hs
@@ -0,0 +1,11 @@
+{-# LANGUAGE MagicHash #-}
+
+import GHC.Exts
+import GHC.IO
+import System.IO
+
+main :: IO ()
+main = do
+  IO $ \s -> keepAlive# () s (\s' f -> unIO f s')
+                             (putStrLn "This should get printed.")
+  hFlush stdout
diff --git a/testsuite/tests/core-to-stg/T22937.stdout b/testsuite/tests/core-to-stg/T22937.stdout
new file mode 100644
index 0000000000..2474f121fc
--- /dev/null
+++ b/testsuite/tests/core-to-stg/T22937.stdout
@@ -0,0 +1 @@
+This should get printed.
diff --git a/testsuite/tests/core-to-stg/all.T b/testsuite/tests/core-to-stg/all.T
index baab982cb4..6cdc87b83e 100644
--- a/testsuite/tests/core-to-stg/all.T
+++ b/testsuite/tests/core-to-stg/all.T
@@ -1,3 +1,4 @@
 # Tests for CorePrep and CoreToStg
 
 test('T19700', normal, compile, ['-O'])
+test('T22937', normal, compile_and_run, [''])