Modern type signature style in UniqSet

2 files changed, 74 insertions, 79 deletions

diff --git a/compiler/basicTypes/VarSet.hs b/compiler/basicTypes/VarSet.hs
index e4f0d25dfb..710cb0db3a 100644
--- a/compiler/basicTypes/VarSet.hs
+++ b/compiler/basicTypes/VarSet.hs
@@ -16,7 +16,7 @@ module VarSet (
         unionVarSet, unionVarSets, mapUnionVarSet,
         intersectVarSet, intersectsVarSet, disjointVarSet,
         isEmptyVarSet, delVarSet, delVarSetList, delVarSetByKey,
-        minusVarSet, filterVarSet,
+        minusVarSet, filterVarSet, mapVarSet,
         anyVarSet, allVarSet,
         transCloVarSet, fixVarSet,
         lookupVarSet_Directly, lookupVarSet, lookupVarSetByName,
@@ -146,8 +146,8 @@ anyVarSet = uniqSetAny
 allVarSet :: (Var -> Bool) -> VarSet -> Bool
 allVarSet = uniqSetAll
 
--- There used to exist mapVarSet, see Note [Unsound mapUniqSet] in UniqSet for
--- why it got removed.
+mapVarSet :: Uniquable b => (a -> b) -> UniqSet a -> UniqSet b
+mapVarSet = mapUniqSet
 
 fixVarSet :: (VarSet -> VarSet)   -- Map the current set to a new set
           -> VarSet -> VarSet
diff --git a/compiler/utils/UniqSet.hs b/compiler/utils/UniqSet.hs
index d9d51f4c75..f29a1e6e1f 100644
--- a/compiler/utils/UniqSet.hs
+++ b/compiler/utils/UniqSet.hs
@@ -57,77 +57,121 @@ import Data.Data
 import qualified Data.Semigroup
 #endif
 
-{-
-************************************************************************
-*                                                                      *
-\subsection{The signature of the module}
-*                                                                      *
-************************************************************************
--}
+-- Note [UniqSet invariant]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~
+-- UniqSet has the following invariant:
+--   The keys in the map are the uniques of the values
+-- It means that to implement mapUniqSet you have to update
+-- both the keys and the values.
+
+newtype UniqSet a = UniqSet {getUniqSet' :: UniqFM a} deriving Data
 
 emptyUniqSet :: UniqSet a
+emptyUniqSet = UniqSet emptyUFM
+
 unitUniqSet :: Uniquable a => a -> UniqSet a
+unitUniqSet x = UniqSet $ unitUFM x x
+
 mkUniqSet :: Uniquable a => [a]  -> UniqSet a
+mkUniqSet = foldl' addOneToUniqSet emptyUniqSet
 
 addOneToUniqSet :: Uniquable a => UniqSet a -> a -> UniqSet a
+addOneToUniqSet (UniqSet set) x = UniqSet (addToUFM set x x)
+
 addListToUniqSet :: Uniquable a => UniqSet a -> [a] -> UniqSet a
+addListToUniqSet = foldl' addOneToUniqSet
 
 delOneFromUniqSet :: Uniquable a => UniqSet a -> a -> UniqSet a
+delOneFromUniqSet (UniqSet s) a = UniqSet (delFromUFM s a)
+
 delOneFromUniqSet_Directly :: UniqSet a -> Unique -> UniqSet a
+delOneFromUniqSet_Directly (UniqSet s) u = UniqSet (delFromUFM_Directly s u)
+
 delListFromUniqSet :: Uniquable a => UniqSet a -> [a] -> UniqSet a
+delListFromUniqSet (UniqSet s) l = UniqSet (delListFromUFM s l)
+
 delListFromUniqSet_Directly :: UniqSet a -> [Unique] -> UniqSet a
+delListFromUniqSet_Directly (UniqSet s) l =
+    UniqSet (delListFromUFM_Directly s l)
 
 unionUniqSets :: UniqSet a -> UniqSet a -> UniqSet a
+unionUniqSets (UniqSet s) (UniqSet t) = UniqSet (plusUFM s t)
+
 unionManyUniqSets :: [UniqSet a] -> UniqSet a
+unionManyUniqSets = foldl' (flip unionUniqSets) emptyUniqSet
+
 minusUniqSet  :: UniqSet a -> UniqSet a -> UniqSet a
+minusUniqSet (UniqSet s) (UniqSet t) = UniqSet (minusUFM s t)
+
 intersectUniqSets :: UniqSet a -> UniqSet a -> UniqSet a
+intersectUniqSets (UniqSet s) (UniqSet t) = UniqSet (intersectUFM s t)
+
 restrictUniqSetToUFM :: UniqSet a -> UniqFM b -> UniqSet a
+restrictUniqSetToUFM (UniqSet s) m = UniqSet (intersectUFM s m)
+
 uniqSetMinusUFM :: UniqSet a -> UniqFM b -> UniqSet a
+uniqSetMinusUFM (UniqSet s) t = UniqSet (minusUFM s t)
 
 elementOfUniqSet :: Uniquable a => a -> UniqSet a -> Bool
+elementOfUniqSet a (UniqSet s) = elemUFM a s
+
 elemUniqSet_Directly :: Unique -> UniqSet a -> Bool
+elemUniqSet_Directly a (UniqSet s) = elemUFM_Directly a s
+
 filterUniqSet :: (a -> Bool) -> UniqSet a -> UniqSet a
+filterUniqSet p (UniqSet s) = UniqSet (filterUFM p s)
+
 filterUniqSet_Directly :: (Unique -> elt -> Bool) -> UniqSet elt -> UniqSet elt
+filterUniqSet_Directly f (UniqSet s) = UniqSet (filterUFM_Directly f s)
+
 partitionUniqSet :: (a -> Bool) -> UniqSet a -> (UniqSet a, UniqSet a)
+partitionUniqSet p (UniqSet s) = coerce (partitionUFM p s)
+
+uniqSetAny :: (a -> Bool) -> UniqSet a -> Bool
+uniqSetAny p (UniqSet s) = anyUFM p s
+
+uniqSetAll :: (a -> Bool) -> UniqSet a -> Bool
+uniqSetAll p (UniqSet s) = allUFM p s
 
 sizeUniqSet :: UniqSet a -> Int
+sizeUniqSet (UniqSet s) = sizeUFM s
+
 isEmptyUniqSet :: UniqSet a -> Bool
+isEmptyUniqSet (UniqSet s) = isNullUFM s
+
 lookupUniqSet :: Uniquable a => UniqSet b -> a -> Maybe b
+lookupUniqSet (UniqSet s) k = lookupUFM s k
+
 lookupUniqSet_Directly :: UniqSet a -> Unique -> Maybe a
+lookupUniqSet_Directly (UniqSet s) k = lookupUFM_Directly s k
 
+-- See Note [Deterministic UniqFM] to learn about nondeterminism.
+-- If you use this please provide a justification why it doesn't introduce
+-- nondeterminism.
 nonDetEltsUniqSet :: UniqSet elt -> [elt]
+nonDetEltsUniqSet = nonDetEltsUFM . getUniqSet'
+
+-- See Note [Deterministic UniqFM] to learn about nondeterminism.
+-- If you use this please provide a justification why it doesn't introduce
+-- nondeterminism.
 nonDetKeysUniqSet :: UniqSet elt -> [Unique]
+nonDetKeysUniqSet = nonDetKeysUFM . getUniqSet'
 
 -- 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
+nonDetFoldUniqSet c n (UniqSet s) = nonDetFoldUFM c n s
 
 -- 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
+nonDetFoldUniqSet_Directly f n (UniqSet s) = nonDetFoldUFM_Directly f n s
 
+-- See Note [UniqSet invariant]
 mapUniqSet :: Uniquable b => (a -> b) -> UniqSet a -> UniqSet b
-
-{-
-************************************************************************
-*                                                                      *
-\subsection{Implementation using ``UniqFM''}
-*                                                                      *
-************************************************************************
--}
-
--- Note [Unsound mapUniqSet]
--- ~~~~~~~~~~~~~~~~~~~~~~~~~
--- UniqSet has the following invariant:
---   The keys in the map are the uniques of the values
--- It means that to implement mapUniqSet you'd have to update
--- both the keys and the values. There used to be an implementation
--- that only updated the values and it's been removed, because it broke
--- the invariant.
-
-newtype UniqSet a = UniqSet {getUniqSet' :: UniqFM a} deriving Data
+mapUniqSet f = mkUniqSet . map f . nonDetEltsUniqSet
 
 -- Two 'UniqSet's are considered equal if they contain the same
 -- uniques.
@@ -139,7 +183,7 @@ 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].
+-- that has that 'Unique'. See Note [UniqSet invariant].
 unsafeUFMToUniqSet :: UniqFM a -> UniqSet a
 unsafeUFMToUniqSet = UniqSet
 
@@ -155,52 +199,3 @@ instance Monoid (UniqSet a) where
 
 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 (UniqSet set) x = UniqSet (addToUFM set x x)
-addListToUniqSet = foldl' addOneToUniqSet
-
-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)
-
-unionManyUniqSets = foldl' (flip unionUniqSets) emptyUniqSet
-
-minusUniqSet (UniqSet s) (UniqSet t) = UniqSet (minusUFM s t)
-uniqSetMinusUFM (UniqSet s) t = UniqSet (minusUFM s t)
-
-
-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 p (UniqSet s) = anyUFM p s
-
-uniqSetAll :: (a -> Bool) -> UniqSet a -> Bool
-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