SpecConstr: Properly create rules for call patterns representing partial applicationswip/andreask/spec-constr-fix

The main fix is that in addVoidWorkerArg we now add the argument to the front.

This fixes #21448.

-------------------------
Metric Decrease:
    T16875
-------------------------

3 files changed, 74 insertions, 28 deletions

diff --git a/compiler/GHC/Core/Opt/SpecConstr.hs b/compiler/GHC/Core/Opt/SpecConstr.hs
index d9429053a0..c4517c1c52 100644
--- a/compiler/GHC/Core/Opt/SpecConstr.hs
+++ b/compiler/GHC/Core/Opt/SpecConstr.hs
@@ -1761,14 +1761,6 @@ spec_one env fn arg_bndrs body (call_pat, rule_number)
               -- changes (#4012).
               rule_name  = mkFastString ("SC:" ++ occNameString fn_occ ++ show rule_number)
               spec_name  = mkInternalName spec_uniq spec_occ fn_loc
---      ; pprTrace "spec_one {" (vcat [ text "function:" <+> ppr fn <+> ppr (idUnique fn)
---                                    , text "sc_count:" <+> ppr (sc_count env)
---                                    , text "pats:" <+> ppr pats
---                                    , text "-->" <+> ppr spec_name
---                                    , text "bndrs" <+> ppr arg_bndrs
---                                    , text "body" <+> ppr body
---                                    , text "how_bound" <+> ppr (sc_how_bound env) ]) $
---        return ()
 
         -- Specialise the body
         -- ; pprTraceM "body_subst_for" $ ppr (spec_occ) $$ ppr (sc_subst body_env)
@@ -1783,9 +1775,10 @@ spec_one env fn arg_bndrs body (call_pat, rule_number)
                   = 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)
-                  | needsVoidWorkerArg fn arg_bndrs spec_lam_args1
+                  | add_void_arg
+                  -- See Note [SpecConst 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.
@@ -1809,17 +1802,63 @@ 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 spec_body
-              rule_rhs   = mkVarApps (Var spec_id) $
-                           dropTail (length extra_bndrs) spec_call_args
+              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
               inline_act = idInlineActivation fn
               this_mod   = sc_module env
               rule       = mkRule this_mod True {- Auto -} True {- Local -}
                                   rule_name inline_act fn_name qvars pats rule_rhs
                            -- See Note [Transfer activation]
+
+        -- ; pprTrace "spec_one {" (vcat [ text "function:" <+> ppr fn <+> ppr (idUnique fn)
+        --                               , text "sc_count:" <+> ppr (sc_count env)
+        --                               , text "pats:" <+> ppr pats
+        --                               , text "call_pat:" <+> ppr call_pat
+        --                               , text "-->" <+> ppr spec_name
+        --                               , text "bndrs" <+> ppr arg_bndrs
+        --                               , text "extra_bndrs" <+> ppr extra_bndrs
+        --                               , text "spec_lam_args" <+> ppr spec_lam_args
+        --                               , text "spec_call_args" <+> ppr spec_call_args
+        --                               , text "rule_rhs" <+> ppr rule_rhs
+        --                               , text "adds_void_worker_arg" <+> ppr adds_void_worker_arg
+        --                               , text "body" <+> ppr body
+        --                               , text "spec_rhs" <+> ppr spec_rhs
+        --                               , text "how_bound" <+> ppr (sc_how_bound env) ]) $
+        --   return ()
         ; return (spec_usg, OS { os_pat = call_pat, os_rule = rule
                                , os_id = spec_id
                                , os_rhs = spec_rhs }) }
 
+{- Note [SpecConst needs to add void args first]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider a function
+    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 @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:
+    $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 easist 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!
+-}
 
 calcSpecInfo :: Id                     -- The original function
              -> CallPat                -- Call pattern
@@ -2251,11 +2290,16 @@ callToPats env bndr_occs call@(Call fn args con_env)
               "SpecConstr: bad covars"
               (ppr bad_covars $$ ppr call) $
           if interesting && isEmptyVarSet bad_covars
-          then
+          then do
               -- pprTraceM "callToPatsOut" (
-              --   text "fun" <> ppr fn $$
+              --         text "fn:" <+> ppr fn $$
+              --         text "args:" <+> ppr args $$
+              --         text "in_scope:" <+> ppr in_scope $$
+              --         -- text "in_scope:" <+> ppr in_scope $$
+              --         text "pat_fvs:" <+> ppr pat_fvs
+              --       )
               --   ppr (CP { cp_qvars = qvars', cp_args = pats })) >>
-                return (Just (CP { cp_qvars = qvars', cp_args = pats }))
+              return (Just (CP { cp_qvars = qvars', cp_args = pats }))
           else return Nothing }
 
     -- argToPat takes an actual argument, and returns an abstracted
@@ -2475,6 +2519,7 @@ setStrUnfolding id str
   = -- trace "setStrUnfolding3"
     id
 
+-- | wildCardPats are always boring
 wildCardPat :: Type -> StrictnessMark -> UniqSM (Bool, CoreArg)
 wildCardPat ty str
   = do { uniq <- getUniqueM
diff --git a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
index 63ac670418..108b9079e6 100644
--- a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
@@ -395,15 +395,16 @@ needsVoidWorkerArg fn_id wrap_args work_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 value argument (but after leading type args).
+-- | Inserts a `Void#` arg before the first argument.
+--
+-- Why as the first argument? See Note [SpecConst needs to add void args first]
+-- in SpecConstr.
 addVoidWorkerArg :: [Var] -> [CbvMark]
                  -> ([Var],     -- Lambda bound args
                      [Var],     -- Args at call site
                      [CbvMark]) -- cbv semantics for the worker args.
 addVoidWorkerArg work_args cbv_marks
-  = (ty_args ++ voidArgId:rest, ty_args ++ voidPrimId:rest, NotMarkedCbv:cbv_marks)
-  where
-    (ty_args, rest) = break isId work_args
+  = (voidArgId : work_args, voidPrimId:work_args, NotMarkedCbv:cbv_marks)
 
 {-
 Note [Protecting the last value argument]
diff --git a/testsuite/tests/simplCore/should_compile/T13143.stderr b/testsuite/tests/simplCore/should_compile/T13143.stderr
index 5ca8a9a503..c80ddb569b 100644
--- a/testsuite/tests/simplCore/should_compile/T13143.stderr
+++ b/testsuite/tests/simplCore/should_compile/T13143.stderr
@@ -1,14 +1,14 @@
 
 ==================== Tidy Core ====================
 Result size of Tidy Core
-  = {terms: 71, types: 40, coercions: 0, joins: 0/0}
+  = {terms: 71, types: 41, coercions: 0, joins: 0/0}
 
 Rec {
 -- RHS size: {terms: 4, types: 3, coercions: 0, joins: 0/0}
 T13143.$wf [InlPrag=NOINLINE, Occ=LoopBreaker]
-  :: forall {a}. (# #) -> a
+  :: (# #) -> forall {a}. a
 [GblId, Arity=1, Str=<B>b, Cpr=b, Unf=OtherCon []]
-T13143.$wf = \ (@a) _ [Occ=Dead] -> T13143.$wf @a GHC.Prim.(##)
+T13143.$wf = \ _ [Occ=Dead] (@a) -> T13143.$wf GHC.Prim.(##) @a
 end Rec }
 
 -- RHS size: {terms: 4, types: 3, coercions: 0, joins: 0/0}
@@ -20,8 +20,8 @@ f [InlPrag=[final]] :: forall a. Int -> a
  Unf=Unf{Src=InlineStable, TopLvl=True, Value=True, ConLike=True,
          WorkFree=True, Expandable=True,
          Guidance=ALWAYS_IF(arity=1,unsat_ok=True,boring_ok=True)
-         Tmpl= \ (@a) _ [Occ=Dead] -> T13143.$wf @a GHC.Prim.(##)}]
-f = \ (@a) _ [Occ=Dead] -> T13143.$wf @a GHC.Prim.(##)
+         Tmpl= \ (@a) _ [Occ=Dead] -> T13143.$wf GHC.Prim.(##) @a}]
+f = \ (@a) _ [Occ=Dead] -> T13143.$wf GHC.Prim.(##) @a
 
 -- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
 T13143.$trModule4 :: GHC.Prim.Addr#
@@ -59,13 +59,13 @@ T13143.$trModule :: GHC.Types.Module
 T13143.$trModule
   = GHC.Types.Module T13143.$trModule3 T13143.$trModule1
 
--- RHS size: {terms: 2, types: 1, coercions: 0, joins: 0/0}
-lvl :: Int
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl :: forall {a}. a
 [GblId, Str=b, Cpr=b]
-lvl = T13143.$wf @Int GHC.Prim.(##)
+lvl = T13143.$wf GHC.Prim.(##)
 
 Rec {
--- RHS size: {terms: 28, types: 7, coercions: 0, joins: 0/0}
+-- RHS size: {terms: 28, types: 8, coercions: 0, joins: 0/0}
 T13143.$wg [InlPrag=[2], Occ=LoopBreaker]
   :: Bool -> Bool -> GHC.Prim.Int# -> GHC.Prim.Int#
 [GblId[StrictWorker([!, !, ~])],