Improve some folds over Uniq[D]FM

* Replace some non-deterministic lazy folds with
  strict folds.
* Replace some O(n log n) folds in deterministic order
  with O(n) non-deterministic folds.
* Replace some folds with set-operations on the underlying
  IntMaps.

This reduces max residency when compiling
`nofib/spectral/simple/Main.hs` with -O0 by about 1%.

Maximum residency when compiling Cabal also seems reduced on the
order of 3-9%.

1 files changed, 21 insertions, 3 deletions

diff --git a/compiler/GHC/Tc/Types/Evidence.hs b/compiler/GHC/Tc/Types/Evidence.hs
index 49ae605feb..8649871670 100644
--- a/compiler/GHC/Tc/Types/Evidence.hs
+++ b/compiler/GHC/Tc/Types/Evidence.hs
@@ -15,8 +15,12 @@ module GHC.Tc.Types.Evidence (
   -- * Evidence bindings
   TcEvBinds(..), EvBindsVar(..),
   EvBindMap(..), emptyEvBindMap, extendEvBinds,
-  lookupEvBind, evBindMapBinds, foldEvBindMap, filterEvBindMap,
+  lookupEvBind, evBindMapBinds,
+  foldEvBindMap, nonDetStrictFoldEvBindMap,
+  filterEvBindMap,
   isEmptyEvBindMap,
+  evBindMapToVarSet,
+  varSetMinusEvBindMap,
   EvBind(..), emptyTcEvBinds, isEmptyTcEvBinds, mkGivenEvBind, mkWantedEvBind,
   evBindVar, isCoEvBindsVar,
 
@@ -55,6 +59,8 @@ module GHC.Tc.Types.Evidence (
 
 import GHC.Prelude
 
+import GHC.Types.Unique.DFM
+import GHC.Types.Unique.FM
 import GHC.Types.Var
 import GHC.Core.Coercion.Axiom
 import GHC.Core.Coercion
@@ -496,10 +502,22 @@ evBindMapBinds = foldEvBindMap consBag emptyBag
 foldEvBindMap :: (EvBind -> a -> a) -> a -> EvBindMap -> a
 foldEvBindMap k z bs = foldDVarEnv k z (ev_bind_varenv bs)
 
+-- See Note [Deterministic UniqFM] to learn about nondeterminism.
+-- If you use this please provide a justification why it doesn't introduce
+-- nondeterminism.
+nonDetStrictFoldEvBindMap :: (EvBind -> a -> a) -> a -> EvBindMap -> a
+nonDetStrictFoldEvBindMap k z bs = nonDetStrictFoldDVarEnv k z (ev_bind_varenv bs)
+
 filterEvBindMap :: (EvBind -> Bool) -> EvBindMap -> EvBindMap
 filterEvBindMap k (EvBindMap { ev_bind_varenv = env })
   = EvBindMap { ev_bind_varenv = filterDVarEnv k env }
 
+evBindMapToVarSet :: EvBindMap -> VarSet
+evBindMapToVarSet (EvBindMap dve) = unsafeUFMToUniqSet (mapUFM evBindVar (udfmToUfm dve))
+
+varSetMinusEvBindMap :: VarSet -> EvBindMap -> VarSet
+varSetMinusEvBindMap vs (EvBindMap dve) = vs `uniqSetMinusUDFM` dve
+
 instance Outputable EvBindMap where
   ppr (EvBindMap m) = ppr m
 
@@ -851,8 +869,8 @@ findNeededEvVars ev_binds seeds
   = transCloVarSet also_needs seeds
   where
    also_needs :: VarSet -> VarSet
-   also_needs needs = nonDetFoldUniqSet add emptyVarSet needs
-     -- It's OK to use nonDetFoldUFM here because we immediately
+   also_needs needs = nonDetStrictFoldUniqSet add emptyVarSet needs
+     -- It's OK to use a non-deterministic fold here because we immediately
      -- forget about the ordering by creating a set
 
    add :: Var -> VarSet -> VarSet