Upgrade UniqSet to a newtype

The fundamental problem with `type UniqSet = UniqFM` is that `UniqSet`
has a key invariant `UniqFM` does not. For example, `fmap` over
`UniqSet` will generally produce nonsense.

* Upgrade `UniqSet` from a type synonym to a newtype.

* Remove unused and shady `extendVarSet_C` and `addOneToUniqSet_C`.

* Use cached unique in `tyConsOfType` by replacing
  `unitNameEnv (tyConName tc) tc` with `unitUniqSet tc`.

Reviewers: austin, hvr, goldfire, simonmar, niteria, bgamari

Reviewed By: niteria

Subscribers: thomie

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

6 files changed, 121 insertions, 43 deletions

diff --git a/compiler/utils/GraphColor.hs b/compiler/utils/GraphColor.hs
index 056ce0daa4..492125b787 100644
--- a/compiler/utils/GraphColor.hs
+++ b/compiler/utils/GraphColor.hs
@@ -309,7 +309,7 @@ selectColor colors graph u
         Just nsConflicts
                         = sequence
                         $ map (lookupNode graph)
-                        $ nonDetEltsUFM
+                        $ nonDetEltsUniqSet
                         $ nodeConflicts node
                         -- See Note [Unique Determinism and code generation]
 
@@ -356,7 +356,7 @@ selectColor colors graph u
 
                 -- it wasn't a preference, but it was still ok
                 | not $ isEmptyUniqSet colors_ok
-                , c : _         <- nonDetEltsUFM colors_ok
+                , c : _         <- nonDetEltsUniqSet colors_ok
                 -- See Note [Unique Determinism and code generation]
                 = Just c
 
diff --git a/compiler/utils/GraphOps.hs b/compiler/utils/GraphOps.hs
index 0985797571..3677e517b5 100644
--- a/compiler/utils/GraphOps.hs
+++ b/compiler/utils/GraphOps.hs
@@ -59,7 +59,7 @@ addNode k node graph
  = let
         -- add back conflict edges from other nodes to this one
         map_conflict =
-          nonDetFoldUFM
+          nonDetFoldUniqSet
             -- It's OK to use nonDetFoldUFM here because the
             -- operation is commutative
             (adjustUFM_C (\n -> n { nodeConflicts =
@@ -69,7 +69,7 @@ addNode k node graph
 
         -- add back coalesce edges from other nodes to this one
         map_coalesce =
-          nonDetFoldUFM
+          nonDetFoldUniqSet
             -- It's OK to use nonDetFoldUFM here because the
             -- operation is commutative
             (adjustUFM_C (\n -> n { nodeCoalesce =
@@ -89,11 +89,11 @@ delNode k graph
         | Just node     <- lookupNode graph k
         = let   -- delete conflict edges from other nodes to this one.
                 graph1  = foldl' (\g k1 -> let Just g' = delConflict k1 k g in g') graph
-                        $ nonDetEltsUFM (nodeConflicts node)
+                        $ nonDetEltsUniqSet (nodeConflicts node)
 
                 -- delete coalesce edge from other nodes to this one.
                 graph2  = foldl' (\g k1 -> let Just g' = delCoalesce k1 k g in g') graph1
-                        $ nonDetEltsUFM (nodeCoalesce node)
+                        $ nonDetEltsUniqSet (nodeCoalesce node)
                         -- See Note [Unique Determinism and code generation]
 
                 -- delete the node
@@ -182,7 +182,7 @@ addConflicts
 addConflicts conflicts getClass
 
         -- just a single node, but no conflicts, create the node anyway.
-        | (u : [])      <- nonDetEltsUFM conflicts
+        | (u : [])      <- nonDetEltsUniqSet conflicts
         = graphMapModify
         $ adjustWithDefaultUFM
                 id
@@ -191,8 +191,8 @@ addConflicts conflicts getClass
 
         | otherwise
         = graphMapModify
-        $ (\fm -> foldl' (\g u  -> addConflictSet1 u getClass conflicts g) fm
-                $ nonDetEltsUFM conflicts)
+        $ \fm -> foldl' (\g u  -> addConflictSet1 u getClass conflicts g) fm
+                $ nonDetEltsUniqSet conflicts
                 -- See Note [Unique Determinism and code generation]
 
 
@@ -318,7 +318,7 @@ coalesceGraph' aggressive triv graph kkPairsAcc
         --
         cList   = [ (nodeId node1, k2)
                         | node1 <- cNodes
-                        , k2    <- nonDetEltsUFM $ nodeCoalesce node1 ]
+                        , k2    <- nonDetEltsUniqSet $ nodeCoalesce node1 ]
                         -- See Note [Unique Determinism and code generation]
 
         -- do the coalescing, returning the new graph and a list of pairs of keys
@@ -472,7 +472,7 @@ freezeNode k
                 else node       -- panic "GraphOps.freezeNode: edge to freeze wasn't in the coalesce set"
                                 -- If the edge isn't actually in the coelesce set then just ignore it.
 
-        fm2     = nonDetFoldUFM (adjustUFM_C (freezeEdge k)) fm1
+        fm2     = nonDetFoldUniqSet (adjustUFM_C (freezeEdge k)) fm1
                     -- It's OK to use nonDetFoldUFM here because the operation
                     -- is commutative
                         $ nodeCoalesce node
@@ -568,7 +568,7 @@ validateGraph doc isColored graph
         , not $ isEmptyUniqSet badEdges
         = pprPanic "GraphOps.validateGraph"
                 (  text "Graph has edges that point to non-existent nodes"
-                $$ text "  bad edges: " <> pprUFM badEdges (vcat . map ppr)
+                $$ text "  bad edges: " <> pprUFM (getUniqSet badEdges) (vcat . map ppr)
                 $$ doc )
 
         -- Check that no conflicting nodes have the same color
@@ -609,7 +609,7 @@ checkNode
 checkNode graph node
         | Just color            <- nodeColor node
         , Just neighbors        <- sequence $ map (lookupNode graph)
-                                $  nonDetEltsUFM $ nodeConflicts node
+                                $  nonDetEltsUniqSet $ nodeConflicts node
             -- See Note [Unique Determinism and code generation]
 
         , neighbourColors       <- catMaybes $ map nodeColor neighbors
diff --git a/compiler/utils/GraphPpr.hs b/compiler/utils/GraphPpr.hs
index f5276842aa..a40e1058d0 100644
--- a/compiler/utils/GraphPpr.hs
+++ b/compiler/utils/GraphPpr.hs
@@ -87,7 +87,7 @@ dotNode colorMap triv node
         excludes
                 = hcat $ punctuate space
                 $ map (\n -> text "-" <> ppr n)
-                $ nonDetEltsUFM $ nodeExclusions node
+                $ nonDetEltsUniqSet $ nodeExclusions node
                 -- See Note [Unique Determinism and code generation]
 
         preferences
@@ -146,13 +146,13 @@ dotNodeEdges visited node
         | otherwise
         = let   dconflicts
                         = map (dotEdgeConflict (nodeId node))
-                        $ nonDetEltsUFM
+                        $ nonDetEltsUniqSet
                         -- See Note [Unique Determinism and code generation]
                         $ minusUniqSet (nodeConflicts node) visited
 
                 dcoalesces
                         = map (dotEdgeCoalesce (nodeId node))
-                        $ nonDetEltsUFM
+                        $ nonDetEltsUniqSet
                         -- See Note [Unique Determinism and code generation]
                         $ minusUniqSet (nodeCoalesce node) visited
 
diff --git a/compiler/utils/UniqDSet.hs b/compiler/utils/UniqDSet.hs
index 90e9996d1a..4e8c7ed97f 100644
--- a/compiler/utils/UniqDSet.hs
+++ b/compiler/utils/UniqDSet.hs
@@ -70,7 +70,7 @@ minusUniqDSet :: UniqDSet a -> UniqDSet a -> UniqDSet a
 minusUniqDSet = minusUDFM
 
 uniqDSetMinusUniqSet :: UniqDSet a -> UniqSet a -> UniqDSet a
-uniqDSetMinusUniqSet = udfmMinusUFM
+uniqDSetMinusUniqSet xs ys = udfmMinusUFM xs (getUniqSet ys)
 
 intersectUniqDSets :: UniqDSet a -> UniqDSet a -> UniqDSet a
 intersectUniqDSets = intersectUDFM
diff --git a/compiler/utils/UniqFM.hs b/compiler/utils/UniqFM.hs
index 49ceb89d90..8214f1704b 100644
--- a/compiler/utils/UniqFM.hs
+++ b/compiler/utils/UniqFM.hs
@@ -233,7 +233,7 @@ plusUFMList = foldl' plusUFM emptyUFM
 minusUFM :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
 minusUFM (UFM x) (UFM y) = UFM (M.difference x y)
 
-intersectUFM :: UniqFM elt -> UniqFM elt -> UniqFM elt
+intersectUFM :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
 intersectUFM (UFM x) (UFM y) = UFM (M.intersection x y)
 
 intersectUFM_C
diff --git a/compiler/utils/UniqSet.hs b/compiler/utils/UniqSet.hs
index 6f58652f80..ede900a842 100644
--- a/compiler/utils/UniqSet.hs
+++ b/compiler/utils/UniqSet.hs
@@ -8,33 +8,54 @@ Based on @UniqFMs@ (as you would expect).
 
 Basically, the things need to be in class @Uniquable@.
 -}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DeriveDataTypeable #-}
 
 module UniqSet (
         -- * Unique set type
         UniqSet,    -- type synonym for UniqFM a
+        getUniqSet,
+        pprUniqSet,
 
         -- ** Manipulating these sets
         emptyUniqSet,
         unitUniqSet,
         mkUniqSet,
-        addOneToUniqSet, addOneToUniqSet_C, addListToUniqSet,
+        addOneToUniqSet, addListToUniqSet,
         delOneFromUniqSet, delOneFromUniqSet_Directly, delListFromUniqSet,
+        delListFromUniqSet_Directly,
         unionUniqSets, unionManyUniqSets,
-        minusUniqSet,
+        minusUniqSet, uniqSetMinusUFM,
         intersectUniqSets,
+        restrictUniqSetToUFM,
         uniqSetAny, uniqSetAll,
         elementOfUniqSet,
         elemUniqSet_Directly,
         filterUniqSet,
+        filterUniqSet_Directly,
         sizeUniqSet,
         isEmptyUniqSet,
         lookupUniqSet,
-        partitionUniqSet
+        lookupUniqSet_Directly,
+        partitionUniqSet,
+        mapUniqSet,
+        unsafeUFMToUniqSet,
+        nonDetEltsUniqSet,
+        nonDetKeysUniqSet,
+        nonDetFoldUniqSet,
+        nonDetFoldUniqSet_Directly
     ) where
 
 import UniqFM
 import Unique
+import Data.Coerce
+import Outputable
 import Data.Foldable (foldl')
+import Data.Data
+#if __GLASGOW_HASKELL__ >= 801
+import qualified Data.Semigroup
+#endif
 
 {-
 ************************************************************************
@@ -49,26 +70,45 @@ unitUniqSet :: Uniquable a => a -> UniqSet a
 mkUniqSet :: Uniquable a => [a]  -> UniqSet a
 
 addOneToUniqSet :: Uniquable a => UniqSet a -> a -> UniqSet a
-addOneToUniqSet_C :: Uniquable a => (a -> a -> a) -> UniqSet a -> a -> UniqSet a
 addListToUniqSet :: Uniquable a => UniqSet a -> [a] -> UniqSet a
 
 delOneFromUniqSet :: Uniquable a => UniqSet a -> a -> UniqSet a
 delOneFromUniqSet_Directly :: UniqSet a -> Unique -> UniqSet a
 delListFromUniqSet :: Uniquable a => UniqSet a -> [a] -> UniqSet a
+delListFromUniqSet_Directly :: UniqSet a -> [Unique] -> UniqSet a
 
 unionUniqSets :: UniqSet a -> UniqSet a -> UniqSet a
 unionManyUniqSets :: [UniqSet a] -> UniqSet a
 minusUniqSet  :: UniqSet a -> UniqSet a -> UniqSet a
 intersectUniqSets :: UniqSet a -> UniqSet a -> UniqSet a
+restrictUniqSetToUFM :: UniqSet a -> UniqFM b -> UniqSet a
+uniqSetMinusUFM :: UniqSet a -> UniqFM b -> UniqSet a
 
 elementOfUniqSet :: Uniquable a => a -> UniqSet a -> Bool
 elemUniqSet_Directly :: Unique -> UniqSet a -> Bool
 filterUniqSet :: (a -> Bool) -> UniqSet a -> UniqSet a
+filterUniqSet_Directly :: (Unique -> elt -> Bool) -> UniqSet elt -> UniqSet elt
 partitionUniqSet :: (a -> Bool) -> UniqSet a -> (UniqSet a, UniqSet a)
 
 sizeUniqSet :: UniqSet a -> Int
 isEmptyUniqSet :: UniqSet a -> Bool
 lookupUniqSet :: Uniquable a => UniqSet b -> a -> Maybe b
+lookupUniqSet_Directly :: UniqSet a -> Unique -> Maybe a
+
+nonDetEltsUniqSet :: UniqSet elt -> [elt]
+nonDetKeysUniqSet :: UniqSet elt -> [Unique]
+
+-- See Note [Deterministic UniqFM] to learn about nondeterminism.
+-- If you use this please provide a justification why it doesn't introduce
+-- nondeterminism.
+nonDetFoldUniqSet :: (elt -> a -> a) -> a -> UniqSet elt -> a
+
+-- See Note [Deterministic UniqFM] to learn about nondeterminism.
+-- If you use this please provide a justification why it doesn't introduce
+-- nondeterminism.
+nonDetFoldUniqSet_Directly:: (Unique -> elt -> a -> a) -> a -> UniqSet elt -> a
+
+mapUniqSet :: Uniquable b => (a -> b) -> UniqSet a -> UniqSet b
 
 {-
 ************************************************************************
@@ -87,36 +127,74 @@ lookupUniqSet :: Uniquable a => UniqSet b -> a -> Maybe b
 -- that only updated the values and it's been removed, because it broke
 -- the invariant.
 
-type UniqSet a = UniqFM a
-
-emptyUniqSet = emptyUFM
-unitUniqSet x = unitUFM x x
+newtype UniqSet a = UniqSet {getUniqSet' :: UniqFM a} deriving Data
+getUniqSet :: UniqSet a -> UniqFM a
+getUniqSet = getUniqSet'
+
+-- | 'unsafeUFMToUniqSet' converts a @'UniqFM' a@ into a @'UniqSet' a@
+-- assuming, without checking, that it maps each 'Unique' to a value
+-- that has that 'Unique'. See Note [Unsound mapUniqSet].
+unsafeUFMToUniqSet :: UniqFM a -> UniqSet a
+unsafeUFMToUniqSet = UniqSet
+
+instance Outputable a => Outputable (UniqSet a) where
+    ppr = pprUniqSet ppr
+#if __GLASGOW_HASKELL__ >= 801
+instance Data.Semigroup.Semigroup (UniqSet a) where
+  (<>) = mappend
+#endif
+instance Monoid (UniqSet a) where
+  mempty = UniqSet mempty
+  UniqSet s `mappend` UniqSet t = UniqSet (s `mappend` t)
+
+pprUniqSet :: (a -> SDoc) -> UniqSet a -> SDoc
+pprUniqSet f (UniqSet s) = pprUniqFM f s
+
+emptyUniqSet = UniqSet emptyUFM
+unitUniqSet x = UniqSet $ unitUFM x x
 mkUniqSet = foldl' addOneToUniqSet emptyUniqSet
 
-addOneToUniqSet set x = addToUFM set x x
-addOneToUniqSet_C f set x = addToUFM_C f set x x
+addOneToUniqSet (UniqSet set) x = UniqSet (addToUFM set x x)
 addListToUniqSet = foldl' addOneToUniqSet
 
-delOneFromUniqSet = delFromUFM
-delOneFromUniqSet_Directly = delFromUFM_Directly
-delListFromUniqSet = delListFromUFM
+delOneFromUniqSet (UniqSet s) a = UniqSet (delFromUFM s a)
+delOneFromUniqSet_Directly (UniqSet s) u = UniqSet (delFromUFM_Directly s u)
+delListFromUniqSet (UniqSet s) l = UniqSet (delListFromUFM s l)
+delListFromUniqSet_Directly (UniqSet s) l =
+    UniqSet (delListFromUFM_Directly s l)
+
+unionUniqSets (UniqSet s) (UniqSet t) = UniqSet (plusUFM s t)
 
-unionUniqSets = plusUFM
 unionManyUniqSets = foldl' (flip unionUniqSets) emptyUniqSet
-minusUniqSet = minusUFM
-intersectUniqSets = intersectUFM
 
-elementOfUniqSet = elemUFM
-elemUniqSet_Directly = elemUFM_Directly
-filterUniqSet = filterUFM
-partitionUniqSet = partitionUFM
+minusUniqSet (UniqSet s) (UniqSet t) = UniqSet (minusUFM s t)
+uniqSetMinusUFM (UniqSet s) t = UniqSet (minusUFM s t)
 
-sizeUniqSet = sizeUFM
-isEmptyUniqSet = isNullUFM
-lookupUniqSet = lookupUFM
+
+intersectUniqSets (UniqSet s) (UniqSet t) = UniqSet (intersectUFM s t)
+restrictUniqSetToUFM (UniqSet s) m = UniqSet (intersectUFM s m)
+
+elementOfUniqSet a (UniqSet s) = elemUFM a s
+elemUniqSet_Directly a (UniqSet s) = elemUFM_Directly a s
+filterUniqSet p (UniqSet s) = UniqSet (filterUFM p s)
+filterUniqSet_Directly f (UniqSet s) = UniqSet (filterUFM_Directly f s)
+
+partitionUniqSet p (UniqSet s) = coerce (partitionUFM p s)
+
+sizeUniqSet (UniqSet s) = sizeUFM s
+isEmptyUniqSet (UniqSet s) = isNullUFM s
+lookupUniqSet (UniqSet s) k = lookupUFM s k
+lookupUniqSet_Directly (UniqSet s) k = lookupUFM_Directly s k
 
 uniqSetAny :: (a -> Bool) -> UniqSet a -> Bool
-uniqSetAny = anyUFM
+uniqSetAny p (UniqSet s) = anyUFM p s
 
 uniqSetAll :: (a -> Bool) -> UniqSet a -> Bool
-uniqSetAll = allUFM
+uniqSetAll p (UniqSet s) = allUFM p s
+
+nonDetFoldUniqSet c n (UniqSet s) = nonDetFoldUFM c n s
+nonDetFoldUniqSet_Directly f n (UniqSet s) = nonDetFoldUFM_Directly f n s
+nonDetEltsUniqSet = nonDetEltsUFM . getUniqSet'
+nonDetKeysUniqSet = nonDetKeysUFM . getUniqSet'
+
+mapUniqSet f = mkUniqSet . map f . nonDetEltsUniqSet