Kill varEnvElts in specImports

We need the order of specialized binds and rules to be deterministic,
so we use a deterministic set here.

Test Plan: ./validate

Reviewers: simonmar, bgamari, austin, simonpj

Reviewed By: simonpj

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D2197

GHC Trac Issues: #4012

3 files changed, 73 insertions, 12 deletions

diff --git a/compiler/basicTypes/VarEnv.hs b/compiler/basicTypes/VarEnv.hs
index 917946f56e..c591ee452c 100644
--- a/compiler/basicTypes/VarEnv.hs
+++ b/compiler/basicTypes/VarEnv.hs
@@ -24,15 +24,20 @@ module VarEnv (
         partitionVarEnv,
 
         -- * Deterministic Var environments (maps)
-        DVarEnv,
+        DVarEnv, DIdEnv,
 
         -- ** Manipulating these environments
         emptyDVarEnv,
+        dVarEnvElts,
         extendDVarEnv,
         lookupDVarEnv,
         foldDVarEnv,
         mapDVarEnv,
         alterDVarEnv,
+        plusDVarEnv_C,
+        unitDVarEnv,
+        delDVarEnv,
+        delDVarEnvList,
 
         -- * The InScopeSet type
         InScopeSet,
@@ -503,11 +508,15 @@ modifyVarEnv_Directly mangle_fn env key
 -- See Note [Deterministic UniqFM] in UniqDFM for explanation why we need
 -- DVarEnv.
 
-type DVarEnv elt   = UniqDFM elt
+type DVarEnv elt = UniqDFM elt
+type DIdEnv elt = DVarEnv elt
 
 emptyDVarEnv :: DVarEnv a
 emptyDVarEnv = emptyUDFM
 
+dVarEnvElts :: DVarEnv a -> [a]
+dVarEnvElts = eltsUDFM
+
 extendDVarEnv :: DVarEnv a -> Var -> a -> DVarEnv a
 extendDVarEnv = addToUDFM
 
@@ -522,3 +531,15 @@ mapDVarEnv = mapUDFM
 
 alterDVarEnv :: (Maybe a -> Maybe a) -> DVarEnv a -> Var -> DVarEnv a
 alterDVarEnv = alterUDFM
+
+plusDVarEnv_C :: (a -> a -> a) -> DVarEnv a -> DVarEnv a -> DVarEnv a
+plusDVarEnv_C = plusUDFM_C
+
+unitDVarEnv :: Var -> a -> DVarEnv a
+unitDVarEnv = unitUDFM
+
+delDVarEnv :: DVarEnv a -> Var -> DVarEnv a
+delDVarEnv = delFromUDFM
+
+delDVarEnvList :: DVarEnv a -> [Var] -> DVarEnv a
+delDVarEnvList = delListFromUDFM
diff --git a/compiler/specialise/Specialise.hs b/compiler/specialise/Specialise.hs
index 09caa0034d..5c76f23af2 100644
--- a/compiler/specialise/Specialise.hs
+++ b/compiler/specialise/Specialise.hs
@@ -35,6 +35,7 @@ import Util
 import Outputable
 import FastString
 import State
+import UniqDFM
 
 import Control.Monad
 #if __GLASGOW_HASKELL__ > 710
@@ -653,7 +654,7 @@ specImports dflags this_mod top_env done callers rule_base cds
     return ([], [])
 
   | otherwise =
-    do { let import_calls = varEnvElts cds
+    do { let import_calls = dVarEnvElts cds
        ; (rules, spec_binds) <- go rule_base import_calls
        ; return (rules, spec_binds) }
   where
@@ -1720,10 +1721,13 @@ type DictBind = (CoreBind, VarSet)
 type DictExpr = CoreExpr
 
 emptyUDs :: UsageDetails
-emptyUDs = MkUD { ud_binds = emptyBag, ud_calls = emptyVarEnv }
+emptyUDs = MkUD { ud_binds = emptyBag, ud_calls = emptyDVarEnv }
 
 ------------------------------------------------------------
-type CallDetails  = IdEnv CallInfoSet
+type CallDetails  = DIdEnv CallInfoSet
+  -- The order of specialized binds and rules depends on how we linearize
+  -- CallDetails, so to get determinism we must use a deterministic set here.
+  -- See Note [Deterministic UniqFM] in UniqDFM
 newtype CallKey   = CallKey [Maybe Type]                        -- Nothing => unconstrained type argument
 
 -- CallInfo uses a Map, thereby ensuring that
@@ -1768,13 +1772,16 @@ instance Ord CallKey where
                   cmp (Just t1) (Just t2) = cmpType t1 t2
 
 unionCalls :: CallDetails -> CallDetails -> CallDetails
-unionCalls c1 c2 = plusVarEnv_C unionCallInfoSet c1 c2
+unionCalls c1 c2 = plusDVarEnv_C unionCallInfoSet c1 c2
 
 unionCallInfoSet :: CallInfoSet -> CallInfoSet -> CallInfoSet
 unionCallInfoSet (CIS f calls1) (CIS _ calls2) = CIS f (calls1 `Map.union` calls2)
 
 callDetailsFVs :: CallDetails -> VarSet
-callDetailsFVs calls = foldVarEnv (unionVarSet . callInfoFVs) emptyVarSet calls
+callDetailsFVs calls =
+  nonDetFoldUDFM (unionVarSet . callInfoFVs) emptyVarSet calls
+  -- It's OK to use nonDetFoldUDFM here because we forget the ordering
+  -- immediately by converting to a nondeterministic set.
 
 callInfoFVs :: CallInfoSet -> VarSet
 callInfoFVs (CIS _ call_info) = Map.foldRight (\(_,fv) vs -> unionVarSet fv vs) emptyVarSet call_info
@@ -1783,7 +1790,7 @@ callInfoFVs (CIS _ call_info) = Map.foldRight (\(_,fv) vs -> unionVarSet fv vs)
 singleCall :: Id -> [Maybe Type] -> [DictExpr] -> UsageDetails
 singleCall id tys dicts
   = MkUD {ud_binds = emptyBag,
-          ud_calls = unitVarEnv id $ CIS id $
+          ud_calls = unitDVarEnv id $ CIS id $
                      Map.singleton (CallKey tys) (dicts, call_fvs) }
   where
     call_fvs = exprsFreeVars dicts `unionVarSet` tys_fvs
@@ -2033,8 +2040,9 @@ callsForMe fn (MkUD { ud_binds = orig_dbs, ud_calls = orig_calls })
     --                 text "Calls for me =" <+> ppr calls_for_me]) $
     (uds_without_me, calls_for_me)
   where
-    uds_without_me = MkUD { ud_binds = orig_dbs, ud_calls = delVarEnv orig_calls fn }
-    calls_for_me = case lookupVarEnv orig_calls fn of
+    uds_without_me = MkUD { ud_binds = orig_dbs
+                          , ud_calls = delDVarEnv orig_calls fn }
+    calls_for_me = case lookupDVarEnv orig_calls fn of
                         Nothing -> []
                         Just (CIS _ calls) -> filter_dfuns (Map.toList calls)
 
@@ -2070,7 +2078,7 @@ splitDictBinds dbs bndr_set
 deleteCallsMentioning :: VarSet -> CallDetails -> CallDetails
 -- Remove calls *mentioning* bs
 deleteCallsMentioning bs calls
-  = mapVarEnv filter_calls calls
+  = mapDVarEnv filter_calls calls
   where
     filter_calls :: CallInfoSet -> CallInfoSet
     filter_calls (CIS f calls) = CIS f (Map.filter keep_call calls)
@@ -2078,7 +2086,7 @@ deleteCallsMentioning bs calls
 
 deleteCallsFor :: [Id] -> CallDetails -> CallDetails
 -- Remove calls *for* bs
-deleteCallsFor bs calls = delVarEnvList calls bs
+deleteCallsFor bs calls = delDVarEnvList calls bs
 
 {-
 ************************************************************************
diff --git a/compiler/utils/UniqDFM.hs b/compiler/utils/UniqDFM.hs
index 1b3cade93a..9dfefa4bdb 100644
--- a/compiler/utils/UniqDFM.hs
+++ b/compiler/utils/UniqDFM.hs
@@ -33,6 +33,7 @@ module UniqDFM (
         alterUDFM,
         mapUDFM,
         plusUDFM,
+        plusUDFM_C,
         lookupUDFM,
         elemUDFM,
         foldUDFM,
@@ -49,6 +50,7 @@ module UniqDFM (
 
         udfmToList,
         udfmToUfm,
+        nonDetFoldUDFM,
         alwaysUnsafeUfmToUdfm,
     ) where
 
@@ -144,12 +146,30 @@ addToUDFM_Directly :: UniqDFM elt -> Unique -> elt -> UniqDFM elt
 addToUDFM_Directly (UDFM m i) u v =
   UDFM (M.insert (getKey u) (TaggedVal v i) m) (i + 1)
 
+addToUDFM_Directly_C
+  :: (elt -> elt -> elt) -> UniqDFM elt -> Unique -> elt -> UniqDFM elt
+addToUDFM_Directly_C f (UDFM m i) u v =
+  UDFM (M.insertWith tf (getKey u) (TaggedVal v i) m) (i + 1)
+  where
+  tf (TaggedVal a j) (TaggedVal b _) = TaggedVal (f a b) j
+
 addListToUDFM_Directly :: UniqDFM elt -> [(Unique,elt)] -> UniqDFM elt
 addListToUDFM_Directly = foldl (\m (k, v) -> addToUDFM_Directly m k v)
 
+addListToUDFM_Directly_C
+  :: (elt -> elt -> elt) -> UniqDFM elt -> [(Unique,elt)] -> UniqDFM elt
+addListToUDFM_Directly_C f = foldl (\m (k, v) -> addToUDFM_Directly_C f m k v)
+
 delFromUDFM :: Uniquable key => UniqDFM elt -> key -> UniqDFM elt
 delFromUDFM (UDFM m i) k = UDFM (M.delete (getKey $ getUnique k) m) i
 
+plusUDFM_C :: (elt -> elt -> elt) -> UniqDFM elt -> UniqDFM elt -> UniqDFM elt
+plusUDFM_C f udfml@(UDFM _ i) udfmr@(UDFM _ j)
+  -- we will use the upper bound on the tag as a proxy for the set size,
+  -- to insert the smaller one into the bigger one
+  | i > j = insertUDFMIntoLeft_C f udfml udfmr
+  | otherwise = insertUDFMIntoLeft_C f udfmr udfml
+
 -- Note [Overflow on plusUDFM]
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 -- There are multiple ways of implementing plusUDFM.
@@ -193,6 +213,11 @@ plusUDFM udfml@(UDFM _ i) udfmr@(UDFM _ j)
 insertUDFMIntoLeft :: UniqDFM elt -> UniqDFM elt -> UniqDFM elt
 insertUDFMIntoLeft udfml udfmr = addListToUDFM_Directly udfml $ udfmToList udfmr
 
+insertUDFMIntoLeft_C
+  :: (elt -> elt -> elt) -> UniqDFM elt -> UniqDFM elt -> UniqDFM elt
+insertUDFMIntoLeft_C f udfml udfmr =
+  addListToUDFM_Directly_C f udfml $ udfmToList udfmr
+
 lookupUDFM :: Uniquable key => UniqDFM elt -> key -> Maybe elt
 lookupUDFM (UDFM m _i) k = taggedFst `fmap` M.lookup (getKey $ getUnique k) m
 
@@ -204,6 +229,13 @@ elemUDFM k (UDFM m _i) = M.member (getKey $ getUnique k) m
 foldUDFM :: (elt -> a -> a) -> a -> UniqDFM elt -> a
 foldUDFM k z m = foldr k z (eltsUDFM m)
 
+-- | Performs a nondeterministic fold over the UniqDFM.
+-- It's O(n), same as the corresponding function on `UniqFM`.
+-- If you use this please provide a justification why it doesn't introduce
+-- nondeterminism.
+nonDetFoldUDFM :: (elt -> a -> a) -> a -> UniqDFM elt -> a
+nonDetFoldUDFM k z (UDFM m _i) = foldr k z $ map taggedFst $ M.elems m
+
 eltsUDFM :: UniqDFM elt -> [elt]
 eltsUDFM (UDFM m _i) =
   map taggedFst $ sortBy (compare `on` taggedSnd) $ M.elems m