Add an Eq instance for UniqSet
[ghc.git] / compiler / utils / UniqFM.hs
1 {-
2 (c) The University of Glasgow 2006
3 (c) The AQUA Project, Glasgow University, 1994-1998
4
5
6 UniqFM: Specialised finite maps, for things with @Uniques@.
7
8 Basically, the things need to be in class @Uniquable@, and we use the
9 @getUnique@ method to grab their @Uniques@.
10
11 (A similar thing to @UniqSet@, as opposed to @Set@.)
12
13 The interface is based on @FiniteMap@s, but the implementation uses
14 @Data.IntMap@, which is both maintained and faster than the past
15 implementation (see commit log).
16
17 The @UniqFM@ interface maps directly to Data.IntMap, only
18 ``Data.IntMap.union'' is left-biased and ``plusUFM'' right-biased
19 and ``addToUFM\_C'' and ``Data.IntMap.insertWith'' differ in the order
20 of arguments of combining function.
21 -}
22
23 {-# LANGUAGE CPP #-}
24 {-# LANGUAGE DeriveDataTypeable #-}
25 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
26 {-# OPTIONS_GHC -Wall #-}
27
28 module UniqFM (
29 -- * Unique-keyed mappings
30 UniqFM, -- abstract type
31
32 -- ** Manipulating those mappings
33 emptyUFM,
34 unitUFM,
35 unitDirectlyUFM,
36 listToUFM,
37 listToUFM_Directly,
38 listToUFM_C,
39 addToUFM,addToUFM_C,addToUFM_Acc,
40 addListToUFM,addListToUFM_C,
41 addToUFM_Directly,
42 addListToUFM_Directly,
43 adjustUFM, alterUFM,
44 adjustUFM_Directly,
45 delFromUFM,
46 delFromUFM_Directly,
47 delListFromUFM,
48 delListFromUFM_Directly,
49 plusUFM,
50 plusUFM_C,
51 plusUFM_CD,
52 plusMaybeUFM_C,
53 plusUFMList,
54 minusUFM,
55 intersectUFM,
56 intersectUFM_C,
57 disjointUFM,
58 equalKeysUFM,
59 nonDetFoldUFM, foldUFM, nonDetFoldUFM_Directly,
60 anyUFM, allUFM, seqEltsUFM,
61 mapUFM, mapUFM_Directly,
62 elemUFM, elemUFM_Directly,
63 filterUFM, filterUFM_Directly, partitionUFM,
64 sizeUFM,
65 isNullUFM,
66 lookupUFM, lookupUFM_Directly,
67 lookupWithDefaultUFM, lookupWithDefaultUFM_Directly,
68 nonDetEltsUFM, eltsUFM, nonDetKeysUFM,
69 ufmToSet_Directly,
70 nonDetUFMToList, ufmToIntMap,
71 pprUniqFM, pprUFM, pprUFMWithKeys, pluralUFM
72 ) where
73
74 import Unique ( Uniquable(..), Unique, getKey )
75 import Outputable
76
77 import Data.List (foldl')
78
79 import qualified Data.IntMap as M
80 #if MIN_VERSION_containers(0,5,9)
81 import qualified Data.IntMap.Merge.Lazy as M
82 import Control.Applicative (Const (..))
83 import qualified Data.Monoid as Mon
84 #endif
85 import qualified Data.IntSet as S
86 import Data.Typeable
87 import Data.Data
88 #if __GLASGOW_HASKELL__ > 710
89 import Data.Semigroup ( Semigroup )
90 import qualified Data.Semigroup as Semigroup
91 #endif
92
93
94 newtype UniqFM ele = UFM (M.IntMap ele)
95 deriving (Data, Eq, Functor, Typeable)
96 -- We used to derive Traversable and Foldable, but they were nondeterministic
97 -- and not obvious at the call site. You can use explicit nonDetEltsUFM
98 -- and fold a list if needed.
99 -- See Note [Deterministic UniqFM] in UniqDFM to learn about determinism.
100
101 emptyUFM :: UniqFM elt
102 emptyUFM = UFM M.empty
103
104 isNullUFM :: UniqFM elt -> Bool
105 isNullUFM (UFM m) = M.null m
106
107 unitUFM :: Uniquable key => key -> elt -> UniqFM elt
108 unitUFM k v = UFM (M.singleton (getKey $ getUnique k) v)
109
110 -- when you've got the Unique already
111 unitDirectlyUFM :: Unique -> elt -> UniqFM elt
112 unitDirectlyUFM u v = UFM (M.singleton (getKey u) v)
113
114 listToUFM :: Uniquable key => [(key,elt)] -> UniqFM elt
115 listToUFM = foldl (\m (k, v) -> addToUFM m k v) emptyUFM
116
117 listToUFM_Directly :: [(Unique, elt)] -> UniqFM elt
118 listToUFM_Directly = foldl (\m (u, v) -> addToUFM_Directly m u v) emptyUFM
119
120 listToUFM_C
121 :: Uniquable key
122 => (elt -> elt -> elt)
123 -> [(key, elt)]
124 -> UniqFM elt
125 listToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v) emptyUFM
126
127 addToUFM :: Uniquable key => UniqFM elt -> key -> elt -> UniqFM elt
128 addToUFM (UFM m) k v = UFM (M.insert (getKey $ getUnique k) v m)
129
130 addListToUFM :: Uniquable key => UniqFM elt -> [(key,elt)] -> UniqFM elt
131 addListToUFM = foldl (\m (k, v) -> addToUFM m k v)
132
133 addListToUFM_Directly :: UniqFM elt -> [(Unique,elt)] -> UniqFM elt
134 addListToUFM_Directly = foldl (\m (k, v) -> addToUFM_Directly m k v)
135
136 addToUFM_Directly :: UniqFM elt -> Unique -> elt -> UniqFM elt
137 addToUFM_Directly (UFM m) u v = UFM (M.insert (getKey u) v m)
138
139 addToUFM_C
140 :: Uniquable key
141 => (elt -> elt -> elt) -- old -> new -> result
142 -> UniqFM elt -- old
143 -> key -> elt -- new
144 -> UniqFM elt -- result
145 -- Arguments of combining function of M.insertWith and addToUFM_C are flipped.
146 addToUFM_C f (UFM m) k v =
147 UFM (M.insertWith (flip f) (getKey $ getUnique k) v m)
148
149 addToUFM_Acc
150 :: Uniquable key
151 => (elt -> elts -> elts) -- Add to existing
152 -> (elt -> elts) -- New element
153 -> UniqFM elts -- old
154 -> key -> elt -- new
155 -> UniqFM elts -- result
156 addToUFM_Acc exi new (UFM m) k v =
157 UFM (M.insertWith (\_new old -> exi v old) (getKey $ getUnique k) (new v) m)
158
159 alterUFM
160 :: Uniquable key
161 => (Maybe elt -> Maybe elt) -- How to adjust
162 -> UniqFM elt -- old
163 -> key -- new
164 -> UniqFM elt -- result
165 alterUFM f (UFM m) k = UFM (M.alter f (getKey $ getUnique k) m)
166
167 addListToUFM_C
168 :: Uniquable key
169 => (elt -> elt -> elt)
170 -> UniqFM elt -> [(key,elt)]
171 -> UniqFM elt
172 addListToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v)
173
174 adjustUFM :: Uniquable key => (elt -> elt) -> UniqFM elt -> key -> UniqFM elt
175 adjustUFM f (UFM m) k = UFM (M.adjust f (getKey $ getUnique k) m)
176
177 adjustUFM_Directly :: (elt -> elt) -> UniqFM elt -> Unique -> UniqFM elt
178 adjustUFM_Directly f (UFM m) u = UFM (M.adjust f (getKey u) m)
179
180 delFromUFM :: Uniquable key => UniqFM elt -> key -> UniqFM elt
181 delFromUFM (UFM m) k = UFM (M.delete (getKey $ getUnique k) m)
182
183 delListFromUFM :: Uniquable key => UniqFM elt -> [key] -> UniqFM elt
184 delListFromUFM = foldl delFromUFM
185
186 delListFromUFM_Directly :: UniqFM elt -> [Unique] -> UniqFM elt
187 delListFromUFM_Directly = foldl delFromUFM_Directly
188
189 delFromUFM_Directly :: UniqFM elt -> Unique -> UniqFM elt
190 delFromUFM_Directly (UFM m) u = UFM (M.delete (getKey u) m)
191
192 -- Bindings in right argument shadow those in the left
193 plusUFM :: UniqFM elt -> UniqFM elt -> UniqFM elt
194 -- M.union is left-biased, plusUFM should be right-biased.
195 plusUFM (UFM x) (UFM y) = UFM (M.union y x)
196 -- Note (M.union y x), with arguments flipped
197 -- M.union is left-biased, plusUFM should be right-biased.
198
199 plusUFM_C :: (elt -> elt -> elt) -> UniqFM elt -> UniqFM elt -> UniqFM elt
200 plusUFM_C f (UFM x) (UFM y) = UFM (M.unionWith f x y)
201
202 -- | `plusUFM_CD f m1 d1 m2 d2` merges the maps using `f` as the
203 -- combinding function and `d1` resp. `d2` as the default value if
204 -- there is no entry in `m1` reps. `m2`. The domain is the union of
205 -- the domains of `m1` and `m2`.
206 --
207 -- Representative example:
208 --
209 -- @
210 -- plusUFM_CD f {A: 1, B: 2} 23 {B: 3, C: 4} 42
211 -- == {A: f 1 42, B: f 2 3, C: f 23 4 }
212 -- @
213 plusUFM_CD
214 :: (elt -> elt -> elt)
215 -> UniqFM elt -- map X
216 -> elt -- default for X
217 -> UniqFM elt -- map Y
218 -> elt -- default for Y
219 -> UniqFM elt
220 plusUFM_CD f (UFM xm) dx (UFM ym) dy
221 = UFM $ M.mergeWithKey
222 (\_ x y -> Just (x `f` y))
223 (M.map (\x -> x `f` dy))
224 (M.map (\y -> dx `f` y))
225 xm ym
226
227 plusMaybeUFM_C :: (elt -> elt -> Maybe elt)
228 -> UniqFM elt -> UniqFM elt -> UniqFM elt
229 plusMaybeUFM_C f (UFM xm) (UFM ym)
230 = UFM $ M.mergeWithKey
231 (\_ x y -> x `f` y)
232 id
233 id
234 xm ym
235
236 plusUFMList :: [UniqFM elt] -> UniqFM elt
237 plusUFMList = foldl' plusUFM emptyUFM
238
239 minusUFM :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
240 minusUFM (UFM x) (UFM y) = UFM (M.difference x y)
241
242 intersectUFM :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
243 intersectUFM (UFM x) (UFM y) = UFM (M.intersection x y)
244
245 intersectUFM_C
246 :: (elt1 -> elt2 -> elt3)
247 -> UniqFM elt1
248 -> UniqFM elt2
249 -> UniqFM elt3
250 intersectUFM_C f (UFM x) (UFM y) = UFM (M.intersectionWith f x y)
251
252 disjointUFM :: UniqFM elt1 -> UniqFM elt2 -> Bool
253 disjointUFM (UFM x) (UFM y) = M.null (M.intersection x y)
254
255 foldUFM :: (elt -> a -> a) -> a -> UniqFM elt -> a
256 foldUFM k z (UFM m) = M.foldr k z m
257
258 mapUFM :: (elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
259 mapUFM f (UFM m) = UFM (M.map f m)
260
261 mapUFM_Directly :: (Unique -> elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
262 mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . getUnique) m)
263
264 filterUFM :: (elt -> Bool) -> UniqFM elt -> UniqFM elt
265 filterUFM p (UFM m) = UFM (M.filter p m)
266
267 filterUFM_Directly :: (Unique -> elt -> Bool) -> UniqFM elt -> UniqFM elt
268 filterUFM_Directly p (UFM m) = UFM (M.filterWithKey (p . getUnique) m)
269
270 partitionUFM :: (elt -> Bool) -> UniqFM elt -> (UniqFM elt, UniqFM elt)
271 partitionUFM p (UFM m) =
272 case M.partition p m of
273 (left, right) -> (UFM left, UFM right)
274
275 sizeUFM :: UniqFM elt -> Int
276 sizeUFM (UFM m) = M.size m
277
278 elemUFM :: Uniquable key => key -> UniqFM elt -> Bool
279 elemUFM k (UFM m) = M.member (getKey $ getUnique k) m
280
281 elemUFM_Directly :: Unique -> UniqFM elt -> Bool
282 elemUFM_Directly u (UFM m) = M.member (getKey u) m
283
284 lookupUFM :: Uniquable key => UniqFM elt -> key -> Maybe elt
285 lookupUFM (UFM m) k = M.lookup (getKey $ getUnique k) m
286
287 -- when you've got the Unique already
288 lookupUFM_Directly :: UniqFM elt -> Unique -> Maybe elt
289 lookupUFM_Directly (UFM m) u = M.lookup (getKey u) m
290
291 lookupWithDefaultUFM :: Uniquable key => UniqFM elt -> elt -> key -> elt
292 lookupWithDefaultUFM (UFM m) v k = M.findWithDefault v (getKey $ getUnique k) m
293
294 lookupWithDefaultUFM_Directly :: UniqFM elt -> elt -> Unique -> elt
295 lookupWithDefaultUFM_Directly (UFM m) v u = M.findWithDefault v (getKey u) m
296
297 eltsUFM :: UniqFM elt -> [elt]
298 eltsUFM (UFM m) = M.elems m
299
300 ufmToSet_Directly :: UniqFM elt -> S.IntSet
301 ufmToSet_Directly (UFM m) = M.keysSet m
302
303 anyUFM :: (elt -> Bool) -> UniqFM elt -> Bool
304 anyUFM p (UFM m) = M.foldr ((||) . p) False m
305
306 allUFM :: (elt -> Bool) -> UniqFM elt -> Bool
307 allUFM p (UFM m) = M.foldr ((&&) . p) True m
308
309 seqEltsUFM :: ([elt] -> ()) -> UniqFM elt -> ()
310 seqEltsUFM seqList = seqList . nonDetEltsUFM
311 -- It's OK to use nonDetEltsUFM here because the type guarantees that
312 -- the only interesting thing this function can do is to force the
313 -- elements.
314
315 -- See Note [Deterministic UniqFM] to learn about nondeterminism.
316 -- If you use this please provide a justification why it doesn't introduce
317 -- nondeterminism.
318 nonDetEltsUFM :: UniqFM elt -> [elt]
319 nonDetEltsUFM (UFM m) = M.elems m
320
321 -- See Note [Deterministic UniqFM] to learn about nondeterminism.
322 -- If you use this please provide a justification why it doesn't introduce
323 -- nondeterminism.
324 nonDetKeysUFM :: UniqFM elt -> [Unique]
325 nonDetKeysUFM (UFM m) = map getUnique $ M.keys m
326
327 -- See Note [Deterministic UniqFM] to learn about nondeterminism.
328 -- If you use this please provide a justification why it doesn't introduce
329 -- nondeterminism.
330 nonDetFoldUFM :: (elt -> a -> a) -> a -> UniqFM elt -> a
331 nonDetFoldUFM k z (UFM m) = M.foldr k z m
332
333 -- See Note [Deterministic UniqFM] to learn about nondeterminism.
334 -- If you use this please provide a justification why it doesn't introduce
335 -- nondeterminism.
336 nonDetFoldUFM_Directly:: (Unique -> elt -> a -> a) -> a -> UniqFM elt -> a
337 nonDetFoldUFM_Directly k z (UFM m) = M.foldrWithKey (k . getUnique) z m
338
339 -- See Note [Deterministic UniqFM] to learn about nondeterminism.
340 -- If you use this please provide a justification why it doesn't introduce
341 -- nondeterminism.
342 nonDetUFMToList :: UniqFM elt -> [(Unique, elt)]
343 nonDetUFMToList (UFM m) = map (\(k, v) -> (getUnique k, v)) $ M.toList m
344
345 ufmToIntMap :: UniqFM elt -> M.IntMap elt
346 ufmToIntMap (UFM m) = m
347
348 -- Determines whether two 'UniqFm's contain the same keys.
349 equalKeysUFM :: UniqFM a -> UniqFM b -> Bool
350 #if MIN_VERSION_containers(0,5,9)
351 equalKeysUFM (UFM m1) (UFM m2) = Mon.getAll $ getConst $
352 M.mergeA (M.traverseMissing (\_ _ -> Const (Mon.All False)))
353 (M.traverseMissing (\_ _ -> Const (Mon.All False)))
354 (M.zipWithAMatched (\_ _ _ -> Const (Mon.All True))) m1 m2
355 #else
356 equalKeysUFM (UFM m1) (UFM m2) = M.keys m1 == M.keys m2
357 #endif
358
359 -- Instances
360
361 #if __GLASGOW_HASKELL__ > 710
362 instance Semigroup (UniqFM a) where
363 (<>) = plusUFM
364 #endif
365
366 instance Monoid (UniqFM a) where
367 mempty = emptyUFM
368 mappend = plusUFM
369
370 -- Output-ery
371
372 instance Outputable a => Outputable (UniqFM a) where
373 ppr ufm = pprUniqFM ppr ufm
374
375 pprUniqFM :: (a -> SDoc) -> UniqFM a -> SDoc
376 pprUniqFM ppr_elt ufm
377 = brackets $ fsep $ punctuate comma $
378 [ ppr uq <+> text ":->" <+> ppr_elt elt
379 | (uq, elt) <- nonDetUFMToList ufm ]
380 -- It's OK to use nonDetUFMToList here because we only use it for
381 -- pretty-printing.
382
383 -- | Pretty-print a non-deterministic set.
384 -- The order of variables is non-deterministic and for pretty-printing that
385 -- shouldn't be a problem.
386 -- Having this function helps contain the non-determinism created with
387 -- nonDetEltsUFM.
388 pprUFM :: UniqFM a -- ^ The things to be pretty printed
389 -> ([a] -> SDoc) -- ^ The pretty printing function to use on the elements
390 -> SDoc -- ^ 'SDoc' where the things have been pretty
391 -- printed
392 pprUFM ufm pp = pp (nonDetEltsUFM ufm)
393
394 -- | Pretty-print a non-deterministic set.
395 -- The order of variables is non-deterministic and for pretty-printing that
396 -- shouldn't be a problem.
397 -- Having this function helps contain the non-determinism created with
398 -- nonDetUFMToList.
399 pprUFMWithKeys
400 :: UniqFM a -- ^ The things to be pretty printed
401 -> ([(Unique, a)] -> SDoc) -- ^ The pretty printing function to use on the elements
402 -> SDoc -- ^ 'SDoc' where the things have been pretty
403 -- printed
404 pprUFMWithKeys ufm pp = pp (nonDetUFMToList ufm)
405
406 -- | Determines the pluralisation suffix appropriate for the length of a set
407 -- in the same way that plural from Outputable does for lists.
408 pluralUFM :: UniqFM a -> SDoc
409 pluralUFM ufm
410 | sizeUFM ufm == 1 = empty
411 | otherwise = char 's'