Literally use Simon's code for tcvs_of_...wip/T14880-reengineered

1 files changed, 49 insertions, 62 deletions

diff --git a/compiler/types/TyCoRep.hs b/compiler/types/TyCoRep.hs
index cfd50b57c1..a467cd748a 100644
--- a/compiler/types/TyCoRep.hs
+++ b/compiler/types/TyCoRep.hs
@@ -1528,85 +1528,72 @@ closeOverKinds tcvs = mapUnionVarSetSet (tyCoVarsOfType . tyVarKind) tcvs `union
 -- Explicitly note that these sets are not closed over kinds
 type TyCoVarSetNotClosed = TyCoVarSet
 
-type TCFV = TyCoVarSet -> TyCoVarSet -> TyCoVarSet
-
-mapUnionTCFV :: (a -> TCFV) -> [a] -> TCFV
-mapUnionTCFV f xs is acc = go xs acc
-  where
-    go [] acc = acc
-    go (x:xs) acc = f x is $ go xs acc
-{-#INLINE mapUnionTCFV #-}
-
-unionTCFV :: TCFV -> TCFV -> TCFV
-unionTCFV f g is acc = g is $ f is acc
-{-#INLINE unionTCFV #-}
-
-emptyTCFV :: TCFV
-emptyTCFV _ acc = acc
-{-#INLINE emptyTCFV #-}
-
--- These functions produce a non-deterministic set. No point in going via FV (which maintains
--- determinism info) and then drop the determinism. This is boring boiler plate code, but this
--- is measurably faster than going via FV.
--- TODO: Add Note to explain why accumulator style is needed here, and why
--- we won't use FV itself.
-tcvs_of_type :: Type -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed
-tcvs_of_type (TyVarTy v)                 is acc
-  | v `elemVarSet` is = acc
+tcvs_of_type :: Type -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
+tcvs_of_type (TyVarTy v) is acc
+  | v `elemVarSet` is  = acc
   | v `elemVarSet` acc = acc
-  | otherwise = tcvs_of_type (tyVarKind v) is (extendVarSet acc v)
-tcvs_of_type (TyConApp _ tys)            is acc = tcvs_of_types tys is acc
-tcvs_of_type (LitTy {})                  _  acc = acc
-tcvs_of_type (AppTy fun arg)             is acc = tcvs_of_type fun is $ tcvs_of_type arg is acc
-tcvs_of_type (FunTy arg res)             is acc = tcvs_of_type arg is $ tcvs_of_type res is acc
+  | otherwise          = tcvs_of_type (tyVarKind v) is (extendVarSet acc v)
+tcvs_of_type (TyConApp _ tys)   is acc = tcvs_of_types tys is acc
+tcvs_of_type (LitTy {})         _  acc = acc
+tcvs_of_type (AppTy fun arg)    is acc = tcvs_of_type fun is (tcvs_of_type arg is acc)
+tcvs_of_type (FunTy arg res)    is acc = tcvs_of_type arg is (tcvs_of_type res is acc)
 tcvs_of_type (ForAllTy (TvBndr tv _) ty) is acc = tcvs_of_type (tyVarKind tv) is $
-                                                  tcvs_of_type ty (extendVarSet is tv) acc
-tcvs_of_type (CastTy ty co)              is acc = tcvs_of_type ty is (tcvs_of_co co is acc)
-tcvs_of_type (CoercionTy co)             is acc = tcvs_of_co co is acc
-
-tcvs_of_types :: [Type] -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed
-tcvs_of_types = mapUnionTCFV tcvs_of_type
--- tcvs_of_types [] _ acc = acc
--- tcvs_of_types (ty:tys) is acc = tcvs_of_type ty is $ tcvs_of_types tys is acc
-
-tcvs_of_co :: Coercion -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed
-tcvs_of_co (Refl _ ty)              is acc = tcvs_of_type ty is acc
-tcvs_of_co (TyConAppCo _ _ cos)     is acc = tcvs_of_cos cos is acc
-tcvs_of_co (AppCo co arg)           is acc = tcvs_of_co co is $ tcvs_of_co arg is acc
-tcvs_of_co (ForAllCo tv kind_co co) is acc = tcvs_of_co kind_co is $ tcvs_of_co co (extendVarSet is tv) acc
-tcvs_of_co (FunCo _ co1 co2)        is acc = tcvs_of_co co1 is $ tcvs_of_co co2 is acc
-tcvs_of_co (CoVarCo v)              is acc = tcvs_of_co_var v is acc
-tcvs_of_co (HoleCo h)               is acc = tcvs_of_co_var (coHoleCoVar h) is acc
+                                                        tcvs_of_type ty (extendVarSet is tv) acc
+tcvs_of_type (CastTy ty co)     is acc = tcvs_of_type ty is (tcvs_of_co co is acc)
+tcvs_of_type (CoercionTy co)    is acc = tcvs_of_co co is acc
+
+tcvs_of_types :: [Type] -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
+tcvs_of_types []       _  acc = acc
+tcvs_of_types (ty:tys) is acc = tcvs_of_type ty is (tcvs_of_types tys is acc)
+
+tcvs_of_co :: Coercion -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
+tcvs_of_co (Refl _ ty)          is acc = tcvs_of_type ty is acc
+tcvs_of_co (TyConAppCo _ _ cos) is acc = tcvs_of_cos cos is acc
+tcvs_of_co (AppCo co arg)       is acc = tcvs_of_co co is $
+                                               tcvs_of_co arg is acc
+tcvs_of_co (ForAllCo tv kind_co co) is acc = tcvs_of_co kind_co is $
+                                                   tcvs_of_co co (extendVarSet is tv) acc
+tcvs_of_co (FunCo _ co1 co2)    is acc = tcvs_of_co co1 is $
+                                               tcvs_of_co co2 is acc
+tcvs_of_co (CoVarCo v)          is acc = tcvs_of_co_var v is acc
+tcvs_of_co (HoleCo h)           is acc = tcvs_of_co_var (coHoleCoVar h) is acc
     -- See Note [CoercionHoles and coercion free variables]
-tcvs_of_co (AxiomInstCo _ _ cos)    is acc = tcvs_of_cos cos is acc
-tcvs_of_co (UnivCo p _ t1 t2)       is acc = (tcvs_of_prov p
-                                                `unionTCFV` tcvs_of_type t1
-                                                `unionTCFV` tcvs_of_type t2)
-                                             is acc
+tcvs_of_co (AxiomInstCo _ _ cos) is acc = tcvs_of_cos cos is acc
+tcvs_of_co (UnivCo p _ t1 t2)    is acc = tcvs_of_prov p is $
+                                                tcvs_of_type t1 is $
+                                                tcvs_of_type t2 is acc
 tcvs_of_co (SymCo co)          is acc = tcvs_of_co co is acc
-tcvs_of_co (TransCo co1 co2)   is acc = (tcvs_of_co co1 `unionTCFV` tcvs_of_co co2) is acc
+tcvs_of_co (TransCo co1 co2)   is acc = tcvs_of_co co1 is $
+                                              tcvs_of_co co2 is acc
 tcvs_of_co (NthCo _ _ co)      is acc = tcvs_of_co co is acc
 tcvs_of_co (LRCo _ co)         is acc = tcvs_of_co co is acc
-tcvs_of_co (InstCo co arg)     is acc = (tcvs_of_co co `unionTCFV` tcvs_of_co arg) is acc
-tcvs_of_co (CoherenceCo c1 c2) is acc = (tcvs_of_co c1 `unionTCFV` tcvs_of_co c2) is acc
+tcvs_of_co (InstCo co arg)     is acc = tcvs_of_co co is $
+                                              tcvs_of_co arg is acc
+tcvs_of_co (CoherenceCo c1 c2) is acc = tcvs_of_co c1 is $
+                                              tcvs_of_co c2 is acc
 tcvs_of_co (KindCo co)         is acc = tcvs_of_co co is acc
 tcvs_of_co (SubCo co)          is acc = tcvs_of_co co is acc
 tcvs_of_co (AxiomRuleCo _ cs)  is acc = tcvs_of_cos cs is acc
 
-tcvs_of_co_var :: CoVar -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed
+tcvs_of_co_var :: CoVar -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
 tcvs_of_co_var v is acc
   | v `elemVarSet` is = acc
   | v `elemVarSet` acc = acc
-  | otherwise = tcvs_of_type (varType v) is (extendVarSet acc v)
+  | otherwise         = tcvs_of_type (varType v) is (extendVarSet acc v)
 
-tcvs_of_cos :: [Coercion] -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed
-tcvs_of_cos = mapUnionTCFV tcvs_of_co
+tcvs_of_cos :: [Coercion] -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
+tcvs_of_cos []       _  acc = acc
+tcvs_of_cos (co:cos) is acc = tcvs_of_co co is (tcvs_of_cos cos is acc)
 
-tcvs_of_prov :: UnivCoProvenance -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed -> TyCoVarSetNotClosed
-tcvs_of_prov UnsafeCoerceProv    is acc = emptyTCFV is acc
+-- tyCoVarsOfProv :: UnivCoProvenance -> TyCoVarSet
+-- tyCoVarsOfProv prov = tcvs_of_prov prov emptyVarSet emptyVarSet
+
+tcvs_of_prov :: UnivCoProvenance -> TyCoVarSet -> TyCoVarSet -> TyCoVarSet
 tcvs_of_prov (PhantomProv co)    is acc = tcvs_of_co co is acc
 tcvs_of_prov (ProofIrrelProv co) is acc = tcvs_of_co co is acc
-tcvs_of_prov (PluginProv _)      is acc = emptyTCFV is acc
+tcvs_of_prov UnsafeCoerceProv    _  acc = acc
+tcvs_of_prov (PluginProv _)      _  acc = acc
+
 
 -- | `tyCoFVsOfType` that returns free variables of a type in a deterministic
 -- set. For explanation of why using `VarSet` is not deterministic see