Improve partial signatures

As #20921 showed, with partial signatures, it is helpful to use the
same algorithm (namely findInferredDiff) for
* picking the constraints to retain for the /group/
  in Solver.decideQuantification
* picking the contraints to retain for the /individual function/
  in Bind.chooseInferredQuantifiers

This is still regrettably declicate, but it's a step forward.

3 files changed, 89 insertions, 21 deletions

diff --git a/compiler/GHC/Core/Ppr.hs b/compiler/GHC/Core/Ppr.hs
index f1791dfebf..853ed990bc 100644
--- a/compiler/GHC/Core/Ppr.hs
+++ b/compiler/GHC/Core/Ppr.hs
@@ -22,6 +22,7 @@ module GHC.Core.Ppr (
         pprCoreExpr, pprParendExpr,
         pprCoreBinding, pprCoreBindings, pprCoreAlt,
         pprCoreBindingWithSize, pprCoreBindingsWithSize,
+        pprCoreBinder, pprCoreBinders,
         pprRules, pprOptCo
     ) where
 
@@ -355,7 +356,6 @@ and @pprCoreExpr@ functions.
 
 Note [Binding-site specific printing]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
 pprCoreBinder and pprTypedLamBinder receive a BindingSite argument to adjust
 the information printed.
 
@@ -393,6 +393,10 @@ pprCoreBinder bind_site bndr
   = getPprDebug $ \debug ->
     pprTypedLamBinder bind_site debug bndr
 
+pprCoreBinders :: [Var] -> SDoc
+-- Print as lambda-binders, i.e. with their type
+pprCoreBinders vs = sep (map (pprCoreBinder LambdaBind) vs)
+
 pprUntypedBinder :: Var -> SDoc
 pprUntypedBinder binder
   | isTyVar binder = text "@" <> ppr binder    -- NB: don't print kind
@@ -642,8 +646,7 @@ pprRule (Rule { ru_name = name, ru_act = act, ru_fn = fn,
                 ru_bndrs = tpl_vars, ru_args = tpl_args,
                 ru_rhs = rhs })
   = hang (doubleQuotes (ftext name) <+> ppr act)
-       4 (sep [text "forall" <+>
-                  sep (map (pprCoreBinder LambdaBind) tpl_vars) <> dot,
+       4 (sep [text "forall" <+> pprCoreBinders tpl_vars <> dot,
                nest 2 (ppr fn <+> sep (map pprArg tpl_args)),
                nest 2 (text "=" <+> pprCoreExpr rhs)
             ])
diff --git a/compiler/GHC/Tc/Gen/Bind.hs b/compiler/GHC/Tc/Gen/Bind.hs
index 93fa9a7e2c..052a2ed2fa 100644
--- a/compiler/GHC/Tc/Gen/Bind.hs
+++ b/compiler/GHC/Tc/Gen/Bind.hs
@@ -789,6 +789,9 @@ mkExport prag_fn insoluble qtvs theta
                 -- This type is just going into tcSubType,
                 -- so Inferred vs. Specified doesn't matter
 
+        ; traceTc "mkExport" (vcat [ ppr poly_id <+> dcolon <+> ppr poly_ty
+                                   , ppr sel_poly_ty ])
+
         ; wrap <- if sel_poly_ty `eqType` poly_ty  -- NB: eqType ignores visibility
                   then return idHsWrapper  -- Fast path; also avoids complaint when we infer
                                            -- an ambiguous type and have AllowAmbiguousType
diff --git a/compiler/GHC/Tc/Solver.hs b/compiler/GHC/Tc/Solver.hs
index aba2ad54a8..5319a52ad0 100644
--- a/compiler/GHC/Tc/Solver.hs
+++ b/compiler/GHC/Tc/Solver.hs
@@ -58,6 +58,7 @@ import GHC.Core.Predicate
 import GHC.Tc.Types.Origin
 import GHC.Tc.Utils.TcType
 import GHC.Core.Type
+import GHC.Core.Ppr
 import GHC.Builtin.Types ( liftedRepTy, manyDataConTy, liftedDataConTy )
 import GHC.Core.Unify    ( tcMatchTyKi )
 import GHC.Utils.Misc
@@ -677,7 +678,6 @@ than one path, this alternative doesn't work.
 
 Note [Safe Haskell Overlapping Instances Implementation]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
 How is this implemented? It's complicated! So we'll step through it all:
 
  1) `InstEnv.lookupInstEnv` -- Performs instance resolution, so this is where
@@ -1081,7 +1081,7 @@ simplifyInfer rhs_tclvl infer_mode sigs name_taus wanteds
             <- setTcLevel rhs_tclvl $
                runTcSWithEvBinds ev_binds_var $
                solveWanteds (mkSimpleWC psig_evs `andWC` wanteds)
-               -- psig_evs : see Note [Add signature contexts as givens]
+               -- psig_evs : see Note [Add signature contexts as wanteds]
 
        -- Find quant_pred_candidates, the predicates that
        -- we'll consider quantifying over
@@ -1117,9 +1117,9 @@ simplifyInfer rhs_tclvl infer_mode sigs name_taus wanteds
          -- All done!
        ; traceTc "} simplifyInfer/produced residual implication for quantification" $
          vcat [ text "quant_pred_candidates =" <+> ppr quant_pred_candidates
-              , text "psig_theta =" <+> ppr psig_theta
-              , text "bound_theta =" <+> ppr bound_theta
-              , text "qtvs ="       <+> ppr qtvs
+              , text "psig_theta ="     <+> ppr psig_theta
+              , text "bound_theta ="    <+> pprCoreBinders bound_theta_vars
+              , text "qtvs ="           <+> ppr qtvs
               , text "definite_error =" <+> ppr definite_error ]
 
        ; return ( qtvs, bound_theta_vars, TcEvBinds ev_binds_var, definite_error ) }
@@ -1177,7 +1177,14 @@ ctsPreds cts = [ ctEvPred ev | ct <- bagToList cts
                              , let ev = ctEvidence ct ]
 
 findInferredDiff :: TcThetaType -> TcThetaType -> TcM TcThetaType
+-- Given a partial type signature f :: (C a, D a, _) => blah
+-- and the inferred constraints (X a, D a, Y a, C a)
+-- compute the difference, which is what will fill in the "_" underscore,
+-- In this case the diff is (X a, Y a).
 findInferredDiff annotated_theta inferred_theta
+  | null annotated_theta   -- Short cut the common case when the user didn't
+  = return inferred_theta  -- write any constraints in the partial signature
+  | otherwise
   = pushTcLevelM_ $
     do { lcl_env   <- TcM.getLclEnv
        ; given_ids <- mapM TcM.newEvVar annotated_theta
@@ -1254,7 +1261,24 @@ We'll use plan InferGen because there are holes in the type.  But:
    fundeps
 
 Solution: in simplifyInfer, we add the constraints from the signature
-as extra Wanteds
+as extra Wanteds.
+
+Why Wanteds?  Wouldn't it be neater to treat them as Givens?  Alas
+that would mess up (GivenInv) in Note [TcLevel invariants].  Consider
+    f :: (Eq a, _) => blah1
+    f = ....g...
+    g :: (Eq b, _) => blah2
+    g = ...f...
+
+Then we have two psig_theta constraints (Eq a[tv], Eq b[tv]), both with
+TyVarTvs inside.  Ultimately a[tv] := b[tv], but only when we've solved
+all those constraints.  And both have level 1, so we can't put them as
+Givens when solving at level 1.
+
+Best to treat them as Wanteds.
+
+But see also #20076, which would be solved if they were Givens.
+
 
 ************************************************************************
 *                                                                      *
@@ -1335,17 +1359,23 @@ decideQuantification infer_mode rhs_tclvl name_taus psigs candidates
        -- into quantified skolems, so we have to zonk again
        ; candidates <- TcM.zonkTcTypes candidates
        ; psig_theta <- TcM.zonkTcTypes (concatMap sig_inst_theta psigs)
-       ; let quantifiable_candidates
-               = pickQuantifiablePreds (mkVarSet qtvs) candidates
-
-             theta = mkMinimalBySCs id $  -- See Note [Minimize by Superclasses]
-                     psig_theta ++ quantifiable_candidates
-             -- NB: add psig_theta back in here, even though it's already
-             -- part of candidates, because we always want to quantify over
-             -- psig_theta, and pickQuantifiableCandidates might have
-             -- dropped some e.g. CallStack constraints.  c.f #14658
-             --                   equalities (a ~ Bool)
-             -- Remember, this is the theta for the residual constraint
+       ; let min_theta = mkMinimalBySCs id $  -- See Note [Minimize by Superclasses]
+                         pickQuantifiablePreds (mkVarSet qtvs) candidates
+
+             min_psig_theta = mkMinimalBySCs id psig_theta
+
+       -- Add psig_theta back in here, even though it's already
+       -- part of candidates, because we always want to quantify over
+       -- psig_theta, and pickQuantifiableCandidates might have
+       -- dropped some e.g. CallStack constraints.  c.f #14658
+       --                   equalities (a ~ Bool)
+       -- It's helpful to use the same "find difference" algorithm here as
+       -- we use in GHC.Tc.Gen.Bind.chooseInferredQuantifiers (#20921)
+       -- See Note [Constraints in partial type signatures]
+       ; theta <- if null psig_theta
+                  then return min_theta  -- Fast path for the non-partial-sig case
+                  else do { diff <- findInferredDiff min_psig_theta min_theta
+                          ; return (min_psig_theta ++ diff) }
 
        ; traceTc "decideQuantification"
            (vcat [ text "infer_mode:" <+> ppr infer_mode
@@ -1357,7 +1387,39 @@ decideQuantification infer_mode rhs_tclvl name_taus psigs candidates
                  , text "theta:"      <+> ppr theta ])
        ; return (qtvs, theta, co_vars) }
 
-------------------
+{- Note [Constraints in partial type signatures]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Suppose we have a partial type signature
+    f :: (Eq a, C a, _) => blah
+
+We will ultimately quantify f over (Eq a, C a, <diff>), where
+
+   * <diff> is the result of
+         findInferredDiff (Eq a, C a) <quant-theta>
+     in GHC.Tc.Gen.Bind.chooseInferredQuantifiers
+
+   * <quant-theta> is the theta returned right here,
+     by decideQuantification
+
+At least for single functions we would like to quantify f over
+precisely the same theta as <quant-theta>, so that we get to take
+the short-cut path in GHC.Tc.Gen.Bind.mkExport, and avoid calling
+tcSubTypeSigma for impedence matching. Why avoid?  Because it falls
+over for ambiguous types (#20921).
+
+We can get precisely the same theta by using the same algorithm,
+findInferredDiff.
+
+All of this goes wrong if we have (a) mutual recursion, (b) mutiple
+partial type signatures, (c) with different constraints, and (d)
+ambiguous types.  Something like
+    f :: forall a. Eq a => F a -> _
+    f x = (undefined :: a) == g x undefined
+    g :: forall b. Show b => F b -> _ -> b
+    g x y = let _ = (f y, show x) in x
+But that's a battle for another day.
+-}
+
 decideMonoTyVars :: InferMode
                  -> [(Name,TcType)]
                  -> [TcIdSigInst]