Add arity to the INLINE pragmas for pattern synonymswip/T21531

The lack of INLNE arity was exposed by #21531.  The fix is
simple enough, if a bit clumsy.

2 files changed, 77 insertions, 34 deletions

diff --git a/compiler/GHC/Tc/Gen/Sig.hs b/compiler/GHC/Tc/Gen/Sig.hs
index 971a47bb99..874870765f 100644
--- a/compiler/GHC/Tc/Gen/Sig.hs
+++ b/compiler/GHC/Tc/Gen/Sig.hs
@@ -21,7 +21,8 @@ module GHC.Tc.Gen.Sig(
        tcInstSig,
 
        TcPragEnv, emptyPragEnv, lookupPragEnv, extendPragEnv,
-       mkPragEnv, tcSpecPrags, tcSpecWrapper, tcImpPrags, addInlinePrags
+       mkPragEnv, tcSpecPrags, tcSpecWrapper, tcImpPrags,
+       addInlinePrags, addInlinePragArity
    ) where
 
 import GHC.Prelude
@@ -66,7 +67,6 @@ import GHC.Unit.Module( getModule )
 import GHC.Utils.Misc as Utils ( singleton )
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
-import GHC.Utils.Trace
 
 import GHC.Data.Maybe( orElse )
 
@@ -577,29 +577,32 @@ mkPragEnv sigs binds
     prs = mapMaybe get_sig sigs
 
     get_sig :: LSig GhcRn -> Maybe (Name, LSig GhcRn)
-    get_sig (L l (SpecSig x lnm@(L _ nm) ty inl))
-      = Just (nm, L l $ SpecSig   x lnm ty (add_arity nm inl))
-    get_sig (L l (InlineSig x lnm@(L _ nm) inl))
-      = Just (nm, L l $ InlineSig x lnm    (add_arity nm inl))
-    get_sig (L l (SCCFunSig x st lnm@(L _ nm) str))
-      = Just (nm, L l $ SCCFunSig x st lnm str)
+    get_sig sig@(L _ (SpecSig _ (L _ nm) _ _))   = Just (nm, add_arity nm sig)
+    get_sig sig@(L _ (InlineSig _ (L _ nm) _))   = Just (nm, add_arity nm sig)
+    get_sig sig@(L _ (SCCFunSig _ _ (L _ nm) _)) = Just (nm, sig)
     get_sig _ = Nothing
 
-    add_arity n inl_prag   -- Adjust inl_sat field to match visible arity of function
-      | isInlinePragma inl_prag
-        -- add arity only for real INLINE pragmas, not INLINABLE
+    add_arity n sig  -- Adjust inl_sat field to match visible arity of function
       = case lookupNameEnv ar_env n of
-          Just ar -> inl_prag { inl_sat = Just ar }
-          Nothing -> warnPprTrace True "mkPragEnv no arity" (ppr n) $
-                     -- There really should be a binding for every INLINE pragma
-                     inl_prag
-      | otherwise
-      = inl_prag
+          Just ar -> addInlinePragArity ar sig
+          Nothing -> sig -- See Note [Pattern synonym inline arity]
 
     -- ar_env maps a local to the arity of its definition
     ar_env :: NameEnv Arity
     ar_env = foldr lhsBindArity emptyNameEnv binds
 
+addInlinePragArity :: Arity -> LSig GhcRn -> LSig GhcRn
+addInlinePragArity ar (L l (InlineSig x nm inl))  = L l (InlineSig x nm (add_inl_arity ar inl))
+addInlinePragArity ar (L l (SpecSig x nm ty inl)) = L l (SpecSig x nm ty (add_inl_arity ar inl))
+addInlinePragArity _ sig = sig
+
+add_inl_arity :: Arity -> InlinePragma -> InlinePragma
+add_inl_arity ar prag@(InlinePragma { inl_inline = inl_spec })
+  | Inline {} <- inl_spec  -- Add arity only for real INLINE pragmas, not INLINABLE
+  = prag { inl_sat = Just ar }
+  | otherwise
+  = prag
+
 lhsBindArity :: LHsBind GhcRn -> NameEnv Arity -> NameEnv Arity
 lhsBindArity (L _ (FunBind { fun_id = id, fun_matches = ms })) env
   = extendNameEnv env (unLoc id) (matchGroupArity ms)
@@ -638,6 +641,25 @@ addInlinePrags poly_id prags_for_me
     pp_inl (L loc prag) = ppr prag <+> parens (ppr loc)
 
 
+{- Note [Pattern synonym inline arity]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider
+    {-# INLINE P #-}
+    pattern P x = (x, True)
+
+The INLINE pragma attaches to both the /matcher/ and the /builder/ for
+the pattern synonym; see Note [Pragmas for pattern synonyms] in
+GHC.Tc.TyCl.PatSyn.  But they have different inline arities (i.e. number
+of binders to which we apply the function before inlining), and we don't
+know what those arities are yet.  So for pattern synonyms we don't set
+the inl_sat field yet; instead we do so (via addInlinePragArity) in
+GHC.Tc.TyCl.PatSyn.tcPatSynMatcher and tcPatSynBuilderBind.
+
+It's a bit messy that we set the arities in different ways.  Perhaps we
+should add the arity later for all binders.  But it works fine like this.
+-}
+
+
 {- *********************************************************************
 *                                                                      *
                    SPECIALISE pragmas
diff --git a/compiler/GHC/Tc/TyCl/PatSyn.hs b/compiler/GHC/Tc/TyCl/PatSyn.hs
index 45810f5d9f..a306a3967e 100644
--- a/compiler/GHC/Tc/TyCl/PatSyn.hs
+++ b/compiler/GHC/Tc/TyCl/PatSyn.hs
@@ -27,7 +27,8 @@ import GHC.Core.Type ( tidyTyCoVarBinders, tidyTypes, tidyType, isManyDataConTy
 import GHC.Core.TyCo.Subst( extendTvSubstWithClone )
 import GHC.Tc.Errors.Types
 import GHC.Tc.Utils.Monad
-import GHC.Tc.Gen.Sig ( TcPragEnv, emptyPragEnv, completeSigFromId, lookupPragEnv, addInlinePrags )
+import GHC.Tc.Gen.Sig ( TcPragEnv, emptyPragEnv, completeSigFromId, lookupPragEnv
+                      , addInlinePrags, addInlinePragArity )
 import GHC.Tc.Utils.Env
 import GHC.Tc.Utils.TcMType
 import GHC.Tc.Utils.Zonk
@@ -782,7 +783,7 @@ tcPatSynMatcher :: LocatedN Name
                 -> TcType
                 -> TcM (PatSynMatcher, LHsBinds GhcTc)
 -- See Note [Matchers and builders for pattern synonyms] in GHC.Core.PatSyn
-tcPatSynMatcher (L loc name) lpat prag_fn
+tcPatSynMatcher (L loc ps_name) lpat prag_fn
                 (univ_tvs, req_theta, req_ev_binds, req_dicts)
                 (ex_tvs, ex_tys, prov_theta, prov_dicts)
                 (args, arg_tys) pat_ty
@@ -802,7 +803,7 @@ tcPatSynMatcher (L loc name) lpat prag_fn
 
              fail_ty  = mkVisFunTyMany unboxedUnitTy res_ty
 
-       ; matcher_name <- newImplicitBinder name mkMatcherOcc
+       ; matcher_name <- newImplicitBinder ps_name mkMatcherOcc
        ; scrutinee    <- newSysLocalId (fsLit "scrut") Many pat_ty
        ; cont         <- newSysLocalId (fsLit "cont")  Many cont_ty
        ; fail         <- newSysLocalId (fsLit "fail")  Many fail_ty
@@ -811,7 +812,7 @@ tcPatSynMatcher (L loc name) lpat prag_fn
        ; let matcher_tau   = mkVisFunTysMany [pat_ty, cont_ty, fail_ty] res_ty
              matcher_sigma = mkInfSigmaTy (rr_tv:res_tv:univ_tvs) req_theta matcher_tau
              matcher_id    = mkExportedVanillaId matcher_name matcher_sigma
-             patsyn_id     = mkExportedVanillaId name matcher_sigma
+             patsyn_id     = mkExportedVanillaId ps_name matcher_sigma
                              -- See Note [Exported LocalIds] in GHC.Types.Id
 
              inst_wrap = mkWpEvApps prov_dicts <.> mkWpTyApps ex_tys
@@ -848,16 +849,21 @@ tcPatSynMatcher (L loc name) lpat prag_fn
                     , mg_ext = MatchGroupTc [] res_ty
                     , mg_origin = Generated
                     }
-             prags = lookupPragEnv prag_fn name
+             matcher_arity = length req_theta + 3
              -- See Note [Pragmas for pattern synonyms]
 
-       ; matcher_prag_id <- addInlinePrags matcher_id prags
+       -- Add INLINE pragmas; see Note [Pragmas for pattern synonyms]
+       -- NB: prag_fn is keyed by the PatSyn Name, not the (internal) matcher name
+       ; matcher_prag_id <- addInlinePrags matcher_id              $
+                            map (addInlinePragArity matcher_arity) $
+                            lookupPragEnv prag_fn ps_name
+
        ; let bind = FunBind{ fun_id = L loc matcher_prag_id
                            , fun_matches = mg
                            , fun_ext = idHsWrapper
                            , fun_tick = [] }
              matcher_bind = unitBag (noLocA bind)
-       ; traceTc "tcPatSynMatcher" (ppr name $$ ppr (idType matcher_id))
+       ; traceTc "tcPatSynMatcher" (ppr ps_name $$ ppr (idType matcher_id))
        ; traceTc "tcPatSynMatcher" (ppr matcher_bind)
 
        ; return ((matcher_name, matcher_sigma, is_unlifted), matcher_bind) }
@@ -905,6 +911,7 @@ mkPatSynBuilder dir (L _ name)
                               mkPhiTy theta $
                               mkVisFunTysMany arg_tys $
                               pat_ty
+
        ; return (Just (builder_name, builder_sigma, need_dummy_arg)) }
 
 tcPatSynBuilderBind :: TcPragEnv
@@ -937,11 +944,17 @@ tcPatSynBuilderBind prag_fn (PSB { psb_id = ps_lname@(L loc ps_name)
     do { -- Bidirectional, so patSynBuilder returns Just
          let builder_id = mkExportedVanillaId builder_name builder_ty
                          -- See Note [Exported LocalIds] in GHC.Types.Id
-             prags = lookupPragEnv prag_fn ps_name
-             -- See Note [Pragmas for pattern synonyms]
-             -- Keyed by the PatSyn Name, not the (internal) builder name
 
-       ; builder_id <- addInlinePrags builder_id prags
+             (_, req_theta, _, prov_theta, arg_tys, _) = patSynSigBndr patsyn
+             builder_arity = length req_theta + length prov_theta
+                             + length arg_tys
+                             + (if need_dummy_arg then 1 else 0)
+
+       -- Add INLINE pragmas; see Note [Pragmas for pattern synonyms]
+       -- NB: prag_fn is keyed by the PatSyn Name, not the (internal) builder name
+       ; builder_id <- addInlinePrags builder_id              $
+                       map (addInlinePragArity builder_arity) $
+                       lookupPragEnv prag_fn ps_name
 
        ; let match_group' | need_dummy_arg = add_dummy_arg match_group
                           | otherwise      = match_group
@@ -955,8 +968,7 @@ tcPatSynBuilderBind prag_fn (PSB { psb_id = ps_lname@(L loc ps_name)
 
        ; traceTc "tcPatSynBuilderBind {" $
          vcat [ ppr patsyn
-              , ppr builder_id <+> dcolon <+> ppr (idType builder_id)
-              , ppr prags ]
+              , ppr builder_id <+> dcolon <+> ppr (idType builder_id) ]
        ; (builder_binds, _) <- tcPolyCheck emptyPragEnv sig (noLocA bind)
        ; traceTc "tcPatSynBuilderBind }" $ ppr builder_binds
        ; return builder_binds } } }
@@ -1216,18 +1228,19 @@ want to avoid difficult to decipher core lint errors!
 
 Note [Pragmas for pattern synonyms]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-INLINE and NOINLINE pragmas are supported for pattern synonyms. They affect both
-the matcher and the builder.
+INLINE and NOINLINE pragmas are supported for pattern synonyms.
+They affect both the matcher and the builder.
 (See Note [Matchers and builders for pattern synonyms] in PatSyn)
 
 For example:
     pattern InlinedPattern x = [x]
     {-# INLINE InlinedPattern #-}
+
     pattern NonInlinedPattern x = [x]
     {-# NOINLINE NonInlinedPattern #-}
 
-For pattern synonyms with explicit builders, only pragma for the entire pattern
-synonym is supported. For example:
+For pattern synonyms with explicit builders, only a pragma for the
+entire pattern synonym is supported. For example:
     pattern HeadC x <- x:xs where
       HeadC x = [x]
       -- This wouldn't compile: {-# INLINE HeadC #-}
@@ -1235,6 +1248,14 @@ synonym is supported. For example:
 
 When no pragma is provided for a pattern, the inlining decision might change
 between different versions of GHC.
+
+Implementation notes.  The prag_fn passed in to tcPatSynDecl will have a binding
+for the /pattern synonym/ Name, thus
+      InlinedPattern :-> INLINE
+From this we cook up an INLINE pragma for the matcher (in tcPatSynMatcher)
+and builder (in tcPatSynBuilderBind), by looking up the /pattern synonym/
+Name in the prag_fn, and then using addInlinePragArity to add the right
+inl_sat field to that INLINE pragma for the matcher or builder respectively.
  -}