Rewrite tyCoVarsOf... in terms of new nondet FV (NDFV)wip/T14880-nondet-fv

NDFV is a non-deterministic variation on FV - essentially, FV without
the list part, but with all the other characteristics. This is to figure
out whether the list part is the cause of performance issues in Trac #14880.

3 files changed, 278 insertions, 1 deletions

diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index 2a4d9755bd..4cd819c10a 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -506,6 +506,7 @@ Library
         Fingerprint
         FiniteMap
         FV
+        NDFV
         GraphBase
         GraphColor
         GraphOps
diff --git a/compiler/types/TyCoRep.hs b/compiler/types/TyCoRep.hs
index ec4607a2fb..a5651a3841 100644
--- a/compiler/types/TyCoRep.hs
+++ b/compiler/types/TyCoRep.hs
@@ -164,6 +164,7 @@ import TyCon
 import Class
 import CoAxiom
 import FV
+import NDFV
 
 -- others
 import BasicTypes ( LeftOrRight(..), TyPrec(..), maybeParen, pickLR )
@@ -1484,7 +1485,7 @@ so, so it's easiest to do it here.
 -- synonym.
 tyCoVarsOfType :: Type -> TyCoVarSet
 -- See Note [Free variables of types]
-tyCoVarsOfType ty = fvVarSet $ tyCoFVsOfType ty
+tyCoVarsOfType ty = ndfvVarSet $ tyCoNDFVsOfType ty
 
 -- | `tyCoFVsOfType` that returns free variables of a type in a deterministic
 -- set. For explanation of why using `VarSet` is not deterministic see
@@ -1525,6 +1526,31 @@ tyCoFVsBndr :: TyVarBinder -> FV -> FV
 tyCoFVsBndr (TvBndr tv _) fvs = (delFV tv fvs)
                                 `unionFV` tyCoFVsOfType (tyVarKind tv)
 
+-- | The worker for `tyCoNDFVsOfType` and `tyCoNDFVsOfTypeList`.
+-- The previous implementation used `unionVarSet` which is O(n+m) and can
+-- make the function quadratic.
+-- It's exported, so that it can be composed with
+-- other functions that compute free variables.
+-- See Note [NDFV naming conventions] in NDFV.
+--
+-- Eta-expanded because that makes it run faster (apparently)
+-- See Note [NDFV eta expansion] in NDFV for explanation.
+tyCoNDFVsOfType :: Type -> NDFV
+-- See Note [Free variables of types]
+tyCoNDFVsOfType (TyVarTy v)        a b c = (unitNDFV v `unionNDFV` tyCoNDFVsOfType (tyVarKind v)) a b c
+tyCoNDFVsOfType (TyConApp _ tys)   a b c = tyCoNDFVsOfTypes tys a b c
+tyCoNDFVsOfType (LitTy {})         a b c = emptyNDFV a b c
+tyCoNDFVsOfType (AppTy fun arg)    a b c = (tyCoNDFVsOfType fun `unionNDFV` tyCoNDFVsOfType arg) a b c
+tyCoNDFVsOfType (FunTy arg res)    a b c = (tyCoNDFVsOfType arg `unionNDFV` tyCoNDFVsOfType res) a b c
+tyCoNDFVsOfType (ForAllTy bndr ty) a b c = tyCoNDFVsBndr bndr (tyCoNDFVsOfType ty)  a b c
+tyCoNDFVsOfType (CastTy ty co)     a b c = (tyCoNDFVsOfType ty `unionNDFV` tyCoNDFVsOfCo co) a b c
+tyCoNDFVsOfType (CoercionTy co)    a b c = tyCoNDFVsOfCo co a b c
+
+tyCoNDFVsBndr :: TyVarBinder -> NDFV -> NDFV
+-- Free vars of (forall b. <thing with fvs>)
+tyCoNDFVsBndr (TvBndr tv _) fvs = (delNDFV tv fvs)
+                                `unionNDFV` tyCoNDFVsOfType (tyVarKind tv)
+
 -- | Returns free variables of types, including kind variables as
 -- a non-deterministic set. For type synonyms it does /not/ expand the
 -- synonym.
@@ -1560,6 +1586,11 @@ tyCoFVsOfTypes :: [Type] -> FV
 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
 
+tyCoNDFVsOfTypes :: [Type] -> NDFV
+-- See Note [Free variables of types]
+tyCoNDFVsOfTypes (ty:tys) fv_cand in_scope acc = (tyCoNDFVsOfType ty `unionNDFV` tyCoNDFVsOfTypes tys) fv_cand in_scope acc
+tyCoNDFVsOfTypes []       fv_cand in_scope acc = emptyNDFV fv_cand in_scope acc
+
 tyCoVarsOfCo :: Coercion -> TyCoVarSet
 -- See Note [Free variables of types]
 tyCoVarsOfCo co = fvVarSet $ tyCoFVsOfCo co
@@ -1607,6 +1638,40 @@ tyCoFVsOfCoVar :: CoVar -> FV
 tyCoFVsOfCoVar v fv_cand in_scope acc
   = (unitFV v `unionFV` tyCoFVsOfType (varType v)) fv_cand in_scope acc
 
+tyCoNDFVsOfCo :: Coercion -> NDFV
+-- Extracts type and coercion variables from a coercion
+-- See Note [Free variables of types]
+tyCoNDFVsOfCo (Refl _ ty)         fv_cand in_scope acc = tyCoNDFVsOfType ty fv_cand in_scope acc
+tyCoNDFVsOfCo (TyConAppCo _ _ cos) fv_cand in_scope acc = tyCoNDFVsOfCos cos fv_cand in_scope acc
+tyCoNDFVsOfCo (AppCo co arg) fv_cand in_scope acc
+  = (tyCoNDFVsOfCo co `unionNDFV` tyCoNDFVsOfCo arg) fv_cand in_scope acc
+tyCoNDFVsOfCo (ForAllCo tv kind_co co) fv_cand in_scope acc
+  = (delNDFV tv (tyCoNDFVsOfCo co) `unionNDFV` tyCoNDFVsOfCo kind_co) fv_cand in_scope acc
+tyCoNDFVsOfCo (FunCo _ co1 co2)    fv_cand in_scope acc
+  = (tyCoNDFVsOfCo co1 `unionNDFV` tyCoNDFVsOfCo co2) fv_cand in_scope acc
+tyCoNDFVsOfCo (CoVarCo v) fv_cand in_scope acc
+  = tyCoNDFVsOfCoVar v fv_cand in_scope acc
+tyCoNDFVsOfCo (HoleCo h) fv_cand in_scope acc
+  = tyCoNDFVsOfCoVar (coHoleCoVar h) fv_cand in_scope acc
+    -- See Note [CoercionHoles and coercion free variables]
+tyCoNDFVsOfCo (AxiomInstCo _ _ cos) fv_cand in_scope acc = tyCoNDFVsOfCos cos fv_cand in_scope acc
+tyCoNDFVsOfCo (UnivCo p _ t1 t2) fv_cand in_scope acc
+  = (tyCoNDFVsOfProv p `unionNDFV` tyCoNDFVsOfType t1
+                     `unionNDFV` tyCoNDFVsOfType t2) fv_cand in_scope acc
+tyCoNDFVsOfCo (SymCo co)          fv_cand in_scope acc = tyCoNDFVsOfCo co fv_cand in_scope acc
+tyCoNDFVsOfCo (TransCo co1 co2)   fv_cand in_scope acc = (tyCoNDFVsOfCo co1 `unionNDFV` tyCoNDFVsOfCo co2) fv_cand in_scope acc
+tyCoNDFVsOfCo (NthCo _ _ co)      fv_cand in_scope acc = tyCoNDFVsOfCo co fv_cand in_scope acc
+tyCoNDFVsOfCo (LRCo _ co)         fv_cand in_scope acc = tyCoNDFVsOfCo co fv_cand in_scope acc
+tyCoNDFVsOfCo (InstCo co arg)     fv_cand in_scope acc = (tyCoNDFVsOfCo co `unionNDFV` tyCoNDFVsOfCo arg) fv_cand in_scope acc
+tyCoNDFVsOfCo (CoherenceCo c1 c2) fv_cand in_scope acc = (tyCoNDFVsOfCo c1 `unionNDFV` tyCoNDFVsOfCo c2) fv_cand in_scope acc
+tyCoNDFVsOfCo (KindCo co)         fv_cand in_scope acc = tyCoNDFVsOfCo co fv_cand in_scope acc
+tyCoNDFVsOfCo (SubCo co)          fv_cand in_scope acc = tyCoNDFVsOfCo co fv_cand in_scope acc
+tyCoNDFVsOfCo (AxiomRuleCo _ cs)  fv_cand in_scope acc = tyCoNDFVsOfCos cs fv_cand in_scope acc
+
+tyCoNDFVsOfCoVar :: CoVar -> NDFV
+tyCoNDFVsOfCoVar v fv_cand in_scope acc
+  = (unitNDFV v `unionNDFV` tyCoNDFVsOfType (varType v)) fv_cand in_scope acc
+
 tyCoVarsOfProv :: UnivCoProvenance -> TyCoVarSet
 tyCoVarsOfProv prov = fvVarSet $ tyCoFVsOfProv prov
 
@@ -1616,6 +1681,12 @@ tyCoFVsOfProv (PhantomProv co)    fv_cand in_scope acc = tyCoFVsOfCo co fv_cand
 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
 
+tyCoNDFVsOfProv :: UnivCoProvenance -> NDFV
+tyCoNDFVsOfProv UnsafeCoerceProv    fv_cand in_scope acc = emptyNDFV fv_cand in_scope acc
+tyCoNDFVsOfProv (PhantomProv co)    fv_cand in_scope acc = tyCoNDFVsOfCo co fv_cand in_scope acc
+tyCoNDFVsOfProv (ProofIrrelProv co) fv_cand in_scope acc = tyCoNDFVsOfCo co fv_cand in_scope acc
+tyCoNDFVsOfProv (PluginProv _)      fv_cand in_scope acc = emptyNDFV fv_cand in_scope acc
+
 tyCoVarsOfCos :: [Coercion] -> TyCoVarSet
 tyCoVarsOfCos cos = fvVarSet $ tyCoFVsOfCos cos
 
@@ -1628,6 +1699,10 @@ 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
 
+tyCoNDFVsOfCos :: [Coercion] -> NDFV
+tyCoNDFVsOfCos []       fv_cand in_scope acc = emptyNDFV fv_cand in_scope acc
+tyCoNDFVsOfCos (co:cos) fv_cand in_scope acc = (tyCoNDFVsOfCo co `unionNDFV` tyCoNDFVsOfCos cos) fv_cand in_scope acc
+
 coVarsOfType :: Type -> CoVarSet
 coVarsOfType (TyVarTy v)         = coVarsOfType (tyVarKind v)
 coVarsOfType (TyConApp _ tys)    = coVarsOfTypes tys
diff --git a/compiler/utils/NDFV.hs b/compiler/utils/NDFV.hs
new file mode 100644
index 0000000000..ce114e55be
--- /dev/null
+++ b/compiler/utils/NDFV.hs
@@ -0,0 +1,201 @@
+{-
+(c) Bartosz Nitka, Facebook 2015
+
+Utilities for efficiently and deterministically computing free variables.
+
+-}
+
+{-# LANGUAGE BangPatterns #-}
+
+module NDFV (
+        -- * Nondeterministic free vars computations
+        NDFV, InterestingVarFun,
+
+        -- * Running the computations
+        ndfvVarListVarSet,
+        -- ndfvVarList,
+        ndfvVarSet,
+        -- ndfvDVarSet,
+
+        -- ** Manipulating those computations
+        unitNDFV,
+        emptyNDFV,
+        mkNDFVs,
+        unionNDFV,
+        unionsNDFV,
+        delNDFV,
+        delNDFVs,
+        filterNDFV,
+        mapUnionNDFV,
+    ) where
+
+import GhcPrelude
+
+import Var
+import VarSet
+
+-- | Predicate on possible free variables: returns @True@ iff the variable is
+-- interesting
+type InterestingVarFun = Var -> Bool
+
+-- Note [Nondeterministic FV]
+-- ~~~~~~~~~~~~~~~~~~~~~~~
+-- When computing free variables, the order in which you get them affects
+-- the results of floating and specialization. If you use UniqFM to collect
+-- them and then turn that into a list, you get them in nondeterministic
+-- order as described in Note [Deterministic UniqFM] in UniqDFM.
+
+-- A naive algorithm for free variables relies on merging sets of variables.
+-- Merging costs O(n+m) for UniqFM and for UniqDFM there's an additional log
+-- factor. It's cheaper to incrementally add to a list and use a set to check
+-- for duplicates.
+type NDFV = InterestingVarFun
+             -- Used for filtering sets as we build them
+          -> VarSet
+             -- Locally bound variables
+          -> VarSet
+          -> VarSet
+
+-- Note [NDFV naming conventions]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+-- To get the performance and determinism that NDFV provides, NDFV computations
+-- need to built up from smaller NDFV computations and then evaluated with
+-- one of `ndfvVarList`, `ndfvDVarSet`, `ndfvVarListVarSet`. That means the functions
+-- returning NDFV need to be exported.
+--
+-- The conventions are:
+--
+-- a) non-deterministic functions:
+--   * a function that returns VarSet
+--       e.g. `tyVarsOfType`
+-- b) deterministic functions:
+--   * a worker that returns NDFV
+--       e.g. `tyNDFVsOfType`
+--   * a function that returns [Var]
+--       e.g. `tyVarsOfTypeList`
+--   * a function that returns DVarSet
+--       e.g. `tyVarsOfTypeDSet`
+--
+-- Where tyVarsOfType, tyVarsOfTypeList, tyVarsOfTypeDSet are implemented
+-- in terms of the worker evaluated with ndfvVarSet, ndfvVarList, ndfvDVarSet
+-- respectively.
+
+-- | Run a free variable computation, returning a list of distinct free
+-- variables in deterministic order and a non-deterministic set containing
+-- those variables.
+ndfvVarListVarSet :: NDFV ->  VarSet
+ndfvVarListVarSet fv = fv (const True) emptyVarSet emptyVarSet
+
+-- | Run a free variable computation, returning a list of distinct free
+-- variables in nondeterministic order.
+---- ndfvVarList :: NDFV -> [Var]
+---- ndfvVarList = fst . ndfvVarListVarSet
+
+-- | Run a free variable computation, returning a deterministic set of free
+-- variables. Note that this is just a wrapper around the version that
+-- returns a deterministic list. If you need a list you should use
+-- `ndfvVarList`.
+---- ndfvDVarSet :: NDFV -> DVarSet
+---- ndfvDVarSet = mkDVarSet . fst . ndfvVarListVarSet
+
+-- | Run a free variable computation, returning a non-deterministic set of
+-- free variables. Don't use if the set will be later converted to a list
+-- and the order of that list will impact the generated code.
+ndfvVarSet :: NDFV -> VarSet
+ndfvVarSet = ndfvVarListVarSet
+
+-- Note [NDFV eta expansion]
+-- ~~~~~~~~~~~~~~~~~~~~~~~
+-- Let's consider an eta-reduced implementation of freeVarsOf using NDFV:
+--
+-- freeVarsOf (App a b) = freeVarsOf a `unionNDFV` freeVarsOf b
+--
+-- If GHC doesn't eta-expand it, after inlining unionNDFV we end up with
+--
+-- freeVarsOf = \x ->
+--   case x of
+--     App a b -> \ndfv_cand in_scope acc ->
+--       freeVarsOf a ndfv_cand in_scope $! freeVarsOf b ndfv_cand in_scope $! acc
+--
+-- which has to create a thunk, resulting in more allocations.
+--
+-- On the other hand if it is eta-expanded:
+--
+-- freeVarsOf (App a b) ndfv_cand in_scope acc =
+--   (freeVarsOf a `unionNDFV` freeVarsOf b) ndfv_cand in_scope acc
+--
+-- after inlining unionNDFV we have:
+--
+-- freeVarsOf = \x ndfv_cand in_scope acc ->
+--   case x of
+--     App a b ->
+--       freeVarsOf a ndfv_cand in_scope $! freeVarsOf b ndfv_cand in_scope $! acc
+--
+-- which saves allocations.
+--
+-- GHC when presented with knowledge about all the call sites, correctly
+-- eta-expands in this case. Unfortunately due to the fact that freeVarsOf gets
+-- exported to be composed with other functions, GHC doesn't have that
+-- information and has to be more conservative here.
+--
+-- Hence functions that get exported and return NDFV need to be manually
+-- eta-expanded. See also #11146.
+
+-- | Add a variable - when free, to the returned free variables.
+-- Ignores duplicates and respects the filtering function.
+unitNDFV :: Id -> NDFV
+unitNDFV var ndfv_cand in_scope acc
+  | var `elemVarSet` in_scope = acc
+  | var `elemVarSet` acc = acc
+  | ndfv_cand var = extendVarSet acc var
+  | otherwise = acc
+{-# INLINE unitNDFV #-}
+
+-- | Return no free variables.
+emptyNDFV :: NDFV
+emptyNDFV _ _ acc = acc
+{-# INLINE emptyNDFV #-}
+
+-- | Union two free variable computations.
+unionNDFV :: NDFV -> NDFV -> NDFV
+unionNDFV fv1 fv2 fv_cand in_scope acc =
+  fv1 fv_cand in_scope $! fv2 fv_cand in_scope $! acc
+{-# INLINE unionNDFV #-}
+
+-- | Mark the variable as not free by putting it in scope.
+delNDFV :: Var -> NDFV -> NDFV
+delNDFV var fv fv_cand !in_scope acc =
+  fv fv_cand (extendVarSet in_scope var) acc
+{-# INLINE delNDFV #-}
+
+-- | Mark many free variables as not free.
+delNDFVs :: VarSet -> NDFV -> NDFV
+delNDFVs vars fv fv_cand !in_scope acc =
+  fv fv_cand (in_scope `unionVarSet` vars) acc
+{-# INLINE delNDFVs #-}
+
+-- | Filter a free variable computation.
+filterNDFV :: InterestingVarFun -> NDFV -> NDFV
+filterNDFV fv_cand2 fv fv_cand1 in_scope acc =
+  fv (\v -> fv_cand1 v && fv_cand2 v) in_scope acc
+{-# INLINE filterNDFV #-}
+
+-- | Map a free variable computation over a list and union the results.
+mapUnionNDFV :: (a -> NDFV) -> [a] -> NDFV
+mapUnionNDFV _f [] _fv_cand _in_scope acc = acc
+mapUnionNDFV f (a:as) fv_cand in_scope acc =
+  f a fv_cand in_scope $! mapUnionNDFV f as fv_cand in_scope $! acc
+  -- NB: preserve ordering of the input list by treating a before as
+{-# INLINABLE mapUnionNDFV #-}
+
+-- | Union many free variable computations.
+unionsNDFV :: [NDFV] -> NDFV
+unionsNDFV fvs fv_cand in_scope acc = mapUnionNDFV id fvs fv_cand in_scope acc
+{-# INLINE unionsNDFV #-}
+
+-- | Add multiple variables - when free, to the returned free variables.
+-- Ignores duplicates and respects the filtering function.
+mkNDFVs :: [Var] -> NDFV
+mkNDFVs vars fv_cand in_scope acc =
+  mapUnionNDFV unitNDFV vars fv_cand in_scope acc
+{-# INLINE mkNDFVs #-}