Close over kinds at the end (#14880, step 2)

4 files changed, 87 insertions, 108 deletions

diff --git a/compiler/typecheck/TcTyDecls.hs b/compiler/typecheck/TcTyDecls.hs
index f64b9f3e73..ae9ead9a78 100644
--- a/compiler/typecheck/TcTyDecls.hs
+++ b/compiler/typecheck/TcTyDecls.hs
@@ -634,7 +634,8 @@ markNominal lcls ty = let nvars = fvVarList (FV.delFVs lcls $ get_ty_vars ty) in
     get_ty_vars (AppTy t1 t2)     = get_ty_vars t1 `unionFV` get_ty_vars t2
     get_ty_vars (FunTy t1 t2)     = get_ty_vars t1 `unionFV` get_ty_vars t2
     get_ty_vars (TyConApp _ tys)  = mapUnionFV get_ty_vars tys
-    get_ty_vars (ForAllTy tvb ty) = tyCoFVsBndr tvb (get_ty_vars ty)
+    get_ty_vars (ForAllTy tvb ty) = FV.delFV (binderVar tvb) (get_ty_vars ty)
+                                    `unionFV` tyCoFVsOfType (binderKind tvb)
     get_ty_vars (LitTy {})        = emptyFV
     get_ty_vars (CastTy ty _)     = get_ty_vars ty
     get_ty_vars (CoercionTy _)    = emptyFV
diff --git a/compiler/types/Coercion.hs b/compiler/types/Coercion.hs
index 3c81935777..492fad18a3 100644
--- a/compiler/types/Coercion.hs
+++ b/compiler/types/Coercion.hs
@@ -68,7 +68,7 @@ module Coercion (
 
         -- ** Free variables
         tyCoVarsOfCo, tyCoVarsOfCos, coVarsOfCo,
-        tyCoFVsOfCo, tyCoFVsOfCos, tyCoVarsOfCoDSet,
+        tyCoFVsOfCo, tyCoVarsOfCoDSet,
         coercionSize,
 
         -- ** Substitution
diff --git a/compiler/types/TyCoRep.hs b/compiler/types/TyCoRep.hs
index 7f470ea7b8..d2cb070313 100644
--- a/compiler/types/TyCoRep.hs
+++ b/compiler/types/TyCoRep.hs
@@ -73,14 +73,14 @@ module TyCoRep (
 
         -- * Free variables
         tyCoVarsOfType, tyCoVarsOfTypeDSet, tyCoVarsOfTypes, tyCoVarsOfTypesDSet,
-        tyCoFVsBndr, tyCoFVsOfType, tyCoVarsOfTypeList,
+        tyCoFVsOfType, tyCoVarsOfTypeList,
         tyCoFVsOfTypes, tyCoVarsOfTypesList,
         closeOverKindsDSet, closeOverKindsFV, closeOverKindsList,
         coVarsOfType, coVarsOfTypes,
         coVarsOfCo, coVarsOfCos,
         tyCoVarsOfCo, tyCoVarsOfCos,
         tyCoVarsOfCoDSet,
-        tyCoFVsOfCo, tyCoFVsOfCos,
+        tyCoFVsOfCo,
         tyCoVarsOfCoList, tyCoVarsOfProv,
         closeOverKinds,
         injectiveVarsOfBinder, injectiveVarsOfType,
@@ -1641,7 +1641,9 @@ mkTyCoInScopeSet tys cos
 
 tyCoVarsOfType :: Type -> TyCoVarSet
 -- See Note [Free variables of types]
-tyCoVarsOfType ty = ty_co_vars_of_type ty emptyVarSet emptyVarSet
+tyCoVarsOfType ty =
+  let unclosed = ty_co_vars_of_type ty emptyVarSet emptyVarSet
+  in closeOverKinds unclosed
 
 tyCoVarsOfTypes :: [Type] -> TyCoVarSet
 tyCoVarsOfTypes tys = ty_co_vars_of_types tys emptyVarSet emptyVarSet
@@ -1650,7 +1652,8 @@ 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 (tyVarKind v) is (extendVarSet acc v)
+  -- | otherwise          = ty_co_vars_of_type (tyVarKind v) is (extendVarSet acc v)
+  | otherwise          = 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)
@@ -1780,26 +1783,33 @@ tyCoVarsOfTypesList tys = fvVarList $ tyCoFVsOfTypes tys
 --
 -- Eta-expanded because that makes it run faster (apparently)
 -- See Note [FV eta expansion] in FV for explanation.
-tyCoFVsOfType :: Type -> FV
+fvs_of_type :: Type -> FV
 -- See Note [Free variables of types]
-tyCoFVsOfType (TyVarTy v)        a b c = (unitFV v `unionFV` tyCoFVsOfType (tyVarKind 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
-
-tyCoFVsBndr :: TyVarBinder -> FV -> FV
+fvs_of_type (TyVarTy v)        a b c = unitFV v a b c
+fvs_of_type (TyConApp _ tys)   a b c = fvs_of_types tys a b c
+fvs_of_type (LitTy {})         a b c = emptyFV a b c
+fvs_of_type (AppTy fun arg)    a b c = (fvs_of_type fun `unionFV` fvs_of_type arg) a b c
+fvs_of_type (FunTy arg res)    a b c = (fvs_of_type arg `unionFV` fvs_of_type res) a b c
+fvs_of_type (ForAllTy bndr ty) a b c = fvs_of_binder bndr (fvs_of_type ty)  a b c
+fvs_of_type (CastTy ty co)     a b c = (fvs_of_type ty `unionFV` fvs_of_co co) a b c
+fvs_of_type (CoercionTy co)    a b c = fvs_of_co co a b c
+
+fvs_of_binder :: TyVarBinder -> FV -> FV
 -- Free vars of (forall b. <thing with fvs>)
-tyCoFVsBndr (TvBndr tv _) fvs = (delFV tv fvs)
-                                `unionFV` tyCoFVsOfType (tyVarKind tv)
+fvs_of_binder (TvBndr tv _) fvs = (delFV tv fvs)
+                                `unionFV` fvs_of_type (tyVarKind tv)
 
-tyCoFVsOfTypes :: [Type] -> FV
+fvs_of_types :: [Type] -> FV
 -- See Note [Free variables of types]
-tyCoFVsOfTypes (ty:tys) fv_cand in_scope acc = (tyCoFVsOfType ty `unionFV` tyCoFVsOfTypes tys) fv_cand in_scope acc
-tyCoFVsOfTypes []       fv_cand in_scope acc = emptyFV fv_cand in_scope acc
+fvs_of_types = mapUnionFV fvs_of_type
+
+tyCoFVsOfType :: Type -> FV
+tyCoFVsOfType ty =
+  let unclosed = fvs_of_type ty
+  in closeOverKindsFV (fvVarList unclosed)
+
+tyCoFVsOfTypes :: [Type] -> FV
+tyCoFVsOfTypes = mapUnionFV tyCoFVsOfType
 
 -- | Get a deterministic set of the vars free in a coercion
 tyCoVarsOfCoDSet :: Coercion -> DTyCoVarSet
@@ -1810,9 +1820,9 @@ tyCoVarsOfCoList :: Coercion -> [TyCoVar]
 -- See Note [Free variables of types]
 tyCoVarsOfCoList co = fvVarList $ tyCoFVsOfCo co
 
-tyCoFVsOfMCo :: MCoercion -> FV
-tyCoFVsOfMCo MRefl    = emptyFV
-tyCoFVsOfMCo (MCo co) = tyCoFVsOfCo co
+fvs_of_mco :: MCoercion -> FV
+fvs_of_mco MRefl    = emptyFV
+fvs_of_mco (MCo co) = fvs_of_co co
 
 tyCoVarsOfCosSet :: CoVarEnv Coercion -> TyCoVarSet
 tyCoVarsOfCosSet cos = tyCoVarsOfCos $ nonDetEltsUFM cos
@@ -1820,101 +1830,68 @@ tyCoVarsOfCosSet cos = tyCoVarsOfCos $ nonDetEltsUFM cos
   -- ordering by returning a set
 
 tyCoFVsOfCo :: Coercion -> FV
+tyCoFVsOfCo co =
+  let unclosed = fvs_of_co co
+  in closeOverKindsFV (fvVarList unclosed)
+
 -- Extracts type and coercion variables from a coercion
 -- See Note [Free variables of types]
-tyCoFVsOfCo (Refl ty) fv_cand in_scope acc
-  = tyCoFVsOfType ty fv_cand in_scope acc
-tyCoFVsOfCo (GRefl _ ty mco) fv_cand in_scope acc
-  = (tyCoFVsOfType ty `unionFV` tyCoFVsOfMCo mco) fv_cand in_scope acc
-tyCoFVsOfCo (TyConAppCo _ _ cos) fv_cand in_scope acc = tyCoFVsOfCos cos fv_cand in_scope acc
-tyCoFVsOfCo (AppCo co arg) fv_cand in_scope acc
-  = (tyCoFVsOfCo co `unionFV` tyCoFVsOfCo arg) fv_cand in_scope acc
-tyCoFVsOfCo (ForAllCo tv kind_co co) fv_cand in_scope acc
-  = (delFV tv (tyCoFVsOfCo co) `unionFV` tyCoFVsOfCo kind_co) fv_cand in_scope acc
-tyCoFVsOfCo (FunCo _ co1 co2)    fv_cand in_scope acc
-  = (tyCoFVsOfCo co1 `unionFV` tyCoFVsOfCo co2) fv_cand in_scope acc
-tyCoFVsOfCo (CoVarCo v) fv_cand in_scope acc
-  = tyCoFVsOfCoVar v fv_cand in_scope acc
-tyCoFVsOfCo (HoleCo h) fv_cand in_scope acc
-  = tyCoFVsOfCoVar (coHoleCoVar h) fv_cand in_scope acc
+fvs_of_co :: Coercion -> FV
+fvs_of_co (Refl ty) fv_cand in_scope acc
+  = fvs_of_type ty fv_cand in_scope acc
+fvs_of_co (GRefl _ ty mco) fv_cand in_scope acc
+  = (fvs_of_type ty `unionFV` fvs_of_mco mco) fv_cand in_scope acc
+fvs_of_co (TyConAppCo _ _ cos) fv_cand in_scope acc = fvs_of_cos cos fv_cand in_scope acc
+fvs_of_co (AppCo co arg) fv_cand in_scope acc
+  = (fvs_of_co co `unionFV` fvs_of_co arg) fv_cand in_scope acc
+fvs_of_co (ForAllCo tv kind_co co) fv_cand in_scope acc
+  = (delFV tv (fvs_of_co co) `unionFV` fvs_of_co kind_co) fv_cand in_scope acc
+fvs_of_co (FunCo _ co1 co2)    fv_cand in_scope acc
+  = (fvs_of_co co1 `unionFV` fvs_of_co co2) fv_cand in_scope acc
+fvs_of_co (CoVarCo v) fv_cand in_scope acc
+  = fvs_of_co_var v fv_cand in_scope acc
+fvs_of_co (HoleCo h) fv_cand in_scope acc
+  = fvs_of_co_var (coHoleCoVar h) fv_cand in_scope acc
     -- See Note [CoercionHoles and coercion free variables]
-tyCoFVsOfCo (AxiomInstCo _ _ cos) fv_cand in_scope acc = tyCoFVsOfCos cos fv_cand in_scope acc
-tyCoFVsOfCo (UnivCo p _ t1 t2) fv_cand in_scope acc
-  = (tyCoFVsOfProv p `unionFV` tyCoFVsOfType t1
-                     `unionFV` tyCoFVsOfType t2) fv_cand in_scope acc
-tyCoFVsOfCo (SymCo co)          fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
-tyCoFVsOfCo (TransCo co1 co2)   fv_cand in_scope acc = (tyCoFVsOfCo co1 `unionFV` tyCoFVsOfCo co2) fv_cand in_scope acc
-tyCoFVsOfCo (NthCo _ _ co)      fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
-tyCoFVsOfCo (LRCo _ co)         fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
-tyCoFVsOfCo (InstCo co arg)     fv_cand in_scope acc = (tyCoFVsOfCo co `unionFV` tyCoFVsOfCo arg) fv_cand in_scope acc
-tyCoFVsOfCo (KindCo co)         fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
-tyCoFVsOfCo (SubCo co)          fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
-tyCoFVsOfCo (AxiomRuleCo _ cs)  fv_cand in_scope acc = tyCoFVsOfCos cs fv_cand in_scope acc
-
-tyCoFVsOfCoVar :: CoVar -> FV
-tyCoFVsOfCoVar v fv_cand in_scope acc
+fvs_of_co (AxiomInstCo _ _ cos) fv_cand in_scope acc = fvs_of_cos cos fv_cand in_scope acc
+fvs_of_co (UnivCo p _ t1 t2) fv_cand in_scope acc
+  = (fvs_of_prov p `unionFV` fvs_of_type t1
+                     `unionFV` fvs_of_type t2) fv_cand in_scope acc
+fvs_of_co (SymCo co)          fv_cand in_scope acc = fvs_of_co co fv_cand in_scope acc
+fvs_of_co (TransCo co1 co2)   fv_cand in_scope acc = (fvs_of_co co1 `unionFV` fvs_of_co co2) fv_cand in_scope acc
+fvs_of_co (NthCo _ _ co)      fv_cand in_scope acc = fvs_of_co co fv_cand in_scope acc
+fvs_of_co (LRCo _ co)         fv_cand in_scope acc = fvs_of_co co fv_cand in_scope acc
+fvs_of_co (InstCo co arg)     fv_cand in_scope acc = (fvs_of_co co `unionFV` fvs_of_co arg) fv_cand in_scope acc
+fvs_of_co (KindCo co)         fv_cand in_scope acc = fvs_of_co co fv_cand in_scope acc
+fvs_of_co (SubCo co)          fv_cand in_scope acc = fvs_of_co co fv_cand in_scope acc
+fvs_of_co (AxiomRuleCo _ cs)  fv_cand in_scope acc = fvs_of_cos cs fv_cand in_scope acc
+
+fvs_of_co_var :: CoVar -> FV
+fvs_of_co_var v fv_cand in_scope acc
   = (unitFV v `unionFV` tyCoFVsOfType (varType v)) fv_cand in_scope acc
 
-tyCoFVsOfProv :: UnivCoProvenance -> FV
-tyCoFVsOfProv UnsafeCoerceProv    fv_cand in_scope acc = emptyFV fv_cand in_scope acc
-tyCoFVsOfProv (PhantomProv co)    fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
-tyCoFVsOfProv (ProofIrrelProv co) fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
-tyCoFVsOfProv (PluginProv _)      fv_cand in_scope acc = emptyFV fv_cand in_scope acc
+fvs_of_prov :: UnivCoProvenance -> FV
+fvs_of_prov UnsafeCoerceProv    fv_cand in_scope acc = emptyFV fv_cand in_scope acc
+fvs_of_prov (PhantomProv co)    fv_cand in_scope acc = fvs_of_co co fv_cand in_scope acc
+fvs_of_prov (ProofIrrelProv co) fv_cand in_scope acc = fvs_of_co co fv_cand in_scope acc
+fvs_of_prov (PluginProv _)      fv_cand in_scope acc = emptyFV fv_cand in_scope acc
 
-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
+fvs_of_cos :: [Coercion] -> FV
+fvs_of_cos = mapUnionFV fvs_of_co
 
 ------------- Extracting the CoVars of a type or coercion -----------
 
-getCoVarSet :: FV -> CoVarSet
-getCoVarSet fv = snd (fv isCoVar emptyVarSet ([], emptyVarSet))
-
 coVarsOfType :: Type -> CoVarSet
-coVarsOfType ty = getCoVarSet (tyCoFVsOfType ty)
+coVarsOfType = filterVarSet isCoVar . tyCoVarsOfType
 
 coVarsOfTypes :: [Type] -> TyCoVarSet
-coVarsOfTypes tys = getCoVarSet (tyCoFVsOfTypes tys)
-
--- coVarsOfMCo :: MCoercion -> CoVarSet
--- coVarsOfMCo MRefl    = emptyVarSet
--- coVarsOfMCo (MCo co) = coVarsOfCo co
+coVarsOfTypes = mapUnionVarSet coVarsOfType
 
 coVarsOfCo :: Coercion -> CoVarSet
-coVarsOfCo co = getCoVarSet (tyCoFVsOfCo co)
--- -- Extract *coercion* variables only.  Tiresome to repeat the code, but easy.
--- coVarsOfCo (Refl ty)             = coVarsOfType ty
--- coVarsOfCo (GRefl _ ty co)       = coVarsOfType ty `unionVarSet` coVarsOfMCo co
--- coVarsOfCo (TyConAppCo _ _ args) = coVarsOfCos args
--- coVarsOfCo (AppCo co arg)      = coVarsOfCo co `unionVarSet` coVarsOfCo arg
--- coVarsOfCo (ForAllCo tv kind_co co)
---   = coVarsOfCo co `delVarSet` tv `unionVarSet` coVarsOfCo kind_co
--- coVarsOfCo (FunCo _ co1 co2)    = coVarsOfCo co1 `unionVarSet` coVarsOfCo co2
--- coVarsOfCo (CoVarCo v)          = coVarsOfCoVar v
--- coVarsOfCo (HoleCo h)           = coVarsOfCoVar (coHoleCoVar h)
---                                   -- See Note [CoercionHoles and coercion free variables]
--- coVarsOfCo (AxiomInstCo _ _ as) = coVarsOfCos as
--- coVarsOfCo (UnivCo p _ t1 t2)   = coVarsOfProv p `unionVarSet` coVarsOfTypes [t1, t2]
--- coVarsOfCo (SymCo co)           = coVarsOfCo co
--- coVarsOfCo (TransCo co1 co2)    = coVarsOfCo co1 `unionVarSet` coVarsOfCo co2
--- coVarsOfCo (NthCo _ _ co)       = coVarsOfCo co
--- coVarsOfCo (LRCo _ co)          = coVarsOfCo co
--- coVarsOfCo (InstCo co arg)      = coVarsOfCo co `unionVarSet` coVarsOfCo arg
--- coVarsOfCo (KindCo co)          = coVarsOfCo co
--- coVarsOfCo (SubCo co)           = coVarsOfCo co
--- coVarsOfCo (AxiomRuleCo _ cs)   = coVarsOfCos cs
-
--- coVarsOfCoVar :: CoVar -> CoVarSet
--- coVarsOfCoVar v = unitVarSet v `unionVarSet` coVarsOfType (varType v)
---
--- coVarsOfProv :: UnivCoProvenance -> CoVarSet
--- coVarsOfProv UnsafeCoerceProv    = emptyVarSet
--- coVarsOfProv (PhantomProv co)    = coVarsOfCo co
--- coVarsOfProv (ProofIrrelProv co) = coVarsOfCo co
--- coVarsOfProv (PluginProv _)      = emptyVarSet
+coVarsOfCo = filterVarSet isCoVar . tyCoVarsOfCo
 
 coVarsOfCos :: [Coercion] -> CoVarSet
-coVarsOfCos cos = getCoVarSet (tyCoFVsOfCos cos)
+coVarsOfCos = mapUnionVarSet coVarsOfCo
 
 ------------- Closing over kinds -----------------
 
@@ -1959,11 +1936,12 @@ injectiveVarsOfBinder (TvBndr tv vis) =
 -- (See @Note [Kind annotations on TyConApps]@ in "TcSplice" for an explanation
 -- of what an injective position is.)
 injectiveVarsOfType :: Type -> FV
-injectiveVarsOfType = go
+injectiveVarsOfType orig_ty =
+  closeOverKindsFV (fvVarList $ go orig_ty)
   where
     go ty                | Just ty' <- coreView ty
                          = go ty'
-    go (TyVarTy v)       = unitFV v `unionFV` go (tyVarKind v)
+    go (TyVarTy v)       = unitFV v
     go (AppTy f a)       = go f `unionFV` go a
     go (FunTy ty1 ty2)   = go ty1 `unionFV` go ty2
     go (TyConApp tc tys) =
@@ -1973,7 +1951,7 @@ injectiveVarsOfType = go
                          filterByList (inj ++ repeat True) tys
                          -- Oversaturated arguments to a tycon are
                          -- always injective, hence the repeat True
-    go (ForAllTy tvb ty) = tyCoFVsBndr tvb $ go (tyVarKind (binderVar tvb))
+    go (ForAllTy tvb ty) = fvs_of_binder tvb $ go (tyVarKind (binderVar tvb))
                                              `unionFV` go ty
     go LitTy{}           = emptyFV
     go (CastTy ty _)     = go ty
diff --git a/compiler/types/Type.hs b/compiler/types/Type.hs
index 180af3862c..30775dc2cf 100644
--- a/compiler/types/Type.hs
+++ b/compiler/types/Type.hs
@@ -128,7 +128,7 @@ module Type (
         liftedTypeKind,
 
         -- * Type free variables
-        tyCoFVsOfType, tyCoFVsBndr,
+        tyCoFVsOfType,
         tyCoVarsOfType, tyCoVarsOfTypes,
         tyCoVarsOfTypeDSet,
         coVarsOfType,