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.

7 files changed, 164 insertions, 22 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]
diff --git a/testsuite/tests/partial-sigs/should_compile/T14715.stderr b/testsuite/tests/partial-sigs/should_compile/T14715.stderr
index 286ca25671..4d3a668241 100644
--- a/testsuite/tests/partial-sigs/should_compile/T14715.stderr
+++ b/testsuite/tests/partial-sigs/should_compile/T14715.stderr
@@ -3,7 +3,7 @@ T14715.hs:13:53: warning: [-Wpartial-type-signatures (in -Wdefault)]
     • Found extra-constraints wildcard standing for ‘Reduce z zq’
       Where: ‘z’, ‘zq’ are rigid type variables bound by
                the inferred type of
-                 bench_mulPublic :: (LiftOf zq ~ z, Reduce z zq) =>
+                 bench_mulPublic :: (z ~ LiftOf zq, Reduce z zq) =>
                                     Cyc zp -> Cyc z -> IO (zp, zq)
                at T14715.hs:13:27-33
     • In the type signature:
diff --git a/testsuite/tests/partial-sigs/should_compile/T20921.hs b/testsuite/tests/partial-sigs/should_compile/T20921.hs
new file mode 100644
index 0000000000..e6f9fb6df1
--- /dev/null
+++ b/testsuite/tests/partial-sigs/should_compile/T20921.hs
@@ -0,0 +1,53 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE PartialTypeSignatures #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+module PsigBug where
+
+import Data.Kind
+  ( Constraint )
+import GHC.TypeLits
+  ( ErrorMessage(..), TypeError )
+import GHC.TypeNats ( Nat )
+
+type family OK (i :: Nat) :: Constraint where
+  OK 1 = ()
+  OK 2 = ()
+  OK n = TypeError (ShowType n :<>: Text "is not OK")
+
+class C (i :: Nat) where
+  foo :: Int
+
+instance C 1 where
+  foo = 1
+instance C 2 where
+  foo = 2
+
+type family Boo (i :: Nat) :: Nat where
+  Boo i = i
+
+bar :: Int
+bar =
+  let
+    quux :: forall (i :: Nat). ( OK (Boo i), _ ) => Int
+    quux = foo @(Boo i)
+  in quux @1 + quux @2
+
+{-
+From RHS of quux
+  [W] C (Boo i)
+  [W] OK (Boo i)  -- Note [Add signature contexts as wanteds]
+
+Simplifies to
+  [W] C i, OK i
+
+Add back in OK (Boo i) (in decideQuantification),
+and mkMinimalBySCs (which does not eliminate (OK i)
+  (OK (Boo i), OK i, C i)
+
+-}
diff --git a/testsuite/tests/partial-sigs/should_compile/T20921.stderr b/testsuite/tests/partial-sigs/should_compile/T20921.stderr
new file mode 100644
index 0000000000..6c2156cc9f
--- /dev/null
+++ b/testsuite/tests/partial-sigs/should_compile/T20921.stderr
@@ -0,0 +1,20 @@
+
+T20921.hs:37:46: warning: [-Wpartial-type-signatures (in -Wdefault)]
+    • Found extra-constraints wildcard standing for ‘C i’
+      Where: ‘i’ is a rigid type variable bound by
+               the inferred type of quux :: (OK (Boo i), C i) => Int
+               at T20921.hs:37:21
+    • In the type signature:
+        quux :: forall (i :: Nat). (OK (Boo i), _) => Int
+      In the expression:
+        let
+          quux :: forall (i :: Nat). (OK (Boo i), _) => Int
+          quux = foo @(Boo i)
+        in quux @1 + quux @2
+      In an equation for ‘bar’:
+          bar
+            = let
+                quux :: forall (i :: Nat). (OK (Boo i), _) => Int
+                quux = foo @(Boo i)
+              in quux @1 + quux @2
+    • Relevant bindings include bar :: Int (bound at T20921.hs:35:1)
diff --git a/testsuite/tests/partial-sigs/should_compile/all.T b/testsuite/tests/partial-sigs/should_compile/all.T
index 6367aa16f5..e38358f88a 100644
--- a/testsuite/tests/partial-sigs/should_compile/all.T
+++ b/testsuite/tests/partial-sigs/should_compile/all.T
@@ -101,3 +101,4 @@ test('T18008', normal, compile, [''])
 test('T16762d', normal, compile, [''])
 test('T14658', normal, compile, [''])
 test('T18646', normal, compile, [''])
+test('T20921', normal, compile, [''])