Fix nasty bug in the type free-var finder, at last

Consider the type
  forall k. b -> k
where
  b :: k -> Type

Here the 'k' in b's kind must be a different 'k' to the forall k,
because 'b' is free in the expression.  So we must return 'k'
among the free vars returned from tyCoVarsOfType applied that
type.  But we weren't.

This is an outright bug, although we don't have a program that
fails because of it.

It's easy to fix, too: see TyCoRep
  Note [Closing over free variable kinds]

This fix has been in the pipeline for ages because it fell into
the Trac #14880 swamp.  But this patch nails it.

8 files changed, 104 insertions, 22 deletions

diff --git a/compiler/types/TyCoRep.hs b/compiler/types/TyCoRep.hs
index d1c4626238..ef894d8d57 100644
--- a/compiler/types/TyCoRep.hs
+++ b/compiler/types/TyCoRep.hs
@@ -79,7 +79,7 @@ module TyCoRep (
         tyCoVarsOfType, tyCoVarsOfTypeDSet, tyCoVarsOfTypes, tyCoVarsOfTypesDSet,
         tyCoFVsBndr, tyCoFVsOfType, tyCoVarsOfTypeList,
         tyCoFVsOfTypes, tyCoVarsOfTypesList,
-        closeOverKindsDSet, closeOverKindsFV, closeOverKindsList,
+        closeOverKindsDSet, closeOverKindsFV, closeOverKindsList, closeOverKinds,
         coVarsOfType, coVarsOfTypes,
         coVarsOfCo, coVarsOfCos,
         tyCoVarsOfCo, tyCoVarsOfCos,
@@ -87,7 +87,6 @@ module TyCoRep (
         tyCoFVsOfCo, tyCoFVsOfCos,
         tyCoVarsOfCoList, tyCoVarsOfProv,
         almostDevoidCoVarOfCo,
-        closeOverKinds,
         injectiveVarsOfBinder, injectiveVarsOfType,
 
         noFreeVarsOfType, noFreeVarsOfCo,
@@ -1619,6 +1618,77 @@ so we profiled several versions, exploring different implementation strategies.
 
 See Trac #14880.
 
+Note [Closing over free variable kinds]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+tyCoVarsOfType and tyCoFVsOfType, while traversing a type, will also close over
+free variable kinds. In previous GHC versions, this happened naively: whenever
+we would encounter an occurrence of a free type variable, we would close over
+its kind. This, however is wrong for two reasons (see Trac #14880):
+
+1. Efficiency. If we have Proxy (a::k) -> Proxy (a::k) -> Proxy (a::k), then
+   we don't want to have to traverse k more than once.
+
+2. Correctness. Imagine we have forall k. b -> k, where b has
+   kind k, for some k bound in an outer scope. If we look at b's kind inside
+   the forall, we'll collect that k is free and then remove k from the set of
+   free variables. This is plain wrong. We must instead compute that b is free
+   and then conclude that b's kind is free.
+
+An obvious first approach is to move the closing-over-kinds from the
+occurrences of a type variable to after finding the free vars - however, this
+turns out to introduce performance regressions, and isn't even entirely
+correct.
+
+In fact, it isn't even important *when* we close over kinds; what matters is
+that we handle each type var exactly once, and that we do it in the right
+context.
+
+So the next approach we tried was to use the "in-scope set" part of FV or the
+equivalent argument in the accumulator-style `ty_co_vars_of_type` function, to
+say "don't bother with variables we have already closed over". This should work
+fine in theory, but the code is complicated and doesn't perform well.
+
+But there is a simpler way, which is implemented here. Consider the two points
+above:
+
+1. Efficiency: we now have an accumulator, so the second time we encounter 'a',
+   we'll ignore it, certainly not looking at its kind - this is why
+   pre-checking set membership before inserting ends up not only being faster,
+   but also being correct.
+
+2. Correctness: we have an "in-scope set" (I think we should call it it a
+  "bound-var set"), specifying variables that are bound by a forall in the type
+  we are traversing; we simply ignore these variables, certainly not looking at
+  their kind.
+
+So now consider:
+
+    forall k. b -> k
+
+where b :: k->Type is free; but of course, it's a different k! When looking at
+b -> k we'll have k in the bound-var set. So we'll ignore the k. But suppose
+this is our first encounter with b; we want the free vars of its kind. But we
+want to behave as if we took the free vars of its kind at the end; that is,
+with no bound vars in scope.
+
+So the solution is easy. The old code was this:
+
+  ty_co_vars_of_type (TyVarTy v) is acc
+    | v `elemVarSet` is  = acc
+    | v `elemVarSet` acc = acc
+    | otherwise          = ty_co_vars_of_type (tyVarKind v) is (extendVarSet acc v)
+
+Now all we need to do is take the free vars of tyVarKind v *with an empty
+bound-var set*, thus:
+
+ty_co_vars_of_type (TyVarTy v) is acc
+  | v `elemVarSet` is  = acc
+  | v `elemVarSet` acc = acc
+  | otherwise          = ty_co_vars_of_type (tyVarKind v) emptyVarSet (extendVarSet acc v)
+                                                          ^^^^^^^^^^^
+
+And that's it.
+
 -}
 
 tyCoVarsOfType :: Type -> TyCoVarSet
@@ -1632,7 +1702,10 @@ ty_co_vars_of_type :: Type -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
 ty_co_vars_of_type (TyVarTy v) is acc
   | v `elemVarSet` is  = acc
   | v `elemVarSet` acc = acc
-  | otherwise          = ty_co_vars_of_type (varType v) is (extendVarSet acc v)
+  | otherwise          = ty_co_vars_of_type (tyVarKind v)
+                            emptyVarSet  -- See Note [Closing over free variable kinds]
+                            (extendVarSet acc v)
+
 ty_co_vars_of_type (TyConApp _ tys)   is acc = ty_co_vars_of_types tys is acc
 ty_co_vars_of_type (LitTy {})         _  acc = acc
 ty_co_vars_of_type (AppTy fun arg)    is acc = ty_co_vars_of_type fun is (ty_co_vars_of_type arg is acc)
@@ -1691,7 +1764,9 @@ ty_co_vars_of_co_var :: CoVar -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
 ty_co_vars_of_co_var v is acc
   | v `elemVarSet` is  = acc
   | v `elemVarSet` acc = acc
-  | otherwise         = ty_co_vars_of_type (varType v) is (extendVarSet acc v)
+  | otherwise          = ty_co_vars_of_type (varType v)
+                            emptyVarSet  -- See Note [Closing over free variable kinds]
+                            (extendVarSet acc v)
 
 ty_co_vars_of_cos :: [Coercion] -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
 ty_co_vars_of_cos []       _  acc = acc
@@ -1761,14 +1836,20 @@ tyCoVarsOfTypesList tys = fvVarList $ tyCoFVsOfTypes tys
 -- See Note [FV eta expansion] in FV for explanation.
 tyCoFVsOfType :: Type -> FV
 -- See Note [Free variables of types]
-tyCoFVsOfType (TyVarTy v)        a b c = (unitFV v `unionFV` tyCoFVsOfType (varType v)) a b c
-tyCoFVsOfType (TyConApp _ tys)   a b c = tyCoFVsOfTypes tys a b c
-tyCoFVsOfType (LitTy {})         a b c = emptyFV a b c
-tyCoFVsOfType (AppTy fun arg)    a b c = (tyCoFVsOfType fun `unionFV` tyCoFVsOfType arg) a b c
-tyCoFVsOfType (FunTy arg res)    a b c = (tyCoFVsOfType arg `unionFV` tyCoFVsOfType res) a b c
-tyCoFVsOfType (ForAllTy bndr ty) a b c = tyCoFVsBndr bndr (tyCoFVsOfType ty)  a b c
-tyCoFVsOfType (CastTy ty co)     a b c = (tyCoFVsOfType ty `unionFV` tyCoFVsOfCo co) a b c
-tyCoFVsOfType (CoercionTy co)    a b c = tyCoFVsOfCo co a b c
+tyCoFVsOfType (TyVarTy v)        f bound_vars (acc_list, acc_set)
+  | not (f v) = (acc_list, acc_set)
+  | v `elemVarSet` bound_vars = (acc_list, acc_set)
+  | v `elemVarSet` acc_set = (acc_list, acc_set)
+  | otherwise = tyCoFVsOfType (tyVarKind v) f
+                               emptyVarSet   -- See Note [Closing over free variable kinds]
+                               (v:acc_list, extendVarSet acc_set v)
+tyCoFVsOfType (TyConApp _ tys)   f bound_vars acc = tyCoFVsOfTypes tys f bound_vars acc
+tyCoFVsOfType (LitTy {})         f bound_vars acc = emptyFV f bound_vars acc
+tyCoFVsOfType (AppTy fun arg)    f bound_vars acc = (tyCoFVsOfType fun `unionFV` tyCoFVsOfType arg) f bound_vars acc
+tyCoFVsOfType (FunTy arg res)    f bound_vars acc = (tyCoFVsOfType arg `unionFV` tyCoFVsOfType res) f bound_vars acc
+tyCoFVsOfType (ForAllTy bndr ty) f bound_vars acc = tyCoFVsBndr bndr (tyCoFVsOfType ty)  f bound_vars acc
+tyCoFVsOfType (CastTy ty co)     f bound_vars acc = (tyCoFVsOfType ty `unionFV` tyCoFVsOfCo co) f bound_vars acc
+tyCoFVsOfType (CoercionTy co)    f bound_vars acc = tyCoFVsOfCo co f bound_vars acc
 
 tyCoFVsBndr :: TyCoVarBinder -> FV -> FV
 -- Free vars of (forall b. <thing with fvs>)
@@ -1844,6 +1925,7 @@ tyCoFVsOfCos :: [Coercion] -> FV
 tyCoFVsOfCos []       fv_cand in_scope acc = emptyFV fv_cand in_scope acc
 tyCoFVsOfCos (co:cos) fv_cand in_scope acc = (tyCoFVsOfCo co `unionFV` tyCoFVsOfCos cos) fv_cand in_scope acc
 
+
 ------------- Extracting the CoVars of a type or coercion -----------
 
 {-
diff --git a/testsuite/tests/partial-sigs/should_compile/T12844.stderr b/testsuite/tests/partial-sigs/should_compile/T12844.stderr
index 7049818a7c..0e01cd30f3 100644
--- a/testsuite/tests/partial-sigs/should_compile/T12844.stderr
+++ b/testsuite/tests/partial-sigs/should_compile/T12844.stderr
@@ -2,7 +2,7 @@
 T12844.hs:12:9: warning: [-Wpartial-type-signatures (in -Wdefault)]
     • Found type wildcard ‘_’
         standing for ‘(Foo rngs, Head rngs ~ '(r, r'))’
-      Where: ‘rngs’, ‘r’, ‘k’, ‘r'’, ‘k1’
+      Where: ‘rngs’, ‘k’, ‘r’, ‘k1’, ‘r'’
                are rigid type variables bound by
                the inferred type of
                  bar :: (Foo rngs, Head rngs ~ '(r, r')) => FooData rngs
diff --git a/testsuite/tests/partial-sigs/should_compile/T15039a.stderr b/testsuite/tests/partial-sigs/should_compile/T15039a.stderr
index d9c8e1056f..1563a2eb23 100644
--- a/testsuite/tests/partial-sigs/should_compile/T15039a.stderr
+++ b/testsuite/tests/partial-sigs/should_compile/T15039a.stderr
@@ -25,7 +25,7 @@ T15039a.hs:22:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
 
 T15039a.hs:25:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
     • Found type wildcard ‘_’ standing for ‘Dict (a ~~ b)’
-      Where: ‘a’, ‘b’, ‘k’ are rigid type variables bound by
+      Where: ‘a’, ‘k’, ‘b’ are rigid type variables bound by
                the type signature for:
                  ex3 :: forall k a (b :: k). Dict (a ~~ b) -> ()
                at T15039a.hs:24:1-43
diff --git a/testsuite/tests/partial-sigs/should_compile/T15039b.stderr b/testsuite/tests/partial-sigs/should_compile/T15039b.stderr
index 5726c7fa65..21ec20ae40 100644
--- a/testsuite/tests/partial-sigs/should_compile/T15039b.stderr
+++ b/testsuite/tests/partial-sigs/should_compile/T15039b.stderr
@@ -26,7 +26,7 @@ T15039b.hs:22:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
 T15039b.hs:25:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
     • Found type wildcard ‘_’
         standing for ‘Dict ((a :: *) ~~ (b :: k))’
-      Where: ‘a’, ‘b’, ‘k’ are rigid type variables bound by
+      Where: ‘a’, ‘k’, ‘b’ are rigid type variables bound by
                the type signature for:
                  ex3 :: forall k a (b :: k). Dict ((a :: *) ~~ (b :: k)) -> ()
                at T15039b.hs:24:1-43
diff --git a/testsuite/tests/partial-sigs/should_compile/T15039c.stderr b/testsuite/tests/partial-sigs/should_compile/T15039c.stderr
index b5a6b149de..40c126f061 100644
--- a/testsuite/tests/partial-sigs/should_compile/T15039c.stderr
+++ b/testsuite/tests/partial-sigs/should_compile/T15039c.stderr
@@ -25,7 +25,7 @@ T15039c.hs:22:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
 
 T15039c.hs:25:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
     • Found type wildcard ‘_’ standing for ‘Dict (a ~~ b)’
-      Where: ‘a’, ‘b’, ‘k’ are rigid type variables bound by
+      Where: ‘a’, ‘k’, ‘b’ are rigid type variables bound by
                the type signature for:
                  ex3 :: forall k a (b :: k). Dict (a ~~ b) -> ()
                at T15039c.hs:24:1-43
diff --git a/testsuite/tests/partial-sigs/should_compile/T15039d.stderr b/testsuite/tests/partial-sigs/should_compile/T15039d.stderr
index 30b22f8fb2..620199a13c 100644
--- a/testsuite/tests/partial-sigs/should_compile/T15039d.stderr
+++ b/testsuite/tests/partial-sigs/should_compile/T15039d.stderr
@@ -27,7 +27,7 @@ T15039d.hs:22:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
 T15039d.hs:25:14: warning: [-Wpartial-type-signatures (in -Wdefault)]
     • Found type wildcard ‘_’
         standing for ‘Dict ((a :: *) ~~ (b :: k))’
-      Where: ‘a’, ‘b’, ‘k’ are rigid type variables bound by
+      Where: ‘a’, ‘k’, ‘b’ are rigid type variables bound by
                the type signature for:
                  ex3 :: forall k a (b :: k). Dict ((a :: *) ~~ (b :: k)) -> ()
                at T15039d.hs:24:1-43
diff --git a/testsuite/tests/partial-sigs/should_run/T15415.stderr b/testsuite/tests/partial-sigs/should_run/T15415.stderr
index c11d54e322..daa791ffa4 100644
--- a/testsuite/tests/partial-sigs/should_run/T15415.stderr
+++ b/testsuite/tests/partial-sigs/should_run/T15415.stderr
@@ -1,8 +1,8 @@
 
 <interactive>:1:7: error:
     Found type wildcard ‘_’ standing for ‘w0 :: k0’
-    Where: ‘w0’ is an ambiguous type variable
-           ‘k0’ is an ambiguous type variable
+    Where: ‘k0’ is an ambiguous type variable
+           ‘w0’ is an ambiguous type variable
     To use the inferred type, enable PartialTypeSignatures
 
 <interactive>:1:17: error:
@@ -16,8 +16,8 @@
 
 <interactive>:1:7: warning: [-Wpartial-type-signatures (in -Wdefault)]
     Found type wildcard ‘_’ standing for ‘w0 :: k0’
-    Where: ‘w0’ is an ambiguous type variable
-           ‘k0’ is an ambiguous type variable
+    Where: ‘k0’ is an ambiguous type variable
+           ‘w0’ is an ambiguous type variable
 
 <interactive>:1:17: warning: [-Wpartial-type-signatures (in -Wdefault)]
     Found type wildcard ‘_’ standing for ‘* -> *’
diff --git a/testsuite/tests/polykinds/T14265.stderr b/testsuite/tests/polykinds/T14265.stderr
index be6868fdc4..43366fccb7 100644
--- a/testsuite/tests/polykinds/T14265.stderr
+++ b/testsuite/tests/polykinds/T14265.stderr
@@ -1,7 +1,7 @@
 
 T14265.hs:7:12: error:
     • Found type wildcard ‘_’ standing for ‘w :: k’
-      Where: ‘w’, ‘k’ are rigid type variables bound by
+      Where: ‘k’, ‘w’ are rigid type variables bound by
                the inferred type of f :: proxy w -> ()
                at T14265.hs:8:1-8
       To use the inferred type, enable PartialTypeSignatures