Bring back COMPLETE sets filtered by result TyCon (#14422)

Commit 2a94228 dramatically simplified the implementation and improved
the performance of COMPLETE sets while making them applicable in more
scenarios at the same time.
But it turned out that there was a change in semantics that (to me
unexpectedly) broke users' expectations (see #14422): They relied on the
"type signature" of a COMPLETE pragma to restrict the scrutinee types of
a pattern match for which they are applicable.

This patch brings back that filtering, so the semantics is the same as
it was in GHC 9.0.
See the updated Note [Implementation of COMPLETE pragmas].

There are a few testsuite output changes (`completesig13`, `T14422`)
which assert this change.

Co-authored-by: Sebastian Graf <sebastian.graf@kit.edu>

10 files changed, 121 insertions, 50 deletions

diff --git a/compiler/GHC/Hs/Binds.hs b/compiler/GHC/Hs/Binds.hs
index 5316046880..91b5dd7724 100644
--- a/compiler/GHC/Hs/Binds.hs
+++ b/compiler/GHC/Hs/Binds.hs
@@ -604,7 +604,7 @@ ppr_sig (CompleteMatchSig _ src cs mty)
       ((hsep (punctuate comma (map ppr (unLoc cs))))
         <+> opt_sig)
   where
-    opt_sig = maybe empty ((\t -> dcolon <+> ppr t) . unLoc) mty
+    opt_sig = maybe empty (\t -> dcolon <+> ppr t) mty
 
 instance OutputableBndrId p
        => Outputable (FixitySig (GhcPass p)) where
diff --git a/compiler/GHC/HsToCore/Pmc/Solver.hs b/compiler/GHC/HsToCore/Pmc/Solver.hs
index 7635d0bb25..b128cc93fd 100644
--- a/compiler/GHC/HsToCore/Pmc/Solver.hs
+++ b/compiler/GHC/HsToCore/Pmc/Solver.hs
@@ -48,6 +48,7 @@ import GHC.Utils.Error ( pprMsgEnvelopeBagWithLoc )
 import GHC.Utils.Misc
 import GHC.Utils.Panic
 import GHC.Data.Bag
+import GHC.Types.CompleteMatch
 import GHC.Types.Error
 import GHC.Types.Unique.Set
 import GHC.Types.Unique.DSet
@@ -147,7 +148,7 @@ vanillaCompleteMatchTC tc =
       -- special case.
       mb_dcs | tc == tYPETyCon = Just []
              | otherwise       = tyConDataCons_maybe tc
-  in mkUniqDSet . map RealDataCon <$> mb_dcs
+  in vanillaCompleteMatch . mkUniqDSet . map RealDataCon <$> mb_dcs
 
 -- | Initialise from 'dsGetCompleteMatches' (containing all COMPLETE pragmas)
 -- if the given 'ResidualCompleteMatches' were empty.
@@ -180,9 +181,9 @@ markMatched :: PmAltCon -> ResidualCompleteMatches -> DsM (Maybe ResidualComplet
 markMatched (PmAltLit _)      _   = pure Nothing -- lits are never part of a COMPLETE set
 markMatched (PmAltConLike cl) rcm = do
   rcm' <- addConLikeMatches cl rcm
-  let go cm = case lookupUniqDSet cm cl of
+  let go cm = case lookupUniqDSet (cmConLikes cm) cl of
         Nothing -> (False, cm)
-        Just _  -> (True,  delOneFromUniqDSet cm cl)
+        Just _  -> (True,  cm { cmConLikes = delOneFromUniqDSet (cmConLikes cm) cl })
   pure $ updRcm go rcm'
 
 {-
@@ -203,10 +204,34 @@ function, it gives rise to a total function. An example is:
   booleanToInt F = 0
   booleanToInt T = 1
 
-COMPLETE sets are represented internally in GHC a set of 'ConLike's. For
+COMPLETE sets are represented internally in GHC as a set of 'ConLike's. For
 example, the pragma {-# COMPLETE F, T #-} would be represented as:
 
-  {F, T}
+  CompleteMatch {F, T} Nothing
+
+What is the Maybe for? Answer: COMPLETE pragmas may optionally specify a
+result *type constructor* (cf. T14422):
+
+  class C f where
+    foo :: f a -> ()
+  pattern P :: C f => f a
+  pattern P <- (foo -> ())
+
+  instance C State where
+    foo _ = ()
+  {-# COMPLETE P :: State #-}
+
+  f :: State a -> ()
+  f P = ()
+  g :: C f => f a -> ()
+  g P = ()
+
+The @:: State@ here means that the types at which the COMPLETE pragma *applies*
+is restricted to scrutinee types that are applications of the 'State' TyCon. So
+it applies to the match in @f@ but not in @g@ above, resulting in a warning for
+the latter but not for the former. The pragma is represented as
+
+  CompleteMatch {P} (Just State)
 
 GHC collects all COMPLETE pragmas from the current module and from imports
 into a field in the DsM environment, which can be accessed with
@@ -228,18 +253,20 @@ we know a particular variable can't be (through negative constructor constraints
 testing). If *any* of the COMPLETE sets become empty, we know that the match
 was exhaustive.
 
-We assume that a COMPLETE set is non-empty if for one of its ConLikes
-we fail to 'guessConLikeUnivTyArgsFromResTy'. That accounts for ill-typed
-COMPLETE sets. So why don't we simply prune those ill-typed COMPLETE sets from
-'ResidualCompleteMatches'? The answer is that additional type constraints might
-make more COMPLETE sets applicable! Example:
+We assume that a COMPLETE set does not apply if for one of its
+ConLikes we fail to 'guessConLikeUnivTyArgsFromResTy' or the
+type of the match variable isn't an application of the optional
+result type constructor from the pragma. Why don't we simply
+prune inapplicable COMPLETE sets from 'ResidualCompleteMatches'?
+The answer is that additional type constraints might make more
+COMPLETE sets applicable! Example:
 
-  f :: a -> a :~: Boolean -> ()
-  f x Refl | T <- x = ()
+  h :: a -> a :~: Boolean -> ()
+  h x Refl | T <- x = ()
            | F <- x = ()
 
 If we eagerly prune {F,T} from the residual matches of @x@, then we don't see
-that the match in the guards of @f@ is exhaustive, where the COMPLETE set
+that the match in the guards of @h@ is exhaustive, where the COMPLETE set
 applies due to refined type information.
 -}
 
@@ -1338,7 +1365,7 @@ anyConLikeSolution p = any (go . paca_con)
     go (PmAltConLike cl) = p cl
     go _                 = False
 
--- | @instCompleteSet fuel nabla nabla cls@ iterates over @cls@ until it finds
+-- | @instCompleteSet fuel nabla x cls@ iterates over @cls@ until it finds
 -- the first inhabited ConLike (as per 'instCon'). Any failed instantiation
 -- attempts of a ConLike are recorded as negative information in the returned
 -- 'Nabla', so that later calls to this function can skip repeatedly fruitless
@@ -1350,23 +1377,26 @@ anyConLikeSolution p = any (go . paca_con)
 -- entirely as an optimisation.
 instCompleteSet :: Int -> Nabla -> Id -> CompleteMatch -> MaybeT DsM Nabla
 instCompleteSet fuel nabla x cs
-  | anyConLikeSolution (`elementOfUniqDSet` cs) (vi_pos vi)
+  | anyConLikeSolution (`elementOfUniqDSet` (cmConLikes cs)) (vi_pos vi)
   -- No need to instantiate a constructor of this COMPLETE set if we already
   -- have a solution!
   = pure nabla
+  | not (completeMatchAppliesAtType (varType x) cs)
+  = pure nabla
   | otherwise
   = go nabla (sorted_candidates cs)
   where
     vi = lookupVarInfo (nabla_tm_st nabla) x
 
     sorted_candidates :: CompleteMatch -> [ConLike]
-    sorted_candidates cs
+    sorted_candidates cm
       -- If there aren't many candidates, we can try to sort them by number of
       -- strict fields, type constraints, etc., so that we are fast in the
       -- common case
       -- (either many simple constructors *or* few "complicated" ones).
       | sizeUniqDSet cs <= 5 = sortBy compareConLikeTestability (uniqDSetToList cs)
       | otherwise            = uniqDSetToList cs
+      where cs = cmConLikes cm
 
     go :: Nabla -> [ConLike] -> MaybeT DsM Nabla
     go _     []         = mzero
@@ -1780,7 +1810,7 @@ generateInhabitingPatterns (x:xs) n nabla = do
 
           -- Test all COMPLETE sets for inhabitants (n inhs at max). Take care of ⊥.
           clss <- pickApplicableCompleteSets rep_ty rcm
-          case NE.nonEmpty (uniqDSetToList <$> clss) of
+          case NE.nonEmpty (uniqDSetToList . cmConLikes <$> clss) of
             Nothing ->
               -- No COMPLETE sets ==> inhabited
               generateInhabitingPatterns xs n newty_nabla
@@ -1831,9 +1861,20 @@ generateInhabitingPatterns (x:xs) n nabla = do
       pure (con_nablas ++ other_cons_nablas)
 
 pickApplicableCompleteSets :: Type -> ResidualCompleteMatches -> DsM [CompleteMatch]
+-- See Note [Implementation of COMPLETE pragmas] on what "applicable" means
 pickApplicableCompleteSets ty rcm = do
   env <- dsGetFamInstEnvs
-  pure $ filter (all (is_valid env) . uniqDSetToList) (getRcm rcm)
+  let applicable :: CompleteMatch -> Bool
+      applicable cm = all (is_valid env) (uniqDSetToList (cmConLikes cm))
+                   && completeMatchAppliesAtType ty cm
+      applicableMatches = filter applicable (getRcm rcm)
+  tracePm "pickApplicableCompleteSets:" $
+    vcat
+      [ ppr ty
+      , ppr rcm
+      , ppr applicableMatches
+      ]
+  return applicableMatches
   where
     is_valid :: FamInstEnvs -> ConLike -> Bool
     is_valid env cl = isJust (guessConLikeUnivTyArgsFromResTy env ty cl)
diff --git a/compiler/GHC/HsToCore/Pmc/Solver/Types.hs b/compiler/GHC/HsToCore/Pmc/Solver/Types.hs
index 1e4e672583..7516a56995 100644
--- a/compiler/GHC/HsToCore/Pmc/Solver/Types.hs
+++ b/compiler/GHC/HsToCore/Pmc/Solver/Types.hs
@@ -61,10 +61,10 @@ import GHC.Builtin.Types
 import GHC.Builtin.Types.Prim
 import GHC.Tc.Solver.Monad (InertSet, emptyInert)
 import GHC.Tc.Utils.TcType (isStringTy)
-import GHC.Types.CompleteMatch (CompleteMatch)
-import GHC.Types.SourceText (mkFractionalLit, FractionalLit, fractionalLitFromRational,
-  FractionalExponentBase(..), SourceText(..))
-
+import GHC.Types.CompleteMatch (CompleteMatch(..))
+import GHC.Types.SourceText (SourceText(..), mkFractionalLit, FractionalLit
+                            , fractionalLitFromRational
+                            , FractionalExponentBase(..))
 import Numeric (fromRat)
 import Data.Foldable (find)
 import Data.Ratio
@@ -368,7 +368,7 @@ eqConLike _                 _                 = PossiblyOverlap
 data PmAltCon = PmAltConLike ConLike
               | PmAltLit     PmLit
 
-data PmAltConSet = PACS !CompleteMatch ![PmLit]
+data PmAltConSet = PACS !(UniqDSet ConLike) ![PmLit]
 
 emptyPmAltConSet :: PmAltConSet
 emptyPmAltConSet = PACS emptyUniqDSet []
diff --git a/compiler/GHC/Iface/Make.hs b/compiler/GHC/Iface/Make.hs
index 1c43e3e6e6..53f0032f28 100644
--- a/compiler/GHC/Iface/Make.hs
+++ b/compiler/GHC/Iface/Make.hs
@@ -73,6 +73,7 @@ import GHC.Types.TypeEnv
 import GHC.Types.SourceFile
 import GHC.Types.TyThing
 import GHC.Types.HpcInfo
+import GHC.Types.CompleteMatch
 
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
@@ -347,8 +348,8 @@ mkIface_ hsc_env
 -}
 
 mkIfaceCompleteMatch :: CompleteMatch -> IfaceCompleteMatch
-mkIfaceCompleteMatch cls =
-  IfaceCompleteMatch (map conLikeName (uniqDSetToList cls))
+mkIfaceCompleteMatch (CompleteMatch cls mtc) =
+  IfaceCompleteMatch (map conLikeName (uniqDSetToList cls)) (toIfaceTyCon <$> mtc)
 
 
 {-
diff --git a/compiler/GHC/Iface/Syntax.hs b/compiler/GHC/Iface/Syntax.hs
index 73e8525589..21b4274cc7 100644
--- a/compiler/GHC/Iface/Syntax.hs
+++ b/compiler/GHC/Iface/Syntax.hs
@@ -325,14 +325,12 @@ data IfaceAnnotation
 
 type IfaceAnnTarget = AnnTarget OccName
 
-newtype IfaceCompleteMatch = IfaceCompleteMatch [IfExtName]
+data IfaceCompleteMatch = IfaceCompleteMatch [IfExtName] (Maybe IfaceTyCon)
 
 instance Outputable IfaceCompleteMatch where
-  ppr (IfaceCompleteMatch cls) = text "COMPLETE" <> colon <+> ppr cls
-
-
-
-
+  ppr (IfaceCompleteMatch cls mtc) = text "COMPLETE" <> colon <+> ppr cls <+> case mtc of
+    Nothing -> empty
+    Just tc -> dcolon <+> ppr tc
 
 -- Here's a tricky case:
 --   * Compile with -O module A, and B which imports A.f
@@ -2493,8 +2491,8 @@ instance Binary IfaceTyConParent where
                 return $ IfDataInstance ax pr ty
 
 instance Binary IfaceCompleteMatch where
-  put_ bh (IfaceCompleteMatch cs) = put_ bh cs
-  get bh = IfaceCompleteMatch <$> get bh
+  put_ bh (IfaceCompleteMatch cs mtc) = put_ bh cs >> put_ bh mtc
+  get bh = IfaceCompleteMatch <$> get bh <*> get bh
 
 
 {-
@@ -2653,7 +2651,7 @@ instance NFData IfaceConAlt where
     IfaceLitAlt lit -> lit `seq` ()
 
 instance NFData IfaceCompleteMatch where
-  rnf (IfaceCompleteMatch f1) = rnf f1
+  rnf (IfaceCompleteMatch f1 mtc) = rnf f1 `seq` rnf mtc
 
 instance NFData IfaceRule where
   rnf (IfaceRule f1 f2 f3 f4 f5 f6 f7 f8) =
diff --git a/compiler/GHC/IfaceToCore.hs b/compiler/GHC/IfaceToCore.hs
index 5a843c5e7e..76079ae8ff 100644
--- a/compiler/GHC/IfaceToCore.hs
+++ b/compiler/GHC/IfaceToCore.hs
@@ -85,6 +85,7 @@ import GHC.Types.Annotations
 import GHC.Types.SourceFile
 import GHC.Types.SourceText
 import GHC.Types.Basic hiding ( SuccessFlag(..) )
+import GHC.Types.CompleteMatch
 import GHC.Types.SrcLoc
 import GHC.Types.TypeEnv
 import GHC.Types.Unique.FM
@@ -1280,8 +1281,10 @@ tcIfaceCompleteMatches :: [IfaceCompleteMatch] -> IfL [CompleteMatch]
 tcIfaceCompleteMatches = mapM tcIfaceCompleteMatch
 
 tcIfaceCompleteMatch :: IfaceCompleteMatch -> IfL CompleteMatch
-tcIfaceCompleteMatch (IfaceCompleteMatch ms) =
-  mkUniqDSet <$> mapM (forkM doc . tcIfaceConLike) ms
+tcIfaceCompleteMatch (IfaceCompleteMatch ms mtc) = do
+  conlikes <- mkUniqDSet <$> mapM (forkM doc . tcIfaceConLike) ms
+  mtc' <- traverse tcIfaceTyCon mtc
+  return (CompleteMatch conlikes mtc')
   where
     doc = text "COMPLETE sig" <+> ppr ms
 
diff --git a/compiler/GHC/IfaceToCore.hs-boot b/compiler/GHC/IfaceToCore.hs-boot
index c21c4a3acb..97124237c7 100644
--- a/compiler/GHC/IfaceToCore.hs-boot
+++ b/compiler/GHC/IfaceToCore.hs-boot
@@ -8,7 +8,7 @@ import GHC.Tc.Types        ( IfL )
 import GHC.Core.InstEnv    ( ClsInst )
 import GHC.Core.FamInstEnv ( FamInst )
 import GHC.Core         ( CoreRule )
-import GHC.Types.CompleteMatch ( CompleteMatch )
+import GHC.Types.CompleteMatch
 import GHC.Types.Annotations ( Annotation )
 import GHC.Types.Name
 import GHC.Fingerprint.Type
diff --git a/compiler/GHC/Tc/Gen/Bind.hs b/compiler/GHC/Tc/Gen/Bind.hs
index caaa8b4894..0ab561a0a7 100644
--- a/compiler/GHC/Tc/Gen/Bind.hs
+++ b/compiler/GHC/Tc/Gen/Bind.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE RankNTypes          #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications    #-}
 {-# LANGUAGE TypeFamilies        #-}
 
 {-
@@ -65,6 +66,7 @@ import GHC.Data.Graph.Directed
 import GHC.Data.Maybe
 import GHC.Utils.Misc
 import GHC.Types.Basic
+import GHC.Types.CompleteMatch
 import GHC.Utils.Outputable as Outputable
 import GHC.Utils.Panic
 import GHC.Builtin.Names( ipClassName )
@@ -203,11 +205,14 @@ tcCompleteSigs sigs =
   let
       doOne :: LSig GhcRn -> TcM (Maybe CompleteMatch)
       -- We don't need to "type-check" COMPLETE signatures anymore; if their
-      -- combinations are invalid it will be found so at match sites. Hence we
-      -- keep '_mtc' only for backwards compatibility.
-      doOne (L loc c@(CompleteMatchSig _ext _src_txt (L _ ns) _mtc))
-        = fmap Just $ setSrcSpan loc $ addErrCtxt (text "In" <+> ppr c) $
-            mkUniqDSet <$> mapM (addLocM tcLookupConLike) ns
+      -- combinations are invalid it will be found so at match sites.
+      -- There it is also where we consider if the type of the pattern match is
+      -- compatible with the result type constructor 'mb_tc'.
+      doOne (L loc c@(CompleteMatchSig _ext _src_txt (L _ ns) mb_tc_nm))
+        = fmap Just $ setSrcSpan loc $ addErrCtxt (text "In" <+> ppr c) $ do
+            cls   <- mkUniqDSet <$> mapM (addLocM tcLookupConLike) ns
+            mb_tc <- traverse @Maybe tcLookupLocatedTyCon mb_tc_nm
+            pure CompleteMatch { cmConLikes = cls, cmResultTyCon = mb_tc }
       doOne _ = return Nothing
 
   -- For some reason I haven't investigated further, the signatures come in
diff --git a/compiler/GHC/ThToHs.hs b/compiler/GHC/ThToHs.hs
index dc10c6fed5..12f65d36ca 100644
--- a/compiler/GHC/ThToHs.hs
+++ b/compiler/GHC/ThToHs.hs
@@ -827,7 +827,7 @@ cvtPragmaD (LineP line file)
        }
 cvtPragmaD (CompleteP cls mty)
   = do { cls' <- noLoc <$> mapM cNameL cls
-       ; mty'  <- traverse tconNameL mty
+       ; mty' <- traverse tconNameL mty
        ; returnJustL $ Hs.SigD noExtField
                    $ CompleteMatchSig noExtField NoSourceText cls' mty' }
 
diff --git a/compiler/GHC/Types/CompleteMatch.hs b/compiler/GHC/Types/CompleteMatch.hs
index 7ad521f738..43216eba12 100644
--- a/compiler/GHC/Types/CompleteMatch.hs
+++ b/compiler/GHC/Types/CompleteMatch.hs
@@ -1,17 +1,40 @@
+{-# LANGUAGE TypeApplications #-}
+
 -- | COMPLETE signature
-module GHC.Types.CompleteMatch
-   ( CompleteMatch
-   , CompleteMatches
-   )
-where
+module GHC.Types.CompleteMatch where
 
+import GHC.Prelude
+import GHC.Core.TyCo.Rep
 import GHC.Types.Unique.DSet
 import GHC.Core.ConLike
+import GHC.Core.TyCon
+import GHC.Core.Type ( splitTyConApp_maybe )
+import GHC.Utils.Outputable
 
 -- | A list of conlikes which represents a complete pattern match.
 -- These arise from @COMPLETE@ signatures.
 -- See also Note [Implementation of COMPLETE pragmas].
-type CompleteMatch = UniqDSet ConLike
+data CompleteMatch = CompleteMatch
+  { cmConLikes :: UniqDSet ConLike -- ^ The set of `ConLike` values
+  , cmResultTyCon :: Maybe TyCon   -- ^ The optional, concrete result TyCon the set applies to
+  }
+
+vanillaCompleteMatch :: UniqDSet ConLike -> CompleteMatch
+vanillaCompleteMatch cls = CompleteMatch { cmConLikes = cls, cmResultTyCon = Nothing }
+
+instance Outputable CompleteMatch where
+  ppr (CompleteMatch cls mty) = case mty of
+    Nothing -> ppr cls
+    Just ty -> ppr cls <> text "@" <> parens (ppr ty)
 
 type CompleteMatches = [CompleteMatch]
 
+completeMatchAppliesAtType :: Type -> CompleteMatch -> Bool
+completeMatchAppliesAtType ty cm = all @Maybe ty_matches (cmResultTyCon cm)
+  where
+    ty_matches sig_tc
+      | Just (tc, _arg_tys) <- splitTyConApp_maybe ty
+      , tc == sig_tc
+      = True
+      | otherwise
+      = False