Fix void-arg-adding mechanism for worker/wrapper

As #22725 shows, in worker/wrapper we must add the void argument
/last/, not first.  See GHC.Core.Opt.WorkWrap.Utils
Note [Worker/wrapper needs to add void arg last].

That led me to to study GHC.Core.Opt.SpecConstr
Note [SpecConstr needs to add void args first] which suggests the
opposite!  And indeed I think it's the other way round for SpecConstr
-- or more precisely the void arg must precede the "extra_bndrs".

That led me to some refactoring of GHC.Core.Opt.SpecConstr.calcSpecInfo.

(cherry picked from commit 964284fcab6e27fe2fa5c279ea008551cbc15dbb)

4 files changed, 133 insertions, 77 deletions

diff --git a/compiler/GHC/Core/Opt/SpecConstr.hs b/compiler/GHC/Core/Opt/SpecConstr.hs
index 157cec6e49..fbdf0269e1 100644
--- a/compiler/GHC/Core/Opt/SpecConstr.hs
+++ b/compiler/GHC/Core/Opt/SpecConstr.hs
@@ -53,6 +53,7 @@ import GHC.Unit.Module.ModGuts
 import GHC.Types.Literal ( litIsLifted )
 import GHC.Types.Id
 import GHC.Types.Id.Info ( IdDetails(..) )
+import GHC.Types.Id.Make ( voidArgId, voidPrimId )
 import GHC.Types.Var.Env
 import GHC.Types.Var.Set
 import GHC.Types.Name
@@ -1924,24 +1925,13 @@ spec_one env fn arg_bndrs body (call_pat, rule_number)
 
                 -- And build the results
         ; (qvars', pats') <- generaliseDictPats qvars pats
-        ; let spec_body_ty   = exprType spec_body
-              (spec_lam_args1, spec_sig, spec_arity1, spec_join_arity1)
-                  = calcSpecInfo fn call_pat extra_bndrs
-                  -- Annotate the variables with the strictness information from
-                  -- the function (see Note [Strictness information in worker binders])
-              add_void_arg = needsVoidWorkerArg fn arg_bndrs spec_lam_args1
-              (spec_lam_args, spec_call_args, spec_arity, spec_join_arity)
-                  | add_void_arg
-                  -- See Note [SpecConstr needs to add void args first]
-                  , (spec_lam_args, spec_call_args, _) <- addVoidWorkerArg spec_lam_args1 []
-                      -- needsVoidWorkerArg: usual w/w hack to avoid generating
-                      -- a spec_rhs of unlifted type and no args.
-                  , !spec_arity      <- spec_arity1 + 1
-                  , !spec_join_arity <- fmap (+ 1) spec_join_arity1
-                  = (spec_lam_args,  spec_call_args, spec_arity,  spec_join_arity)
-                  | otherwise
-                  = (spec_lam_args1, spec_lam_args1, spec_arity1, spec_join_arity1)
+        ; let spec_body_ty = exprType spec_body
+              (spec_lam_args, spec_call_args, spec_sig)
+                  = calcSpecInfo fn arg_bndrs call_pat extra_bndrs
 
+              spec_arity = count isId spec_lam_args
+              spec_join_arity | isJoinId fn = Just (length spec_call_args)
+                              | otherwise   = Nothing
               spec_id    = asWorkerLikeId $
                            mkLocalId spec_name ManyTy
                                      (mkLamTypes spec_lam_args spec_body_ty)
@@ -1953,11 +1943,7 @@ spec_one env fn arg_bndrs body (call_pat, rule_number)
 
         -- Conditionally use result of new worker-wrapper transform
               spec_rhs = mkLams spec_lam_args (mkSeqs cbv_args spec_body_ty spec_body)
-              rule_rhs = mkVarApps (Var spec_id) $
-                              -- This will give us all the arguments we quantify over
-                              -- in the rule plus the void argument if present
-                              -- since `length(qvars) + void + length(extra_bndrs) = length spec_call_args`
-                              dropTail (length extra_bndrs) spec_call_args
+              rule_rhs = mkVarApps (Var spec_id) spec_call_args
               inline_act = idInlineActivation fn
               this_mod   = sc_module $ sc_opts env
               rule       = mkRule this_mod True {- Auto -} True {- Local -}
@@ -2020,33 +2006,55 @@ mkSeqs seqees res_ty rhs =
         = rhs
 
 
-{- Note [SpecConstr needs to add void args first]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+{- Note [SpecConstr void argument insertion]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Consider a function
+    f :: Bool -> forall t. blah
     f start @t = e
 We want to specialize for a partially applied call `f True`.
 See also Note [SpecConstr call patterns], second Wrinkle.
 Naively we would expect to get
+    $sf :: forall t. blah
     $sf @t = $se
     RULE: f True = $sf
-The specialized function only takes a single type argument
-so we add a void argument to prevent it from turning into
-a thunk. See Note [Protecting the last value argument] for details
-why. Normally we would add the void argument after the
-type argument giving us:
+The specialized function only takes a single type argument so we add a
+void argument to prevent it from turning into a thunk. See Note
+[Protecting the last value argument] for details why. Normally we
+would add the void argument after the type argument giving us:
+
     $sf :: forall t. Void# -> bla
     $sf @t void = $se
     RULE: f True = $sf void# (wrong)
-But if you look closely this wouldn't typecheck!
-If we substitute `f True` with `$sf void#` we expect the type argument to be applied first
-but we apply void# first.
-The easiest fix seems to be just to add the void argument to the front of the arguments.
-Now we get:
+
+But if you look closely this wouldn't typecheck!  If we substitute `f
+True` with `$sf void#` we expect the type argument to be applied first
+but we apply void# first.  The easiest fix seems to be just to add the
+void argument to the front of the arguments.  Now we get:
+
     $sf :: Void# -> forall t. bla
     $sf void @t = $se
     RULE: f True = $sf void#
+
 And now we can substitute `f True` with `$sf void#` with everything working out nicely!
 
+More precisely, in `calcSpecInfo`
+(i)  we need the void arg to /precede/ the `extra_bndrs`, but
+(ii) it must still /follow/ `qvar_bndrs`.
+
+Example to illustrate (ii):
+  f :: forall r (a :: TYPE r). Bool -> a
+  f = /\r. /\(a::TYPE r). \b. body
+
+  {- Specialise for f _ _ True -}
+
+  $sf :: forall r (a :: TYPE r). Void# -> a
+  $sf = /\r. /\(a::TYPE r). \v. body[True/b]
+  RULE: forall r (a :: TYPE r). f @r @a True = $sf @r @a void#
+
+The void argument must follow the foralls, lest the forall be
+ill-kinded.  See Note [Worker/wrapper needs to add void arg last] in
+GHC.Core.Opt.WorkWrap.Utils.
+
 Note [generaliseDictPats]
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 Consider these two rules (#21831, item 2):
@@ -2075,36 +2083,45 @@ And /now/ "SPEC:foo" is clearly more specific: we can instantiate the new
 "SC:foo" to match the (prefix of) "SPEC:foo".
 -}
 
-calcSpecInfo :: Id                     -- The original function
-             -> CallPat                -- Call pattern
-             -> [Var]                  -- Extra bndrs
-             -> ( [Var]                     -- Demand-decorated binders
-                , DmdSig                    -- Strictness of specialised thing
-                , Arity, Maybe JoinArity )  -- Arities of specialised thing
+calcSpecInfo :: Id           -- The original function
+             -> [InVar]      -- Lambda binders of original RHS
+             -> CallPat      -- Call pattern
+             -> [Var]        -- Extra bndrs
+             -> ( [Var]           -- Demand-decorated lambda binders
+                                  --   for RHS of specialised function
+                , [Var]           -- Args for call site
+                , DmdSig )        -- Strictness of specialised thing
 -- Calculate bits of IdInfo for the specialised function
 -- See Note [Transfer strictness]
 -- See Note [Strictness information in worker binders]
-calcSpecInfo fn (CP { cp_qvars = qvars, cp_args = pats }) extra_bndrs
-  | isJoinId fn    -- Join points have strictness and arity for LHS only
-  = ( bndrs_w_dmds
-    , mkClosedDmdSig qvar_dmds div
-    , count isId qvars
-    , Just (length qvars) )
-  | otherwise
-  = ( bndrs_w_dmds
-    , mkClosedDmdSig (qvar_dmds ++ extra_dmds) div
-    , count isId qvars + count isId extra_bndrs
-    , Nothing )
+calcSpecInfo fn arg_bndrs (CP { cp_qvars = qvars, cp_args = pats }) extra_bndrs
+  = ( spec_lam_bndrs_w_dmds
+    , spec_call_args
+    , mkClosedDmdSig [idDemandInfo b | b <- spec_lam_bndrs_w_dmds, isId b] div )
   where
     DmdSig (DmdType _ fn_dmds div) = idDmdSig fn
 
-    val_pats   = filterOut isTypeArg pats -- value args at call sites, used to determine how many demands to drop
-                                          -- from the original functions demand and for setting up dmd_env.
+    val_pats   = filterOut isTypeArg pats
+                 -- Value args at call sites, used to determine how many demands to drop
+                 -- from the original functions demand and for setting up dmd_env.
+    dmd_env    = go emptyVarEnv fn_dmds val_pats
     qvar_dmds  = [ lookupVarEnv dmd_env qv `orElse` topDmd | qv <- qvars, isId qv ]
     extra_dmds = dropList val_pats fn_dmds
 
-    bndrs_w_dmds =  set_dmds qvars       qvar_dmds
-                 ++ set_dmds extra_bndrs extra_dmds
+    -- Annotate the variables with the strictness information from
+    -- the function (see Note [Strictness information in worker binders])
+    qvars_w_dmds          = set_dmds qvars       qvar_dmds
+    extras_w_dmds         = set_dmds extra_bndrs extra_dmds
+    spec_lam_bndrs_w_dmds = final_qvars_w_dmds ++ extras_w_dmds
+
+    (final_qvars_w_dmds, spec_call_args)
+       | needsVoidWorkerArg fn arg_bndrs (qvars ++ extra_bndrs)
+         -- Usual w/w hack to avoid generating
+         -- a spec_rhs of unlifted or ill-kinded type and no args.
+         -- See Note [SpecConstr void argument insertion]
+       = ( qvars_w_dmds ++ [voidArgId], qvars ++ [voidPrimId] )
+       | otherwise
+       = ( qvars_w_dmds,                qvars )
 
     set_dmds :: [Var] -> [Demand] -> [Var]
     set_dmds [] _   = []
@@ -2112,8 +2129,6 @@ calcSpecInfo fn (CP { cp_qvars = qvars, cp_args = pats }) extra_bndrs
     set_dmds (v:vs) ds@(d:ds') | isTyVar v = v                   : set_dmds vs ds
                                | otherwise = setIdDemandInfo v d : set_dmds vs ds'
 
-    dmd_env = go emptyVarEnv fn_dmds val_pats
-
     go :: DmdEnv -> [Demand] -> [CoreExpr] -> DmdEnv
     -- We've filtered out all the type patterns already
     go env (d:ds) (pat : pats)     = go (go_one env d pat) ds pats
@@ -2127,7 +2142,6 @@ calcSpecInfo fn (CP { cp_qvars = qvars, cp_args = pats }) extra_bndrs
       = go env ds args
     go_one env _  _ = env
 
-
 {-
 Note [spec_usg includes rhs_usg]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
index 49c4aee18c..f599975355 100644
--- a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
@@ -9,7 +9,7 @@ A library for the ``worker\/wrapper'' back-end to the strictness analyser
 
 module GHC.Core.Opt.WorkWrap.Utils
    ( WwOpts(..), mkWwBodies, mkWWstr, mkWWstr_one
-   , needsVoidWorkerArg, addVoidWorkerArg
+   , needsVoidWorkerArg
    , DataConPatContext(..)
    , UnboxingDecision(..), canUnboxArg
    , findTypeShape, IsRecDataConResult(..), isRecDataCon
@@ -377,25 +377,34 @@ We use the state-token type which generates no code.
 -- Note [Preserving float barriers].
 needsVoidWorkerArg :: Id -> [Var] -> [Var] -> Bool
 needsVoidWorkerArg fn_id wrap_args work_args
-  =  not (isJoinId fn_id) && no_value_arg -- See Note [Protecting the last value argument]
-  || needs_float_barrier                  -- See Note [Preserving float barriers]
+  =  thunk_problem         -- See Note [Protecting the last value argument]
+  || needs_float_barrier   -- See Note [Preserving float barriers]
   where
-    no_value_arg        = all (not . isId) work_args
+    -- thunk_problem: see Note [Protecting the last value argument]
+    -- For join points we are only worried about (4), not (1-4).
+    -- And (4) can't happen if (null work_args)
+    --     (We could be more clever, by looking at the result type, but
+    --      this approach is simple and conservative.)
+    thunk_problem | isJoinId fn_id = no_value_arg && not (null work_args)
+                  | otherwise      = no_value_arg
+    no_value_arg = not (any isId work_args)
+
+    -- needs_float_barrier: see Note [Preserving float barriers]
+    needs_float_barrier = wrap_had_barrier && not work_has_barrier
     is_float_barrier v  = isId v && hasNoOneShotInfo (idOneShotInfo v)
     wrap_had_barrier    = any is_float_barrier wrap_args
     work_has_barrier    = any is_float_barrier work_args
-    needs_float_barrier = wrap_had_barrier && not work_has_barrier
 
--- | Inserts a `Void#` arg before the first argument.
---
--- Why as the first argument? See Note [SpecConstr needs to add void args first]
--- in SpecConstr.
+-- | Inserts a `Void#` arg as the last argument.
+-- Why last? See Note [Worker/wrapper needs to add void arg last]
 addVoidWorkerArg :: [Var] -> [StrictnessMark]
-                 -> ([Var],     -- Lambda bound args
-                     [Var],     -- Args at call site
-                     [StrictnessMark]) -- str semantics for the worker args.
+                 -> ( [Var]     -- Lambda bound args
+                    , [Var]     -- Args at call site
+                    , [StrictnessMark]) -- str semantics for the worker args
 addVoidWorkerArg work_args str_marks
-  = (voidArgId : work_args, voidPrimId:work_args, NotMarkedStrict:str_marks)
+  = ( work_args ++ [voidArgId]
+    , work_args ++ [voidPrimId]
+    , str_marks ++ [NotMarkedStrict] )
 
 {-
 Note [Protecting the last value argument]
@@ -403,8 +412,8 @@ Note [Protecting the last value argument]
 If the user writes (\_ -> E), they might be intentionally disallowing
 the sharing of E. Since absence analysis and worker-wrapper are keen
 to remove such unused arguments, we add in a void argument to prevent
-the function from becoming a thunk.  Three reasons why turning a function
-into a thunk might be bad:
+the function from becoming a thunk.  Here are several reasons why turning
+a function into a thunk might be bad:
 
 1) It can create a space leak. e.g.
       f x = let y () = [1..x]
@@ -423,7 +432,19 @@ into a thunk might be bad:
        g = \x. 30#
    Removing the \x would leave an unlifted binding.
 
-NB: none of these apply to a join point.
+4) It can create a worker of ill-kinded type (#22275).  Consider
+     f :: forall r (a :: TYPE r). () -> a
+     f x = f x
+   Here `x` is absent, but if we simply drop it we'd end up with
+     $wf :: forall r (a :: TYPE r). a
+   But alas $wf's type is ill-kinded: the kind of (/\r (a::TYPE r).a)
+   is (TYPE r), which mentions the bound variable `r`.  See also
+   Note [Worker/wrapper needs to add void arg last]
+
+See also Note [Preserving float barriers]
+
+NB: Of these, only (1-3) don't apply to a join point, which can be
+unlifted even if the RHS is not ok-for-speculation.
 
 Note [Preserving float barriers]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -457,7 +478,7 @@ which some are absent or one-shot and the resulting worker arguments:
 
   * \a{Abs}.\b{os}.\c{os}... ==> \b{os}.\c{os}.\(_::Void#)...
     Wrapper arg `a` was the only float barrier and had been dropped. Hence Void#
-  * \a{Abs,os}.\b{os}.\c... ==> \b{os}.\c...
+p  * \a{Abs,os}.\b{os}.\c... ==> \b{os}.\c...
     Worker arg `c` is a float barrier.
   * \a.\b{Abs}.\c{os}... ==> \a.\c{os}...
     Worker arg `a` is a float barrier.
@@ -469,6 +490,27 @@ which some are absent or one-shot and the resulting worker arguments:
 
 Executable examples in T21150.
 
+Note [Worker/wrapper needs to add void arg last]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider point (4) of Note [Protecting the last value argument]
+
+  f :: forall r (a :: TYPE r). () -> a
+  f x = f x
+
+As pointed out in (4) we need to add a void argument.  But if we add
+it /first/ we'd get
+
+  $wf :: Void# -> forall r (a :: TYPE r). a
+  $wf = ...
+
+But alas $wf's type is /still/ still-kinded, just as before in (4).
+Solution is simple: put the void argument /last/:
+
+  $wf :: forall r (a :: TYPE r). Void# -> a
+  $wf = ...
+
+c.f Note [SpecConstr void argument insertion] in GHC.Core.Opt.SpecConstr
+
 Note [Join points and beta-redexes]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Originally, the worker would invoke the original function by calling it with
diff --git a/compiler/GHC/Core/Type.hs b/compiler/GHC/Core/Type.hs
index 76326b6c50..9eb9d86f6e 100644
--- a/compiler/GHC/Core/Type.hs
+++ b/compiler/GHC/Core/Type.hs
@@ -1337,7 +1337,7 @@ ty_con_app_fun_maybe many_ty_co tc args
       | otherwise
       = Nothing
 
-mkFunctionType :: Mult -> Type -> Type -> Type
+mkFunctionType :: HasDebugCallStack => Mult -> Type -> Type -> Type
 -- ^ This one works out the FunTyFlag from the argument type
 -- See GHC.Types.Var Note [FunTyFlag]
 mkFunctionType mult arg_ty res_ty
diff --git a/compiler/GHC/Core/Utils.hs b/compiler/GHC/Core/Utils.hs
index c88ddb3d55..889bfbf5b4 100644
--- a/compiler/GHC/Core/Utils.hs
+++ b/compiler/GHC/Core/Utils.hs
@@ -159,7 +159,7 @@ coreAltsType :: [CoreAlt] -> Type
 coreAltsType (alt:_) = coreAltType alt
 coreAltsType []      = panic "coreAltsType"
 
-mkLamType  :: Var -> Type -> Type
+mkLamType  :: HasDebugCallStack => Var -> Type -> Type
 -- ^ Makes a @(->)@ type or an implicit forall type, depending
 -- on whether it is given a type variable or a term variable.
 -- This is used, for example, when producing the type of a lambda.