d8e08f599ab10d7ff3d0cc9a7f0c3c421eb673a1
[ghc.git] / compiler / utils / UniqFM.lhs
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 \begin{code}
23 {-# LANGUAGE DeriveTraversable, DeriveDataTypeable, GeneralizedNewtypeDeriving #-}
24 {-# OPTIONS_GHC -Wall #-}
25
26 module UniqFM (
27         -- * Unique-keyed mappings
28         UniqFM,       -- abstract type
29
30         -- ** Manipulating those mappings
31         emptyUFM,
32         unitUFM,
33         unitDirectlyUFM,
34         listToUFM,
35         listToUFM_Directly,
36         listToUFM_C,
37         addToUFM,addToUFM_C,addToUFM_Acc,
38         addListToUFM,addListToUFM_C,
39         addToUFM_Directly,
40         addListToUFM_Directly,
41         adjustUFM, alterUFM,
42         adjustUFM_Directly,
43         delFromUFM,
44         delFromUFM_Directly,
45         delListFromUFM,
46         plusUFM,
47         plusUFM_C,
48         plusUFM_CD,
49         minusUFM,
50         intersectUFM,
51         intersectUFM_C,
52         foldUFM, foldUFM_Directly,
53         mapUFM, mapUFM_Directly,
54         elemUFM, elemUFM_Directly,
55         filterUFM, filterUFM_Directly, partitionUFM,
56         sizeUFM,
57         isNullUFM,
58         lookupUFM, lookupUFM_Directly,
59         lookupWithDefaultUFM, lookupWithDefaultUFM_Directly,
60         eltsUFM, keysUFM, splitUFM,
61         ufmToSet_Directly,
62         ufmToList,
63         joinUFM, pprUniqFM
64     ) where
65
66 import FastString
67 import Unique           ( Uniquable(..), Unique, getKey )
68 import Outputable
69
70 import Compiler.Hoopl   hiding (Unique)
71
72 import Data.Function (on)
73 import qualified Data.IntMap as M
74 import qualified Data.IntSet as S
75 import qualified Data.Foldable as Foldable
76 import qualified Data.Traversable as Traversable
77 import Data.Typeable
78 import Data.Data
79 import Data.Monoid
80 \end{code}
81
82 %************************************************************************
83 %*                                                                      *
84 \subsection{The signature of the module}
85 %*                                                                      *
86 %************************************************************************
87
88 \begin{code}
89 emptyUFM        :: UniqFM elt
90 isNullUFM       :: UniqFM elt -> Bool
91 unitUFM         :: Uniquable key => key -> elt -> UniqFM elt
92 unitDirectlyUFM -- got the Unique already
93                 :: Unique -> elt -> UniqFM elt
94 listToUFM       :: Uniquable key => [(key,elt)] -> UniqFM elt
95 listToUFM_Directly
96                 :: [(Unique, elt)] -> UniqFM elt
97 listToUFM_C     :: Uniquable key => (elt -> elt -> elt)
98                            -> [(key, elt)]
99                            -> UniqFM elt
100
101 addToUFM        :: Uniquable key => UniqFM elt -> key -> elt  -> UniqFM elt
102 addListToUFM    :: Uniquable key => UniqFM elt -> [(key,elt)] -> UniqFM elt
103 addListToUFM_Directly :: UniqFM elt -> [(Unique,elt)] -> UniqFM elt
104 addToUFM_Directly
105                 :: UniqFM elt -> Unique -> elt -> UniqFM elt
106
107 addToUFM_C      :: Uniquable key => (elt -> elt -> elt) -- old -> new -> result
108                            -> UniqFM elt                -- old
109                            -> key -> elt                -- new
110                            -> UniqFM elt                -- result
111
112 addToUFM_Acc    :: Uniquable key =>
113                               (elt -> elts -> elts)     -- Add to existing
114                            -> (elt -> elts)             -- New element
115                            -> UniqFM elts               -- old
116                            -> key -> elt                -- new
117                            -> UniqFM elts               -- result
118
119 alterUFM        :: Uniquable key =>
120                               (Maybe elt -> Maybe elt)  -- How to adjust
121                            -> UniqFM elt                -- old
122                            -> key                       -- new
123                            -> UniqFM elt                -- result
124
125 addListToUFM_C  :: Uniquable key => (elt -> elt -> elt)
126                            -> UniqFM elt -> [(key,elt)]
127                            -> UniqFM elt
128
129 adjustUFM       :: Uniquable key => (elt -> elt) -> UniqFM elt -> key -> UniqFM elt
130 adjustUFM_Directly :: (elt -> elt) -> UniqFM elt -> Unique -> UniqFM elt
131
132 delFromUFM      :: Uniquable key => UniqFM elt -> key    -> UniqFM elt
133 delListFromUFM  :: Uniquable key => UniqFM elt -> [key] -> UniqFM elt
134 delFromUFM_Directly :: UniqFM elt -> Unique -> UniqFM elt
135
136 -- Bindings in right argument shadow those in the left
137 plusUFM         :: UniqFM elt -> UniqFM elt -> UniqFM elt
138
139 plusUFM_C       :: (elt -> elt -> elt)
140                 -> UniqFM elt -> UniqFM elt -> UniqFM elt
141
142 -- | `plusUFM_CD f m1 d1 m2 d2` merges the maps using `f` as the
143 -- combinding function and `d1` resp. `d2` as the default value if
144 -- there is no entry in `m1` reps. `m2`. The domain is the union of
145 -- the domains of `m1` and `m2`.
146 --
147 -- Representative example:
148 --
149 -- @
150 -- plusUFM_CD f {A: 1, B: 2} 23 {B: 3, C: 4} 42
151 --    == {A: f 1 42, B: f 2 3, C: f 23 4 }
152 -- @
153 plusUFM_CD      :: (elt -> elt -> elt)
154                 -> UniqFM elt -> elt -> UniqFM elt -> elt -> UniqFM elt
155
156 minusUFM        :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
157
158 intersectUFM    :: UniqFM elt -> UniqFM elt -> UniqFM elt
159 intersectUFM_C  :: (elt1 -> elt2 -> elt3)
160                 -> UniqFM elt1 -> UniqFM elt2 -> UniqFM elt3
161
162 foldUFM         :: (elt -> a -> a) -> a -> UniqFM elt -> a
163 foldUFM_Directly:: (Unique -> elt -> a -> a) -> a -> UniqFM elt -> a
164 mapUFM          :: (elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
165 mapUFM_Directly :: (Unique -> elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
166 filterUFM       :: (elt -> Bool) -> UniqFM elt -> UniqFM elt
167 filterUFM_Directly :: (Unique -> elt -> Bool) -> UniqFM elt -> UniqFM elt
168 partitionUFM    :: (elt -> Bool) -> UniqFM elt -> (UniqFM elt, UniqFM elt)
169
170 sizeUFM         :: UniqFM elt -> Int
171 --hashUFM               :: UniqFM elt -> Int
172 elemUFM         :: Uniquable key => key -> UniqFM elt -> Bool
173 elemUFM_Directly:: Unique -> UniqFM elt -> Bool
174
175 splitUFM        :: Uniquable key => UniqFM elt -> key -> (UniqFM elt, Maybe elt, UniqFM elt)
176                    -- Splits a UFM into things less than, equal to, and greater than the key
177 lookupUFM       :: Uniquable key => UniqFM elt -> key -> Maybe elt
178 lookupUFM_Directly  -- when you've got the Unique already
179                 :: UniqFM elt -> Unique -> Maybe elt
180 lookupWithDefaultUFM
181                 :: Uniquable key => UniqFM elt -> elt -> key -> elt
182 lookupWithDefaultUFM_Directly
183                 :: UniqFM elt -> elt -> Unique -> elt
184 keysUFM         :: UniqFM elt -> [Unique]       -- Get the keys
185 eltsUFM         :: UniqFM elt -> [elt]
186 ufmToSet_Directly :: UniqFM elt -> S.IntSet
187 ufmToList       :: UniqFM elt -> [(Unique, elt)]
188
189 \end{code}
190
191 %************************************************************************
192 %*                                                                      *
193 \subsection{Monoid interface}
194 %*                                                                      *
195 %************************************************************************
196
197 \begin{code}
198 instance Monoid (UniqFM a) where
199     mempty = emptyUFM
200     mappend = plusUFM
201 \end{code}
202
203 %************************************************************************
204 %*                                                                      *
205 \subsection{Implementation using ``Data.IntMap''}
206 %*                                                                      *
207 %************************************************************************
208
209 \begin{code}
210 newtype UniqFM ele = UFM { unUFM :: M.IntMap ele }
211   deriving (Typeable,Data, Traversable.Traversable, Functor)
212
213 instance Eq ele => Eq (UniqFM ele) where
214     (==) = (==) `on` unUFM
215
216 {-
217 instance Functor UniqFM where
218    fmap f = fmap f . unUFM
219
220 instance Traversable.Traversable UniqFM where
221     traverse f = Traversable.traverse f . unUFM
222 -}
223
224 instance Foldable.Foldable UniqFM where
225     foldMap f = Foldable.foldMap f . unUFM
226
227 emptyUFM = UFM M.empty
228 isNullUFM (UFM m) = M.null m
229 unitUFM k v = UFM (M.singleton (getKey $ getUnique k) v)
230 unitDirectlyUFM u v = UFM (M.singleton (getKey u) v)
231 listToUFM = foldl (\m (k, v) -> addToUFM m k v) emptyUFM
232 listToUFM_Directly = foldl (\m (u, v) -> addToUFM_Directly m u v) emptyUFM
233 listToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v) emptyUFM
234
235 alterUFM f (UFM m) k = UFM (M.alter f (getKey $ getUnique k) m)
236 addToUFM (UFM m) k v   = UFM (M.insert (getKey $ getUnique k) v m)
237 addListToUFM = foldl (\m (k, v) -> addToUFM m k v)
238 addListToUFM_Directly = foldl (\m (k, v) -> addToUFM_Directly m k v)
239 addToUFM_Directly (UFM m) u v = UFM (M.insert (getKey u) v m)
240
241 -- Arguments of combining function of M.insertWith and addToUFM_C are flipped.
242 addToUFM_C f (UFM m) k v =
243   UFM (M.insertWith (flip f) (getKey $ getUnique k) v m)
244 addToUFM_Acc exi new (UFM m) k v =
245   UFM (M.insertWith (\_new old -> exi v old) (getKey $ getUnique k) (new v) m)
246 addListToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v)
247
248 adjustUFM f (UFM m) k = UFM (M.adjust f (getKey $ getUnique k) m)
249 adjustUFM_Directly f (UFM m) u = UFM (M.adjust f (getKey u) m)
250
251 delFromUFM (UFM m) k = UFM (M.delete (getKey $ getUnique k) m)
252 delListFromUFM = foldl delFromUFM
253 delFromUFM_Directly (UFM m) u = UFM (M.delete (getKey u) m)
254
255 -- M.union is left-biased, plusUFM should be right-biased.
256 plusUFM (UFM x) (UFM y) = UFM (M.union y x)
257      -- Note (M.union y x), with arguments flipped
258      -- M.union is left-biased, plusUFM should be right-biased.
259
260 plusUFM_C f (UFM x) (UFM y) = UFM (M.unionWith f x y)
261
262 plusUFM_CD f (UFM xm) dx (UFM ym) dy
263 {-
264 The following implementation should be used as soon as we can expect
265 containers-0.5; presumably from GHC 7.9 on:
266     = UFM $ M.mergeWithKey
267         (\_ x y -> Just (x `f` y))
268         (M.map (\x -> x `f` dy))
269         (M.map (\y -> dx `f` y))
270         xm ym
271 -}
272     = UFM $ M.intersectionWith f xm ym
273         `M.union` M.map (\x -> x  `f` dy) xm
274         `M.union` M.map (\y -> dx `f`  y) ym
275 minusUFM (UFM x) (UFM y) = UFM (M.difference x y)
276 intersectUFM (UFM x) (UFM y) = UFM (M.intersection x y)
277 intersectUFM_C f (UFM x) (UFM y) = UFM (M.intersectionWith f x y)
278
279 foldUFM k z (UFM m) = M.fold k z m
280 foldUFM_Directly k z (UFM m) = M.foldWithKey (k . getUnique) z m
281 mapUFM f (UFM m) = UFM (M.map f m)
282 mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . getUnique) m)
283 filterUFM p (UFM m) = UFM (M.filter p m)
284 filterUFM_Directly p (UFM m) = UFM (M.filterWithKey (p . getUnique) m)
285 partitionUFM p (UFM m) = case M.partition p m of
286                            (left, right) -> (UFM left, UFM right)
287
288 sizeUFM (UFM m) = M.size m
289 elemUFM k (UFM m) = M.member (getKey $ getUnique k) m
290 elemUFM_Directly u (UFM m) = M.member (getKey u) m
291
292 splitUFM (UFM m) k = case M.splitLookup (getKey $ getUnique k) m of
293                        (less, equal, greater) -> (UFM less, equal, UFM greater)
294 lookupUFM (UFM m) k = M.lookup (getKey $ getUnique k) m
295 lookupUFM_Directly (UFM m) u = M.lookup (getKey u) m
296 lookupWithDefaultUFM (UFM m) v k = M.findWithDefault v (getKey $ getUnique k) m
297 lookupWithDefaultUFM_Directly (UFM m) v u = M.findWithDefault v (getKey u) m
298 keysUFM (UFM m) = map getUnique $ M.keys m
299 eltsUFM (UFM m) = M.elems m
300 ufmToSet_Directly (UFM m) = M.keysSet m
301 ufmToList (UFM m) = map (\(k, v) -> (getUnique k, v)) $ M.toList m
302
303 -- Hoopl
304 joinUFM :: JoinFun v -> JoinFun (UniqFM v)
305 joinUFM eltJoin l (OldFact old) (NewFact new) = foldUFM_Directly add (NoChange, old) new
306     where add k new_v (ch, joinmap) =
307             case lookupUFM_Directly joinmap k of
308                 Nothing -> (SomeChange, addToUFM_Directly joinmap k new_v)
309                 Just old_v -> case eltJoin l (OldFact old_v) (NewFact new_v) of
310                                 (SomeChange, v') -> (SomeChange, addToUFM_Directly joinmap k v')
311                                 (NoChange, _) -> (ch, joinmap)
312
313 \end{code}
314
315 %************************************************************************
316 %*                                                                      *
317 \subsection{Output-ery}
318 %*                                                                      *
319 %************************************************************************
320
321 \begin{code}
322 instance Outputable a => Outputable (UniqFM a) where
323     ppr ufm = pprUniqFM ppr ufm
324
325 pprUniqFM :: (a -> SDoc) -> UniqFM a -> SDoc
326 pprUniqFM ppr_elt ufm
327   = brackets $ fsep $ punctuate comma $
328     [ ppr uq <+> ptext (sLit ":->") <+> ppr_elt elt
329     | (uq, elt) <- ufmToList ufm ]
330 \end{code}