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