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