Revert "Batch merge"
[ghc.git] / compiler / utils / TrieMap.hs
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 -}
5
6 {-# LANGUAGE RankNTypes #-}
7 {-# LANGUAGE TypeFamilies #-}
8 {-# LANGUAGE FlexibleContexts #-}
9 {-# LANGUAGE TypeSynonymInstances #-}
10 {-# LANGUAGE FlexibleInstances #-}
11 {-# LANGUAGE UndecidableInstances #-}
12 module TrieMap(
13 -- * Maps over 'Maybe' values
14 MaybeMap,
15 -- * Maps over 'List' values
16 ListMap,
17 -- * Maps over 'Literal's
18 LiteralMap,
19 -- * 'TrieMap' class
20 TrieMap(..), insertTM, deleteTM,
21
22 -- * Things helpful for adding additional Instances.
23 (>.>), (|>), (|>>), XT,
24 foldMaybe,
25 -- * Map for leaf compression
26 GenMap,
27 lkG, xtG, mapG, fdG,
28 xtList, lkList
29
30 ) where
31
32 import GhcPrelude
33
34 import Literal
35 import UniqDFM
36 import Unique( Unique )
37
38 import qualified Data.Map as Map
39 import qualified Data.IntMap as IntMap
40 import Outputable
41 import Control.Monad( (>=>) )
42
43 {-
44 This module implements TrieMaps, which are finite mappings
45 whose key is a structured value like a CoreExpr or Type.
46
47 This file implements tries over general data structures.
48 Implementation for tries over Core Expressions/Types are
49 available in coreSyn/TrieMap.
50
51 The regular pattern for handling TrieMaps on data structures was first
52 described (to my knowledge) in Connelly and Morris's 1995 paper "A
53 generalization of the Trie Data Structure"; there is also an accessible
54 description of the idea in Okasaki's book "Purely Functional Data
55 Structures", Section 10.3.2
56
57 ************************************************************************
58 * *
59 The TrieMap class
60 * *
61 ************************************************************************
62 -}
63
64 type XT a = Maybe a -> Maybe a -- How to alter a non-existent elt (Nothing)
65 -- or an existing elt (Just)
66
67 class TrieMap m where
68 type Key m :: *
69 emptyTM :: m a
70 lookupTM :: forall b. Key m -> m b -> Maybe b
71 alterTM :: forall b. Key m -> XT b -> m b -> m b
72 mapTM :: (a->b) -> m a -> m b
73
74 foldTM :: (a -> b -> b) -> m a -> b -> b
75 -- The unusual argument order here makes
76 -- it easy to compose calls to foldTM;
77 -- see for example fdE below
78
79 insertTM :: TrieMap m => Key m -> a -> m a -> m a
80 insertTM k v m = alterTM k (\_ -> Just v) m
81
82 deleteTM :: TrieMap m => Key m -> m a -> m a
83 deleteTM k m = alterTM k (\_ -> Nothing) m
84
85 ----------------------
86 -- Recall that
87 -- Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c
88
89 (>.>) :: (a -> b) -> (b -> c) -> a -> c
90 -- Reverse function composition (do f first, then g)
91 infixr 1 >.>
92 (f >.> g) x = g (f x)
93 infixr 1 |>, |>>
94
95 (|>) :: a -> (a->b) -> b -- Reverse application
96 x |> f = f x
97
98 ----------------------
99 (|>>) :: TrieMap m2
100 => (XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
101 -> (m2 a -> m2 a)
102 -> m1 (m2 a) -> m1 (m2 a)
103 (|>>) f g = f (Just . g . deMaybe)
104
105 deMaybe :: TrieMap m => Maybe (m a) -> m a
106 deMaybe Nothing = emptyTM
107 deMaybe (Just m) = m
108
109 {-
110 ************************************************************************
111 * *
112 IntMaps
113 * *
114 ************************************************************************
115 -}
116
117 instance TrieMap IntMap.IntMap where
118 type Key IntMap.IntMap = Int
119 emptyTM = IntMap.empty
120 lookupTM k m = IntMap.lookup k m
121 alterTM = xtInt
122 foldTM k m z = IntMap.foldr k z m
123 mapTM f m = IntMap.map f m
124
125 xtInt :: Int -> XT a -> IntMap.IntMap a -> IntMap.IntMap a
126 xtInt k f m = IntMap.alter f k m
127
128 instance Ord k => TrieMap (Map.Map k) where
129 type Key (Map.Map k) = k
130 emptyTM = Map.empty
131 lookupTM = Map.lookup
132 alterTM k f m = Map.alter f k m
133 foldTM k m z = Map.foldr k z m
134 mapTM f m = Map.map f m
135
136
137 {-
138 Note [foldTM determinism]
139 ~~~~~~~~~~~~~~~~~~~~~~~~~
140 We want foldTM to be deterministic, which is why we have an instance of
141 TrieMap for UniqDFM, but not for UniqFM. Here's an example of some things that
142 go wrong if foldTM is nondeterministic. Consider:
143
144 f a b = return (a <> b)
145
146 Depending on the order that the typechecker generates constraints you
147 get either:
148
149 f :: (Monad m, Monoid a) => a -> a -> m a
150
151 or:
152
153 f :: (Monoid a, Monad m) => a -> a -> m a
154
155 The generated code will be different after desugaring as the dictionaries
156 will be bound in different orders, leading to potential ABI incompatibility.
157
158 One way to solve this would be to notice that the typeclasses could be
159 sorted alphabetically.
160
161 Unfortunately that doesn't quite work with this example:
162
163 f a b = let x = a <> a; y = b <> b in x
164
165 where you infer:
166
167 f :: (Monoid m, Monoid m1) => m1 -> m -> m1
168
169 or:
170
171 f :: (Monoid m1, Monoid m) => m1 -> m -> m1
172
173 Here you could decide to take the order of the type variables in the type
174 according to depth first traversal and use it to order the constraints.
175
176 The real trouble starts when the user enables incoherent instances and
177 the compiler has to make an arbitrary choice. Consider:
178
179 class T a b where
180 go :: a -> b -> String
181
182 instance (Show b) => T Int b where
183 go a b = show a ++ show b
184
185 instance (Show a) => T a Bool where
186 go a b = show a ++ show b
187
188 f = go 10 True
189
190 GHC is free to choose either dictionary to implement f, but for the sake of
191 determinism we'd like it to be consistent when compiling the same sources
192 with the same flags.
193
194 inert_dicts :: DictMap is implemented with a TrieMap. In getUnsolvedInerts it
195 gets converted to a bag of (Wanted) Cts using a fold. Then in
196 solve_simple_wanteds it's merged with other WantedConstraints. We want the
197 conversion to a bag to be deterministic. For that purpose we use UniqDFM
198 instead of UniqFM to implement the TrieMap.
199
200 See Note [Deterministic UniqFM] in UniqDFM for more details on how it's made
201 deterministic.
202 -}
203
204 instance TrieMap UniqDFM where
205 type Key UniqDFM = Unique
206 emptyTM = emptyUDFM
207 lookupTM k m = lookupUDFM m k
208 alterTM k f m = alterUDFM f m k
209 foldTM k m z = foldUDFM k z m
210 mapTM f m = mapUDFM f m
211
212 {-
213 ************************************************************************
214 * *
215 Maybes
216 * *
217 ************************************************************************
218
219 If m is a map from k -> val
220 then (MaybeMap m) is a map from (Maybe k) -> val
221 -}
222
223 data MaybeMap m a = MM { mm_nothing :: Maybe a, mm_just :: m a }
224
225 instance TrieMap m => TrieMap (MaybeMap m) where
226 type Key (MaybeMap m) = Maybe (Key m)
227 emptyTM = MM { mm_nothing = Nothing, mm_just = emptyTM }
228 lookupTM = lkMaybe lookupTM
229 alterTM = xtMaybe alterTM
230 foldTM = fdMaybe
231 mapTM = mapMb
232
233 mapMb :: TrieMap m => (a->b) -> MaybeMap m a -> MaybeMap m b
234 mapMb f (MM { mm_nothing = mn, mm_just = mj })
235 = MM { mm_nothing = fmap f mn, mm_just = mapTM f mj }
236
237 lkMaybe :: (forall b. k -> m b -> Maybe b)
238 -> Maybe k -> MaybeMap m a -> Maybe a
239 lkMaybe _ Nothing = mm_nothing
240 lkMaybe lk (Just x) = mm_just >.> lk x
241
242 xtMaybe :: (forall b. k -> XT b -> m b -> m b)
243 -> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a
244 xtMaybe _ Nothing f m = m { mm_nothing = f (mm_nothing m) }
245 xtMaybe tr (Just x) f m = m { mm_just = mm_just m |> tr x f }
246
247 fdMaybe :: TrieMap m => (a -> b -> b) -> MaybeMap m a -> b -> b
248 fdMaybe k m = foldMaybe k (mm_nothing m)
249 . foldTM k (mm_just m)
250
251 {-
252 ************************************************************************
253 * *
254 Lists
255 * *
256 ************************************************************************
257 -}
258
259 data ListMap m a
260 = LM { lm_nil :: Maybe a
261 , lm_cons :: m (ListMap m a) }
262
263 instance TrieMap m => TrieMap (ListMap m) where
264 type Key (ListMap m) = [Key m]
265 emptyTM = LM { lm_nil = Nothing, lm_cons = emptyTM }
266 lookupTM = lkList lookupTM
267 alterTM = xtList alterTM
268 foldTM = fdList
269 mapTM = mapList
270
271 instance (TrieMap m, Outputable a) => Outputable (ListMap m a) where
272 ppr m = text "List elts" <+> ppr (foldTM (:) m [])
273
274 mapList :: TrieMap m => (a->b) -> ListMap m a -> ListMap m b
275 mapList f (LM { lm_nil = mnil, lm_cons = mcons })
276 = LM { lm_nil = fmap f mnil, lm_cons = mapTM (mapTM f) mcons }
277
278 lkList :: TrieMap m => (forall b. k -> m b -> Maybe b)
279 -> [k] -> ListMap m a -> Maybe a
280 lkList _ [] = lm_nil
281 lkList lk (x:xs) = lm_cons >.> lk x >=> lkList lk xs
282
283 xtList :: TrieMap m => (forall b. k -> XT b -> m b -> m b)
284 -> [k] -> XT a -> ListMap m a -> ListMap m a
285 xtList _ [] f m = m { lm_nil = f (lm_nil m) }
286 xtList tr (x:xs) f m = m { lm_cons = lm_cons m |> tr x |>> xtList tr xs f }
287
288 fdList :: forall m a b. TrieMap m
289 => (a -> b -> b) -> ListMap m a -> b -> b
290 fdList k m = foldMaybe k (lm_nil m)
291 . foldTM (fdList k) (lm_cons m)
292
293 foldMaybe :: (a -> b -> b) -> Maybe a -> b -> b
294 foldMaybe _ Nothing b = b
295 foldMaybe k (Just a) b = k a b
296
297 {-
298 ************************************************************************
299 * *
300 Basic maps
301 * *
302 ************************************************************************
303 -}
304
305 type LiteralMap a = Map.Map Literal a
306
307 {-
308 ************************************************************************
309 * *
310 GenMap
311 * *
312 ************************************************************************
313
314 Note [Compressed TrieMap]
315 ~~~~~~~~~~~~~~~~~~~~~~~~~
316
317 The GenMap constructor augments TrieMaps with leaf compression. This helps
318 solve the performance problem detailed in #9960: suppose we have a handful
319 H of entries in a TrieMap, each with a very large key, size K. If you fold over
320 such a TrieMap you'd expect time O(H). That would certainly be true of an
321 association list! But with TrieMap we actually have to navigate down a long
322 singleton structure to get to the elements, so it takes time O(K*H). This
323 can really hurt on many type-level computation benchmarks:
324 see for example T9872d.
325
326 The point of a TrieMap is that you need to navigate to the point where only one
327 key remains, and then things should be fast. So the point of a SingletonMap
328 is that, once we are down to a single (key,value) pair, we stop and
329 just use SingletonMap.
330
331 'EmptyMap' provides an even more basic (but essential) optimization: if there is
332 nothing in the map, don't bother building out the (possibly infinite) recursive
333 TrieMap structure!
334
335 Compressed triemaps are heavily used by CoreMap. So we have to mark some things
336 as INLINEABLE to permit specialization.
337 -}
338
339 data GenMap m a
340 = EmptyMap
341 | SingletonMap (Key m) a
342 | MultiMap (m a)
343
344 instance (Outputable a, Outputable (m a)) => Outputable (GenMap m a) where
345 ppr EmptyMap = text "Empty map"
346 ppr (SingletonMap _ v) = text "Singleton map" <+> ppr v
347 ppr (MultiMap m) = ppr m
348
349 -- TODO undecidable instance
350 instance (Eq (Key m), TrieMap m) => TrieMap (GenMap m) where
351 type Key (GenMap m) = Key m
352 emptyTM = EmptyMap
353 lookupTM = lkG
354 alterTM = xtG
355 foldTM = fdG
356 mapTM = mapG
357
358 --We want to be able to specialize these functions when defining eg
359 --tries over (GenMap CoreExpr) which requires INLINEABLE
360
361 {-# INLINEABLE lkG #-}
362 lkG :: (Eq (Key m), TrieMap m) => Key m -> GenMap m a -> Maybe a
363 lkG _ EmptyMap = Nothing
364 lkG k (SingletonMap k' v') | k == k' = Just v'
365 | otherwise = Nothing
366 lkG k (MultiMap m) = lookupTM k m
367
368 {-# INLINEABLE xtG #-}
369 xtG :: (Eq (Key m), TrieMap m) => Key m -> XT a -> GenMap m a -> GenMap m a
370 xtG k f EmptyMap
371 = case f Nothing of
372 Just v -> SingletonMap k v
373 Nothing -> EmptyMap
374 xtG k f m@(SingletonMap k' v')
375 | k' == k
376 -- The new key matches the (single) key already in the tree. Hence,
377 -- apply @f@ to @Just v'@ and build a singleton or empty map depending
378 -- on the 'Just'/'Nothing' response respectively.
379 = case f (Just v') of
380 Just v'' -> SingletonMap k' v''
381 Nothing -> EmptyMap
382 | otherwise
383 -- We've hit a singleton tree for a different key than the one we are
384 -- searching for. Hence apply @f@ to @Nothing@. If result is @Nothing@ then
385 -- we can just return the old map. If not, we need a map with *two*
386 -- entries. The easiest way to do that is to insert two items into an empty
387 -- map of type @m a@.
388 = case f Nothing of
389 Nothing -> m
390 Just v -> emptyTM |> alterTM k' (const (Just v'))
391 >.> alterTM k (const (Just v))
392 >.> MultiMap
393 xtG k f (MultiMap m) = MultiMap (alterTM k f m)
394
395 {-# INLINEABLE mapG #-}
396 mapG :: TrieMap m => (a -> b) -> GenMap m a -> GenMap m b
397 mapG _ EmptyMap = EmptyMap
398 mapG f (SingletonMap k v) = SingletonMap k (f v)
399 mapG f (MultiMap m) = MultiMap (mapTM f m)
400
401 {-# INLINEABLE fdG #-}
402 fdG :: TrieMap m => (a -> b -> b) -> GenMap m a -> b -> b
403 fdG _ EmptyMap = \z -> z
404 fdG k (SingletonMap _ v) = \z -> k v z
405 fdG k (MultiMap m) = foldTM k m