1{-
2(c) The University of Glasgow 2006
3(c) The AQUA Project, Glasgow University, 1994-1998
4
5\section[UniqSet]{Specialised sets, for things with @Uniques@}
6
7Based on @UniqFMs@ (as you would expect).
8
9Basically, the things need to be in class @Uniquable@.
10-}
11{-# LANGUAGE GeneralizedNewtypeDeriving #-}
12{-# LANGUAGE DeriveDataTypeable #-}
13
14module UniqSet (
15        -- * Unique set type
16        UniqSet,    -- type synonym for UniqFM a
17        getUniqSet,
18        pprUniqSet,
19
20        -- ** Manipulating these sets
21        emptyUniqSet,
22        unitUniqSet,
23        mkUniqSet,
24        addOneToUniqSet, addListToUniqSet,
25        delOneFromUniqSet, delOneFromUniqSet_Directly, delListFromUniqSet,
26        delListFromUniqSet_Directly,
27        unionUniqSets, unionManyUniqSets,
28        minusUniqSet, uniqSetMinusUFM,
29        intersectUniqSets,
30        restrictUniqSetToUFM,
31        uniqSetAny, uniqSetAll,
32        elementOfUniqSet,
33        elemUniqSet_Directly,
34        filterUniqSet,
35        filterUniqSet_Directly,
36        sizeUniqSet,
37        isEmptyUniqSet,
38        lookupUniqSet,
39        lookupUniqSet_Directly,
40        partitionUniqSet,
41        mapUniqSet,
42        unsafeUFMToUniqSet,
43        nonDetEltsUniqSet,
44        nonDetKeysUniqSet,
45        nonDetFoldUniqSet,
46        nonDetFoldUniqSet_Directly
47    ) where
48
49import GhcPrelude
50
51import UniqFM
52import Unique
53import Data.Coerce
54import Outputable
55import Data.Data
56import qualified Data.Semigroup as Semi
57
58-- Note [UniqSet invariant]
59-- ~~~~~~~~~~~~~~~~~~~~~~~~~
60-- UniqSet has the following invariant:
61--   The keys in the map are the uniques of the values
62-- It means that to implement mapUniqSet you have to update
63-- both the keys and the values.
64
65newtype UniqSet a = UniqSet {getUniqSet' :: UniqFM a}
66                  deriving (Data, Semi.Semigroup, Monoid)
67
68emptyUniqSet :: UniqSet a
69emptyUniqSet = UniqSet emptyUFM
70
71unitUniqSet :: Uniquable a => a -> UniqSet a
72unitUniqSet x = UniqSet $ unitUFM x x
73
74mkUniqSet :: Uniquable a => [a]  -> UniqSet a
75mkUniqSet = foldl' addOneToUniqSet emptyUniqSet
76
77addOneToUniqSet :: Uniquable a => UniqSet a -> a -> UniqSet a
78addOneToUniqSet (UniqSet set) x = UniqSet (addToUFM set x x)
79
80addListToUniqSet :: Uniquable a => UniqSet a -> [a] -> UniqSet a
81addListToUniqSet = foldl' addOneToUniqSet
82
83delOneFromUniqSet :: Uniquable a => UniqSet a -> a -> UniqSet a
84delOneFromUniqSet (UniqSet s) a = UniqSet (delFromUFM s a)
85
86delOneFromUniqSet_Directly :: UniqSet a -> Unique -> UniqSet a
87delOneFromUniqSet_Directly (UniqSet s) u = UniqSet (delFromUFM_Directly s u)
88
89delListFromUniqSet :: Uniquable a => UniqSet a -> [a] -> UniqSet a
90delListFromUniqSet (UniqSet s) l = UniqSet (delListFromUFM s l)
91
92delListFromUniqSet_Directly :: UniqSet a -> [Unique] -> UniqSet a
93delListFromUniqSet_Directly (UniqSet s) l =
94    UniqSet (delListFromUFM_Directly s l)
95
96unionUniqSets :: UniqSet a -> UniqSet a -> UniqSet a
97unionUniqSets (UniqSet s) (UniqSet t) = UniqSet (plusUFM s t)
98
99unionManyUniqSets :: [UniqSet a] -> UniqSet a
100unionManyUniqSets = foldl' (flip unionUniqSets) emptyUniqSet
101
102minusUniqSet  :: UniqSet a -> UniqSet a -> UniqSet a
103minusUniqSet (UniqSet s) (UniqSet t) = UniqSet (minusUFM s t)
104
105intersectUniqSets :: UniqSet a -> UniqSet a -> UniqSet a
106intersectUniqSets (UniqSet s) (UniqSet t) = UniqSet (intersectUFM s t)
107
108restrictUniqSetToUFM :: UniqSet a -> UniqFM b -> UniqSet a
109restrictUniqSetToUFM (UniqSet s) m = UniqSet (intersectUFM s m)
110
111uniqSetMinusUFM :: UniqSet a -> UniqFM b -> UniqSet a
112uniqSetMinusUFM (UniqSet s) t = UniqSet (minusUFM s t)
113
114elementOfUniqSet :: Uniquable a => a -> UniqSet a -> Bool
115elementOfUniqSet a (UniqSet s) = elemUFM a s
116
117elemUniqSet_Directly :: Unique -> UniqSet a -> Bool
118elemUniqSet_Directly a (UniqSet s) = elemUFM_Directly a s
119
120filterUniqSet :: (a -> Bool) -> UniqSet a -> UniqSet a
121filterUniqSet p (UniqSet s) = UniqSet (filterUFM p s)
122
123filterUniqSet_Directly :: (Unique -> elt -> Bool) -> UniqSet elt -> UniqSet elt
124filterUniqSet_Directly f (UniqSet s) = UniqSet (filterUFM_Directly f s)
125
126partitionUniqSet :: (a -> Bool) -> UniqSet a -> (UniqSet a, UniqSet a)
127partitionUniqSet p (UniqSet s) = coerce (partitionUFM p s)
128
129uniqSetAny :: (a -> Bool) -> UniqSet a -> Bool
130uniqSetAny p (UniqSet s) = anyUFM p s
131
132uniqSetAll :: (a -> Bool) -> UniqSet a -> Bool
133uniqSetAll p (UniqSet s) = allUFM p s
134
135sizeUniqSet :: UniqSet a -> Int
136sizeUniqSet (UniqSet s) = sizeUFM s
137
138isEmptyUniqSet :: UniqSet a -> Bool
139isEmptyUniqSet (UniqSet s) = isNullUFM s
140
141lookupUniqSet :: Uniquable a => UniqSet b -> a -> Maybe b
142lookupUniqSet (UniqSet s) k = lookupUFM s k
143
144lookupUniqSet_Directly :: UniqSet a -> Unique -> Maybe a
145lookupUniqSet_Directly (UniqSet s) k = lookupUFM_Directly s k
146
147-- See Note [Deterministic UniqFM] to learn about nondeterminism.
148-- If you use this please provide a justification why it doesn't introduce
149-- nondeterminism.
150nonDetEltsUniqSet :: UniqSet elt -> [elt]
151nonDetEltsUniqSet = nonDetEltsUFM . getUniqSet'
152
153-- See Note [Deterministic UniqFM] to learn about nondeterminism.
154-- If you use this please provide a justification why it doesn't introduce
155-- nondeterminism.
156nonDetKeysUniqSet :: UniqSet elt -> [Unique]
157nonDetKeysUniqSet = nonDetKeysUFM . getUniqSet'
158
159-- See Note [Deterministic UniqFM] to learn about nondeterminism.
160-- If you use this please provide a justification why it doesn't introduce
161-- nondeterminism.
162nonDetFoldUniqSet :: (elt -> a -> a) -> a -> UniqSet elt -> a
163nonDetFoldUniqSet c n (UniqSet s) = nonDetFoldUFM c n s
164
165-- See Note [Deterministic UniqFM] to learn about nondeterminism.
166-- If you use this please provide a justification why it doesn't introduce
167-- nondeterminism.
168nonDetFoldUniqSet_Directly:: (Unique -> elt -> a -> a) -> a -> UniqSet elt -> a
169nonDetFoldUniqSet_Directly f n (UniqSet s) = nonDetFoldUFM_Directly f n s
170
171-- See Note [UniqSet invariant]
172mapUniqSet :: Uniquable b => (a -> b) -> UniqSet a -> UniqSet b
173mapUniqSet f = mkUniqSet . map f . nonDetEltsUniqSet
174
175-- Two 'UniqSet's are considered equal if they contain the same
176-- uniques.
177instance Eq (UniqSet a) where
178  UniqSet a == UniqSet b = equalKeysUFM a b
179
180getUniqSet :: UniqSet a -> UniqFM a
181getUniqSet = getUniqSet'
182
183-- | 'unsafeUFMToUniqSet' converts a @'UniqFM' a@ into a @'UniqSet' a@
184-- assuming, without checking, that it maps each 'Unique' to a value
185-- that has that 'Unique'. See Note [UniqSet invariant].
186unsafeUFMToUniqSet :: UniqFM a -> UniqSet a
187unsafeUFMToUniqSet = UniqSet
188
189instance Outputable a => Outputable (UniqSet a) where
190    ppr = pprUniqSet ppr
191
192pprUniqSet :: (a -> SDoc) -> UniqSet a -> SDoc
193-- It's OK to use nonDetUFMToList here because we only use it for
194-- pretty-printing.
195pprUniqSet f = braces . pprWithCommas f . nonDetEltsUniqSet
196