Be more careful in chooseInferredQuantifiers

This fixes #22065. We were failing to retain a quantifier that
was mentioned in the kind of another retained quantifier.

Easy to fix.

7 files changed, 146 insertions, 30 deletions

diff --git a/compiler/GHC/Tc/Gen/Bind.hs b/compiler/GHC/Tc/Gen/Bind.hs
index 6d960adddf..d13b3b197c 100644
--- a/compiler/GHC/Tc/Gen/Bind.hs
+++ b/compiler/GHC/Tc/Gen/Bind.hs
@@ -43,6 +43,7 @@ import GHC.Tc.Solver
 import GHC.Tc.Types.Evidence
 import GHC.Tc.Types.Constraint
 import GHC.Core.Predicate
+import GHC.Core.TyCo.Ppr( pprTyVars )
 import GHC.Tc.Gen.HsType
 import GHC.Tc.Gen.Pat
 import GHC.Tc.Utils.TcMType
@@ -59,7 +60,7 @@ import GHC.Types.SourceText
 import GHC.Types.Id
 import GHC.Types.Var as Var
 import GHC.Types.Var.Set
-import GHC.Types.Var.Env( TidyEnv )
+import GHC.Types.Var.Env( TidyEnv, TyVarEnv, mkVarEnv, lookupVarEnv )
 import GHC.Unit.Module
 import GHC.Types.Name
 import GHC.Types.Name.Set
@@ -934,7 +935,8 @@ chooseInferredQuantifiers residual inferred_theta tau_tvs qtvs
        ; let psig_qtvs    = map binderVar psig_qtv_bndrs
              psig_qtv_set = mkVarSet psig_qtvs
              psig_qtv_prs = psig_qtv_nms `zip` psig_qtvs
-
+             psig_bndr_map :: TyVarEnv InvisTVBinder
+             psig_bndr_map = mkVarEnv [ (binderVar tvb, tvb) | tvb <- psig_qtv_bndrs ]
 
             -- Check whether the quantified variables of the
             -- partial signature have been unified together
@@ -950,32 +952,35 @@ chooseInferredQuantifiers residual inferred_theta tau_tvs qtvs
 
        ; annotated_theta      <- zonkTcTypes annotated_theta
        ; (free_tvs, my_theta) <- choose_psig_context psig_qtv_set annotated_theta wcx
+                                 -- NB: free_tvs includes tau_tvs
+
+       ; let (_,final_qtvs) = foldr (choose_qtv psig_bndr_map) (free_tvs, []) qtvs
+                              -- Pulling from qtvs maintains original order
+                              -- NB: qtvs is already in dependency order
 
-       ; let keep_me    = free_tvs `unionVarSet` psig_qtv_set
-             final_qtvs = [ mkTyVarBinder vis tv
-                          | tv <- qtvs -- Pulling from qtvs maintains original order
-                          , tv `elemVarSet` keep_me
-                          , let vis = case lookupVarBndr tv psig_qtv_bndrs of
-                                  Just spec -> spec
-                                  Nothing   -> InferredSpec ]
+       ; traceTc "chooseInferredQuantifiers" $
+         vcat [ text "qtvs" <+> pprTyVars qtvs
+              , text "psig_qtv_bndrs" <+> ppr psig_qtv_bndrs
+              , text "free_tvs" <+> ppr free_tvs
+              , text "final_tvs" <+> ppr final_qtvs ]
 
        ; return (final_qtvs, my_theta) }
   where
-    report_dup_tyvar_tv_err (n1,n2)
-      = addErrTc (TcRnPartialTypeSigTyVarMismatch n1 n2 fn_name hs_ty)
-
-    report_mono_sig_tv_err (n,tv)
-      = addErrTc (TcRnPartialTypeSigBadQuantifier n fn_name m_unif_ty hs_ty)
-      where
-        m_unif_ty = listToMaybe
-                      [ rhs
-                      -- recall that residuals are always implications
-                      | residual_implic <- bagToList $ wc_impl residual
-                      , residual_ct <- bagToList $ wc_simple (ic_wanted residual_implic)
-                      , let residual_pred = ctPred residual_ct
-                      , Just (Nominal, lhs, rhs) <- [ getEqPredTys_maybe residual_pred ]
-                      , Just lhs_tv <- [ tcGetTyVar_maybe lhs ]
-                      , lhs_tv == tv ]
+    choose_qtv :: TyVarEnv InvisTVBinder -> TcTyVar
+             -> (TcTyVarSet, [InvisTVBinder]) -> (TcTyVarSet, [InvisTVBinder])
+    -- Pick which of the original qtvs should be retained
+    -- Keep it if (a) it is mentioned in the body of the type (free_tvs)
+    --            (b) it is a forall'd variable of the partial signature (psig_qtv_bndrs)
+    --            (c) it is mentioned in the kind of a retained qtv (#22065)
+    choose_qtv psig_bndr_map tv (free_tvs, qtvs)
+       | Just psig_bndr <- lookupVarEnv psig_bndr_map tv
+       = (free_tvs', psig_bndr : qtvs)
+       | tv `elemVarSet` free_tvs
+       = (free_tvs', mkTyVarBinder InferredSpec tv : qtvs)
+       | otherwise  -- Do not pick it
+       = (free_tvs, qtvs)
+       where
+         free_tvs' = free_tvs `unionVarSet` tyCoVarsOfType (tyVarKind tv)
 
     choose_psig_context :: VarSet -> TcThetaType -> Maybe TcType
                         -> TcM (VarSet, TcThetaType)
@@ -1019,6 +1024,22 @@ chooseInferredQuantifiers residual inferred_theta tau_tvs qtvs
              -- Return (annotated_theta ++ diff_theta)
              -- See Note [Extra-constraints wildcards]
 
+    report_dup_tyvar_tv_err (n1,n2)
+      = addErrTc (TcRnPartialTypeSigTyVarMismatch n1 n2 fn_name hs_ty)
+
+    report_mono_sig_tv_err (n,tv)
+      = addErrTc (TcRnPartialTypeSigBadQuantifier n fn_name m_unif_ty hs_ty)
+      where
+        m_unif_ty = listToMaybe
+                      [ rhs
+                      -- recall that residuals are always implications
+                      | residual_implic <- bagToList $ wc_impl residual
+                      , residual_ct <- bagToList $ wc_simple (ic_wanted residual_implic)
+                      , let residual_pred = ctPred residual_ct
+                      , Just (Nominal, lhs, rhs) <- [ getEqPredTys_maybe residual_pred ]
+                      , Just lhs_tv <- [ tcGetTyVar_maybe lhs ]
+                      , lhs_tv == tv ]
+
     mk_ctuple preds = mkBoxedTupleTy preds
        -- Hack alert!  See GHC.Tc.Gen.HsType:
        -- Note [Extra-constraint holes in partial type signatures]
diff --git a/compiler/GHC/Types/Var.hs b/compiler/GHC/Types/Var.hs
index cd32cf19c8..dbb739ce1a 100644
--- a/compiler/GHC/Types/Var.hs
+++ b/compiler/GHC/Types/Var.hs
@@ -79,7 +79,7 @@ module GHC.Types.Var (
         mkTyVarBinder, mkTyVarBinders,
         isTyVarBinder,
         tyVarSpecToBinder, tyVarSpecToBinders, tyVarReqToBinder, tyVarReqToBinders,
-        mapVarBndr, mapVarBndrs, lookupVarBndr,
+        mapVarBndr, mapVarBndrs,
 
         -- ** Constructing TyVar's
         mkTyVar, mkTcTyVar,
@@ -696,11 +696,6 @@ mapVarBndr f (Bndr v fl) = Bndr (f v) fl
 mapVarBndrs :: (var -> var') -> [VarBndr var flag] -> [VarBndr var' flag]
 mapVarBndrs f = map (mapVarBndr f)
 
-lookupVarBndr :: Eq var => var -> [VarBndr var flag] -> Maybe flag
-lookupVarBndr var bndrs = lookup var zipped_bndrs
-  where
-    zipped_bndrs = map (\(Bndr v f) -> (v,f)) bndrs
-
 instance Outputable tv => Outputable (VarBndr tv ArgFlag) where
   ppr (Bndr v Required)  = ppr v
   ppr (Bndr v Specified) = char '@' <> ppr v
diff --git a/testsuite/tests/partial-sigs/should_compile/T16152.hs b/testsuite/tests/partial-sigs/should_compile/T16152.hs
new file mode 100644
index 0000000000..a679833dfe
--- /dev/null
+++ b/testsuite/tests/partial-sigs/should_compile/T16152.hs
@@ -0,0 +1,8 @@
+{-# Language PartialTypeSignatures #-}
+{-# Language PolyKinds             #-}
+{-# Language ScopedTypeVariables   #-}
+
+module T16152 where
+
+top :: forall f. _
+top = undefined
diff --git a/testsuite/tests/partial-sigs/should_compile/T16152.stderr b/testsuite/tests/partial-sigs/should_compile/T16152.stderr
new file mode 100644
index 0000000000..554cf04766
--- /dev/null
+++ b/testsuite/tests/partial-sigs/should_compile/T16152.stderr
@@ -0,0 +1,7 @@
+
+T16152.hs:7:18: warning: [-Wpartial-type-signatures (in -Wdefault)]
+    • Found type wildcard ‘_’ standing for ‘w’
+      Where: ‘w’ is a rigid type variable bound by
+               the inferred type of top :: w
+               at T16152.hs:8:1-15
+    • In the type signature: top :: forall f. _
diff --git a/testsuite/tests/partial-sigs/should_compile/T22065.hs b/testsuite/tests/partial-sigs/should_compile/T22065.hs
new file mode 100644
index 0000000000..b74fb76510
--- /dev/null
+++ b/testsuite/tests/partial-sigs/should_compile/T22065.hs
@@ -0,0 +1,30 @@
+{-# Options_GHC -dcore-lint #-}
+{-# Language PartialTypeSignatures #-}
+
+module T22065 where
+
+data Foo where
+  Apply :: (x -> Int) -> x -> Foo
+
+foo :: Foo
+foo = Apply f x :: forall a. _ where
+
+  f :: [_] -> Int
+  f = length @[] @_
+
+  x :: [_]
+  x = mempty @[_]
+
+{-
+Smaller version I used when debuggging
+
+apply :: (x->Int) -> x -> Bool
+apply = apply
+
+foo :: Bool
+foo = apply f x :: forall a. _
+    where
+      f = length @[]
+      x = mempty
+
+-}
diff --git a/testsuite/tests/partial-sigs/should_compile/T22065.stderr b/testsuite/tests/partial-sigs/should_compile/T22065.stderr
new file mode 100644
index 0000000000..42cde6935d
--- /dev/null
+++ b/testsuite/tests/partial-sigs/should_compile/T22065.stderr
@@ -0,0 +1,53 @@
+
+T22065.hs:10:30: warning: [-Wpartial-type-signatures (in -Wdefault)]
+    • Found type wildcard ‘_’ standing for ‘Foo’
+    • In an expression type signature: forall a. _
+      In the expression: Apply f x :: forall a. _
+      In an equation for ‘foo’:
+          foo
+            = Apply f x :: forall a. _
+            where
+                f :: [_] -> Int
+                f = length @[] @_
+                x :: [_]
+                x = mempty @[_]
+    • Relevant bindings include
+        f :: forall {w}. [w] -> Int (bound at T22065.hs:13:3)
+        x :: forall {w}. [w] (bound at T22065.hs:16:3)
+        foo :: Foo (bound at T22065.hs:10:1)
+
+T22065.hs:12:9: warning: [-Wpartial-type-signatures (in -Wdefault)]
+    • Found type wildcard ‘_’ standing for ‘w’
+      Where: ‘w’ is a rigid type variable bound by
+               the inferred type of f :: [w] -> Int
+               at T22065.hs:13:3-19
+    • In the type ‘[_] -> Int’
+      In the type signature: f :: [_] -> Int
+      In an equation for ‘foo’:
+          foo
+            = Apply f x :: forall a. _
+            where
+                f :: [_] -> Int
+                f = length @[] @_
+                x :: [_]
+                x = mempty @[_]
+    • Relevant bindings include
+        x :: forall {w}. [w] (bound at T22065.hs:16:3)
+        foo :: Foo (bound at T22065.hs:10:1)
+
+T22065.hs:15:9: warning: [-Wpartial-type-signatures (in -Wdefault)]
+    • Found type wildcard ‘_’ standing for ‘w’
+      Where: ‘w’ is a rigid type variable bound by
+               the inferred type of x :: [w]
+               at T22065.hs:16:3-17
+    • In the type ‘[_]’
+      In the type signature: x :: [_]
+      In an equation for ‘foo’:
+          foo
+            = Apply f x :: forall a. _
+            where
+                f :: [_] -> Int
+                f = length @[] @_
+                x :: [_]
+                x = mempty @[_]
+    • Relevant bindings include foo :: Foo (bound at T22065.hs:10:1)
diff --git a/testsuite/tests/partial-sigs/should_compile/all.T b/testsuite/tests/partial-sigs/should_compile/all.T
index 3a7e9355a0..935d62231c 100644
--- a/testsuite/tests/partial-sigs/should_compile/all.T
+++ b/testsuite/tests/partial-sigs/should_compile/all.T
@@ -105,3 +105,5 @@ test('T20921', normal, compile, [''])
 test('T21719', normal, compile, [''])
 test('InstanceGivenOverlap3', expect_broken(20076), compile, [''])
 test('T21667', normal, compile, [''])
+test('T22065', normal, compile, [''])
+test('T16152', normal, compile, [''])