2b23b256cb9742bbd83f92b20e0304a2a9698a23
[ghc.git] / compiler / coreSyn / 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 CoreMap, emptyCoreMap, extendCoreMap, lookupCoreMap, foldCoreMap,
14 TypeMap, emptyTypeMap, extendTypeMap, lookupTypeMap, foldTypeMap,
15 LooseTypeMap,
16 MaybeMap,
17 ListMap,
18 TrieMap(..), insertTM, deleteTM
19 ) where
20
21 import CoreSyn
22 import Coercion
23 import Literal
24 import Name
25 import Type
26 import TyCoRep
27 import Var
28 import UniqFM
29 import Unique( Unique )
30 import FastString(FastString)
31
32 import qualified Data.Map as Map
33 import qualified Data.IntMap as IntMap
34 import VarEnv
35 import NameEnv
36 import Outputable
37 import Control.Monad( (>=>) )
38
39 {-
40 This module implements TrieMaps, which are finite mappings
41 whose key is a structured value like a CoreExpr or Type.
42
43 The code is very regular and boilerplate-like, but there is
44 some neat handling of *binders*. In effect they are deBruijn
45 numbered on the fly.
46
47 The regular pattern for handling TrieMaps on data structures was first
48 described (to my knowledge) in Connelly and Morris's 1995 paper "A
49 generalization of the Trie Data Structure"; there is also an accessible
50 description of the idea in Okasaki's book "Purely Functional Data
51 Structures", Section 10.3.2
52
53 ************************************************************************
54 * *
55 The TrieMap class
56 * *
57 ************************************************************************
58 -}
59
60 type XT a = Maybe a -> Maybe a -- How to alter a non-existent elt (Nothing)
61 -- or an existing elt (Just)
62
63 class TrieMap m where
64 type Key m :: *
65 emptyTM :: m a
66 lookupTM :: forall b. Key m -> m b -> Maybe b
67 alterTM :: forall b. Key m -> XT b -> m b -> m b
68 mapTM :: (a->b) -> m a -> m b
69
70 foldTM :: (a -> b -> b) -> m a -> b -> b
71 -- The unusual argument order here makes
72 -- it easy to compose calls to foldTM;
73 -- see for example fdE below
74
75 insertTM :: TrieMap m => Key m -> a -> m a -> m a
76 insertTM k v m = alterTM k (\_ -> Just v) m
77
78 deleteTM :: TrieMap m => Key m -> m a -> m a
79 deleteTM k m = alterTM k (\_ -> Nothing) m
80
81 ----------------------
82 -- Recall that
83 -- Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c
84
85 (>.>) :: (a -> b) -> (b -> c) -> a -> c
86 -- Reverse function composition (do f first, then g)
87 infixr 1 >.>
88 (f >.> g) x = g (f x)
89 infixr 1 |>, |>>
90
91 (|>) :: a -> (a->b) -> b -- Reverse application
92 x |> f = f x
93
94 ----------------------
95 (|>>) :: TrieMap m2
96 => (XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
97 -> (m2 a -> m2 a)
98 -> m1 (m2 a) -> m1 (m2 a)
99 (|>>) f g = f (Just . g . deMaybe)
100
101 deMaybe :: TrieMap m => Maybe (m a) -> m a
102 deMaybe Nothing = emptyTM
103 deMaybe (Just m) = m
104
105 {-
106 ************************************************************************
107 * *
108 IntMaps
109 * *
110 ************************************************************************
111 -}
112
113 instance TrieMap IntMap.IntMap where
114 type Key IntMap.IntMap = Int
115 emptyTM = IntMap.empty
116 lookupTM k m = IntMap.lookup k m
117 alterTM = xtInt
118 foldTM k m z = IntMap.fold k z m
119 mapTM f m = IntMap.map f m
120
121 xtInt :: Int -> XT a -> IntMap.IntMap a -> IntMap.IntMap a
122 xtInt k f m = IntMap.alter f k m
123
124 instance Ord k => TrieMap (Map.Map k) where
125 type Key (Map.Map k) = k
126 emptyTM = Map.empty
127 lookupTM = Map.lookup
128 alterTM k f m = Map.alter f k m
129 foldTM k m z = Map.fold k z m
130 mapTM f m = Map.map f m
131
132 instance TrieMap UniqFM where
133 type Key UniqFM = Unique
134 emptyTM = emptyUFM
135 lookupTM k m = lookupUFM m k
136 alterTM k f m = alterUFM f m k
137 foldTM k m z = foldUFM k z m
138 mapTM f m = mapUFM f m
139
140 {-
141 ************************************************************************
142 * *
143 Maybes
144 * *
145 ************************************************************************
146
147 If m is a map from k -> val
148 then (MaybeMap m) is a map from (Maybe k) -> val
149 -}
150
151 data MaybeMap m a = MM { mm_nothing :: Maybe a, mm_just :: m a }
152
153 instance TrieMap m => TrieMap (MaybeMap m) where
154 type Key (MaybeMap m) = Maybe (Key m)
155 emptyTM = MM { mm_nothing = Nothing, mm_just = emptyTM }
156 lookupTM = lkMaybe lookupTM
157 alterTM = xtMaybe alterTM
158 foldTM = fdMaybe
159 mapTM = mapMb
160
161 mapMb :: TrieMap m => (a->b) -> MaybeMap m a -> MaybeMap m b
162 mapMb f (MM { mm_nothing = mn, mm_just = mj })
163 = MM { mm_nothing = fmap f mn, mm_just = mapTM f mj }
164
165 lkMaybe :: (forall b. k -> m b -> Maybe b)
166 -> Maybe k -> MaybeMap m a -> Maybe a
167 lkMaybe _ Nothing = mm_nothing
168 lkMaybe lk (Just x) = mm_just >.> lk x
169
170 xtMaybe :: (forall b. k -> XT b -> m b -> m b)
171 -> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a
172 xtMaybe _ Nothing f m = m { mm_nothing = f (mm_nothing m) }
173 xtMaybe tr (Just x) f m = m { mm_just = mm_just m |> tr x f }
174
175 fdMaybe :: TrieMap m => (a -> b -> b) -> MaybeMap m a -> b -> b
176 fdMaybe k m = foldMaybe k (mm_nothing m)
177 . foldTM k (mm_just m)
178
179 {-
180 ************************************************************************
181 * *
182 Lists
183 * *
184 ************************************************************************
185 -}
186
187 data ListMap m a
188 = LM { lm_nil :: Maybe a
189 , lm_cons :: m (ListMap m a) }
190
191 instance TrieMap m => TrieMap (ListMap m) where
192 type Key (ListMap m) = [Key m]
193 emptyTM = LM { lm_nil = Nothing, lm_cons = emptyTM }
194 lookupTM = lkList lookupTM
195 alterTM = xtList alterTM
196 foldTM = fdList
197 mapTM = mapList
198
199 mapList :: TrieMap m => (a->b) -> ListMap m a -> ListMap m b
200 mapList f (LM { lm_nil = mnil, lm_cons = mcons })
201 = LM { lm_nil = fmap f mnil, lm_cons = mapTM (mapTM f) mcons }
202
203 lkList :: TrieMap m => (forall b. k -> m b -> Maybe b)
204 -> [k] -> ListMap m a -> Maybe a
205 lkList _ [] = lm_nil
206 lkList lk (x:xs) = lm_cons >.> lk x >=> lkList lk xs
207
208 xtList :: TrieMap m => (forall b. k -> XT b -> m b -> m b)
209 -> [k] -> XT a -> ListMap m a -> ListMap m a
210 xtList _ [] f m = m { lm_nil = f (lm_nil m) }
211 xtList tr (x:xs) f m = m { lm_cons = lm_cons m |> tr x |>> xtList tr xs f }
212
213 fdList :: forall m a b. TrieMap m
214 => (a -> b -> b) -> ListMap m a -> b -> b
215 fdList k m = foldMaybe k (lm_nil m)
216 . foldTM (fdList k) (lm_cons m)
217
218 foldMaybe :: (a -> b -> b) -> Maybe a -> b -> b
219 foldMaybe _ Nothing b = b
220 foldMaybe k (Just a) b = k a b
221
222 {-
223 ************************************************************************
224 * *
225 Basic maps
226 * *
227 ************************************************************************
228 -}
229
230 lkNamed :: NamedThing n => n -> NameEnv a -> Maybe a
231 lkNamed n env = lookupNameEnv env (getName n)
232
233 xtNamed :: NamedThing n => n -> XT a -> NameEnv a -> NameEnv a
234 xtNamed tc f m = alterNameEnv f m (getName tc)
235
236 ------------------------
237 type LiteralMap a = Map.Map Literal a
238
239 emptyLiteralMap :: LiteralMap a
240 emptyLiteralMap = emptyTM
241
242 lkLit :: Literal -> LiteralMap a -> Maybe a
243 lkLit = lookupTM
244
245 xtLit :: Literal -> XT a -> LiteralMap a -> LiteralMap a
246 xtLit = alterTM
247
248 {-
249 ************************************************************************
250 * *
251 GenMap
252 * *
253 ************************************************************************
254
255 Note [Compressed TrieMap]
256 ~~~~~~~~~~~~~~~~~~~~~~~~~
257
258 The GenMap constructor augments TrieMaps with leaf compression. This helps
259 solve the performance problem detailed in #9960: suppose we have a handful
260 H of entries in a TrieMap, each with a very large key, size K. If you fold over
261 such a TrieMap you'd expect time O(H). That would certainly be true of an
262 association list! But with TrieMap we actually have to navigate down a long
263 singleton structure to get to the elements, so it takes time O(K*H). This
264 can really hurt on many type-level computation benchmarks:
265 see for example T9872d.
266
267 The point of a TrieMap is that you need to navigate to the point where only one
268 key remains, and then things should be fast. So the point of a SingletonMap
269 is that, once we are down to a single (key,value) pair, we stop and
270 just use SingletonMap.
271
272 'EmptyMap' provides an even more basic (but essential) optimization: if there is
273 nothing in the map, don't bother building out the (possibly infinite) recursive
274 TrieMap structure!
275 -}
276
277 data GenMap m a
278 = EmptyMap
279 | SingletonMap (Key m) a
280 | MultiMap (m a)
281
282 instance (Outputable a, Outputable (m a)) => Outputable (GenMap m a) where
283 ppr EmptyMap = text "Empty map"
284 ppr (SingletonMap _ v) = text "Singleton map" <+> ppr v
285 ppr (MultiMap m) = ppr m
286
287 -- TODO undecidable instance
288 instance (Eq (Key m), TrieMap m) => TrieMap (GenMap m) where
289 type Key (GenMap m) = Key m
290 emptyTM = EmptyMap
291 lookupTM = lkG
292 alterTM = xtG
293 foldTM = fdG
294 mapTM = mapG
295
296 -- NB: Be careful about RULES and type families (#5821). So we should make sure
297 -- to specify @Key TypeMapX@ (and not @DeBruijn Type@, the reduced form)
298
299 {-# SPECIALIZE lkG :: Key TypeMapX -> TypeMapG a -> Maybe a #-}
300 {-# SPECIALIZE lkG :: Key CoercionMapX -> CoercionMapG a -> Maybe a #-}
301 {-# SPECIALIZE lkG :: Key CoreMapX -> CoreMapG a -> Maybe a #-}
302 lkG :: (Eq (Key m), TrieMap m) => Key m -> GenMap m a -> Maybe a
303 lkG _ EmptyMap = Nothing
304 lkG k (SingletonMap k' v') | k == k' = Just v'
305 | otherwise = Nothing
306 lkG k (MultiMap m) = lookupTM k m
307
308 {-# SPECIALIZE xtG :: Key TypeMapX -> XT a -> TypeMapG a -> TypeMapG a #-}
309 {-# SPECIALIZE xtG :: Key CoercionMapX -> XT a -> CoercionMapG a -> CoercionMapG a #-}
310 {-# SPECIALIZE xtG :: Key CoreMapX -> XT a -> CoreMapG a -> CoreMapG a #-}
311 xtG :: (Eq (Key m), TrieMap m) => Key m -> XT a -> GenMap m a -> GenMap m a
312 xtG k f EmptyMap
313 = case f Nothing of
314 Just v -> SingletonMap k v
315 Nothing -> EmptyMap
316 xtG k f m@(SingletonMap k' v')
317 | k' == k
318 -- The new key matches the (single) key already in the tree. Hence,
319 -- apply @f@ to @Just v'@ and build a singleton or empty map depending
320 -- on the 'Just'/'Nothing' response respectively.
321 = case f (Just v') of
322 Just v'' -> SingletonMap k' v''
323 Nothing -> EmptyMap
324 | otherwise
325 -- We've hit a singleton tree for a different key than the one we are
326 -- searching for. Hence apply @f@ to @Nothing@. If result is @Nothing@ then
327 -- we can just return the old map. If not, we need a map with *two*
328 -- entries. The easiest way to do that is to insert two items into an empty
329 -- map of type @m a@.
330 = case f Nothing of
331 Nothing -> m
332 Just v -> emptyTM |> alterTM k' (const (Just v'))
333 >.> alterTM k (const (Just v))
334 >.> MultiMap
335 xtG k f (MultiMap m) = MultiMap (alterTM k f m)
336
337 {-# SPECIALIZE mapG :: (a -> b) -> TypeMapG a -> TypeMapG b #-}
338 {-# SPECIALIZE mapG :: (a -> b) -> CoercionMapG a -> CoercionMapG b #-}
339 {-# SPECIALIZE mapG :: (a -> b) -> CoreMapG a -> CoreMapG b #-}
340 mapG :: TrieMap m => (a -> b) -> GenMap m a -> GenMap m b
341 mapG _ EmptyMap = EmptyMap
342 mapG f (SingletonMap k v) = SingletonMap k (f v)
343 mapG f (MultiMap m) = MultiMap (mapTM f m)
344
345 {-# SPECIALIZE fdG :: (a -> b -> b) -> TypeMapG a -> b -> b #-}
346 {-# SPECIALIZE fdG :: (a -> b -> b) -> CoercionMapG a -> b -> b #-}
347 {-# SPECIALIZE fdG :: (a -> b -> b) -> CoreMapG a -> b -> b #-}
348 fdG :: TrieMap m => (a -> b -> b) -> GenMap m a -> b -> b
349 fdG _ EmptyMap = \z -> z
350 fdG k (SingletonMap _ v) = \z -> k v z
351 fdG k (MultiMap m) = foldTM k m
352
353 {-
354 ************************************************************************
355 * *
356 CoreMap
357 * *
358 ************************************************************************
359
360 Note [Binders]
361 ~~~~~~~~~~~~~~
362 * In general we check binders as late as possible because types are
363 less likely to differ than expression structure. That's why
364 cm_lam :: CoreMapG (TypeMapG a)
365 rather than
366 cm_lam :: TypeMapG (CoreMapG a)
367
368 * We don't need to look at the type of some binders, notalby
369 - the case binder in (Case _ b _ _)
370 - the binders in an alternative
371 because they are totally fixed by the context
372
373 Note [Empty case alternatives]
374 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
375 * For a key (Case e b ty (alt:alts)) we don't need to look the return type
376 'ty', because every alternative has that type.
377
378 * For a key (Case e b ty []) we MUST look at the return type 'ty', because
379 otherwise (Case (error () "urk") _ Int []) would compare equal to
380 (Case (error () "urk") _ Bool [])
381 which is utterly wrong (Trac #6097)
382
383 We could compare the return type regardless, but the wildly common case
384 is that it's unnecessary, so we have two fields (cm_case and cm_ecase)
385 for the two possibilities. Only cm_ecase looks at the type.
386
387 See also Note [Empty case alternatives] in CoreSyn.
388 -}
389
390 -- | @CoreMap a@ is a map from 'CoreExpr' to @a@. If you are a client, this
391 -- is the type you want.
392 newtype CoreMap a = CoreMap (CoreMapG a)
393
394 instance TrieMap CoreMap where
395 type Key CoreMap = CoreExpr
396 emptyTM = CoreMap emptyTM
397 lookupTM k (CoreMap m) = lookupTM (deBruijnize k) m
398 alterTM k f (CoreMap m) = CoreMap (alterTM (deBruijnize k) f m)
399 foldTM k (CoreMap m) = foldTM k m
400 mapTM f (CoreMap m) = CoreMap (mapTM f m)
401
402 -- | @CoreMapG a@ is a map from @DeBruijn CoreExpr@ to @a@. The extended
403 -- key makes it suitable for recursive traversal, since it can track binders,
404 -- but it is strictly internal to this module. If you are including a 'CoreMap'
405 -- inside another 'TrieMap', this is the type you want.
406 type CoreMapG = GenMap CoreMapX
407
408 -- | @CoreMapX a@ is the base map from @DeBruijn CoreExpr@ to @a@, but without
409 -- the 'GenMap' optimization.
410 data CoreMapX a
411 = CM { cm_var :: VarMap a
412 , cm_lit :: LiteralMap a
413 , cm_co :: CoercionMapG a
414 , cm_type :: TypeMapG a
415 , cm_cast :: CoreMapG (CoercionMapG a)
416 , cm_tick :: CoreMapG (TickishMap a)
417 , cm_app :: CoreMapG (CoreMapG a)
418 , cm_lam :: CoreMapG (BndrMap a) -- Note [Binders]
419 , cm_letn :: CoreMapG (CoreMapG (BndrMap a))
420 , cm_letr :: ListMap CoreMapG (CoreMapG (ListMap BndrMap a))
421 , cm_case :: CoreMapG (ListMap AltMap a)
422 , cm_ecase :: CoreMapG (TypeMapG a) -- Note [Empty case alternatives]
423 }
424
425 instance Eq (DeBruijn CoreExpr) where
426 D env1 e1 == D env2 e2 = go e1 e2 where
427 go (Var v1) (Var v2) = case (lookupCME env1 v1, lookupCME env2 v2) of
428 (Just b1, Just b2) -> b1 == b2
429 (Nothing, Nothing) -> v1 == v2
430 _ -> False
431 go (Lit lit1) (Lit lit2) = lit1 == lit2
432 go (Type t1) (Type t2) = D env1 t1 == D env2 t2
433 go (Coercion co1) (Coercion co2) = D env1 co1 == D env2 co2
434 go (Cast e1 co1) (Cast e2 co2) = D env1 co1 == D env2 co2 && go e1 e2
435 go (App f1 a1) (App f2 a2) = go f1 f2 && go a1 a2
436 -- This seems a bit dodgy, see 'eqTickish'
437 go (Tick n1 e1) (Tick n2 e2) = n1 == n2 && go e1 e2
438
439 go (Lam b1 e1) (Lam b2 e2)
440 = D env1 (varType b1) == D env2 (varType b2)
441 && D (extendCME env1 b1) e1 == D (extendCME env2 b2) e2
442
443 go (Let (NonRec v1 r1) e1) (Let (NonRec v2 r2) e2)
444 = go r1 r2
445 && D (extendCME env1 v1) e1 == D (extendCME env2 v2) e2
446
447 go (Let (Rec ps1) e1) (Let (Rec ps2) e2)
448 = length ps1 == length ps2
449 && D env1' rs1 == D env2' rs2
450 && D env1' e1 == D env2' e2
451 where
452 (bs1,rs1) = unzip ps1
453 (bs2,rs2) = unzip ps2
454 env1' = extendCMEs env1 bs1
455 env2' = extendCMEs env2 bs2
456
457 go (Case e1 b1 t1 a1) (Case e2 b2 t2 a2)
458 | null a1 -- See Note [Empty case alternatives]
459 = null a2 && go e1 e2 && D env1 t1 == D env2 t2
460 | otherwise
461 = go e1 e2 && D (extendCME env1 b1) a1 == D (extendCME env2 b2) a2
462
463 go _ _ = False
464
465 emptyE :: CoreMapX a
466 emptyE = CM { cm_var = emptyTM, cm_lit = emptyLiteralMap
467 , cm_co = emptyTM, cm_type = emptyTM
468 , cm_cast = emptyTM, cm_app = emptyTM
469 , cm_lam = emptyTM, cm_letn = emptyTM
470 , cm_letr = emptyTM, cm_case = emptyTM
471 , cm_ecase = emptyTM, cm_tick = emptyTM }
472
473 instance TrieMap CoreMapX where
474 type Key CoreMapX = DeBruijn CoreExpr
475 emptyTM = emptyE
476 lookupTM = lkE
477 alterTM = xtE
478 foldTM = fdE
479 mapTM = mapE
480
481 --------------------------
482 mapE :: (a->b) -> CoreMapX a -> CoreMapX b
483 mapE f (CM { cm_var = cvar, cm_lit = clit
484 , cm_co = cco, cm_type = ctype
485 , cm_cast = ccast , cm_app = capp
486 , cm_lam = clam, cm_letn = cletn
487 , cm_letr = cletr, cm_case = ccase
488 , cm_ecase = cecase, cm_tick = ctick })
489 = CM { cm_var = mapTM f cvar, cm_lit = mapTM f clit
490 , cm_co = mapTM f cco, cm_type = mapTM f ctype
491 , cm_cast = mapTM (mapTM f) ccast, cm_app = mapTM (mapTM f) capp
492 , cm_lam = mapTM (mapTM f) clam, cm_letn = mapTM (mapTM (mapTM f)) cletn
493 , cm_letr = mapTM (mapTM (mapTM f)) cletr, cm_case = mapTM (mapTM f) ccase
494 , cm_ecase = mapTM (mapTM f) cecase, cm_tick = mapTM (mapTM f) ctick }
495
496 --------------------------
497 lookupCoreMap :: CoreMap a -> CoreExpr -> Maybe a
498 lookupCoreMap cm e = lookupTM e cm
499
500 extendCoreMap :: CoreMap a -> CoreExpr -> a -> CoreMap a
501 extendCoreMap m e v = alterTM e (\_ -> Just v) m
502
503 foldCoreMap :: (a -> b -> b) -> b -> CoreMap a -> b
504 foldCoreMap k z m = foldTM k m z
505
506 emptyCoreMap :: CoreMap a
507 emptyCoreMap = emptyTM
508
509 instance Outputable a => Outputable (CoreMap a) where
510 ppr m = text "CoreMap elts" <+> ppr (foldTM (:) m [])
511
512 -------------------------
513 fdE :: (a -> b -> b) -> CoreMapX a -> b -> b
514 fdE k m
515 = foldTM k (cm_var m)
516 . foldTM k (cm_lit m)
517 . foldTM k (cm_co m)
518 . foldTM k (cm_type m)
519 . foldTM (foldTM k) (cm_cast m)
520 . foldTM (foldTM k) (cm_tick m)
521 . foldTM (foldTM k) (cm_app m)
522 . foldTM (foldTM k) (cm_lam m)
523 . foldTM (foldTM (foldTM k)) (cm_letn m)
524 . foldTM (foldTM (foldTM k)) (cm_letr m)
525 . foldTM (foldTM k) (cm_case m)
526 . foldTM (foldTM k) (cm_ecase m)
527
528 -- lkE: lookup in trie for expressions
529 lkE :: DeBruijn CoreExpr -> CoreMapX a -> Maybe a
530 lkE (D env expr) cm = go expr cm
531 where
532 go (Var v) = cm_var >.> lkVar env v
533 go (Lit l) = cm_lit >.> lkLit l
534 go (Type t) = cm_type >.> lkG (D env t)
535 go (Coercion c) = cm_co >.> lkG (D env c)
536 go (Cast e c) = cm_cast >.> lkG (D env e) >=> lkG (D env c)
537 go (Tick tickish e) = cm_tick >.> lkG (D env e) >=> lkTickish tickish
538 go (App e1 e2) = cm_app >.> lkG (D env e2) >=> lkG (D env e1)
539 go (Lam v e) = cm_lam >.> lkG (D (extendCME env v) e)
540 >=> lkBndr env v
541 go (Let (NonRec b r) e) = cm_letn >.> lkG (D env r)
542 >=> lkG (D (extendCME env b) e) >=> lkBndr env b
543 go (Let (Rec prs) e) = let (bndrs,rhss) = unzip prs
544 env1 = extendCMEs env bndrs
545 in cm_letr
546 >.> lkList (lkG . D env1) rhss
547 >=> lkG (D env1 e)
548 >=> lkList (lkBndr env1) bndrs
549 go (Case e b ty as) -- See Note [Empty case alternatives]
550 | null as = cm_ecase >.> lkG (D env e) >=> lkG (D env ty)
551 | otherwise = cm_case >.> lkG (D env e)
552 >=> lkList (lkA (extendCME env b)) as
553
554 xtE :: DeBruijn CoreExpr -> XT a -> CoreMapX a -> CoreMapX a
555 xtE (D env (Var v)) f m = m { cm_var = cm_var m
556 |> xtVar env v f }
557 xtE (D env (Type t)) f m = m { cm_type = cm_type m
558 |> xtG (D env t) f }
559 xtE (D env (Coercion c)) f m = m { cm_co = cm_co m
560 |> xtG (D env c) f }
561 xtE (D _ (Lit l)) f m = m { cm_lit = cm_lit m |> xtLit l f }
562 xtE (D env (Cast e c)) f m = m { cm_cast = cm_cast m |> xtG (D env e)
563 |>> xtG (D env c) f }
564 xtE (D env (Tick t e)) f m = m { cm_tick = cm_tick m |> xtG (D env e)
565 |>> xtTickish t f }
566 xtE (D env (App e1 e2)) f m = m { cm_app = cm_app m |> xtG (D env e2)
567 |>> xtG (D env e1) f }
568 xtE (D env (Lam v e)) f m = m { cm_lam = cm_lam m
569 |> xtG (D (extendCME env v) e)
570 |>> xtBndr env v f }
571 xtE (D env (Let (NonRec b r) e)) f m = m { cm_letn = cm_letn m
572 |> xtG (D (extendCME env b) e)
573 |>> xtG (D env r)
574 |>> xtBndr env b f }
575 xtE (D env (Let (Rec prs) e)) f m = m { cm_letr =
576 let (bndrs,rhss) = unzip prs
577 env1 = extendCMEs env bndrs
578 in cm_letr m
579 |> xtList (xtG . D env1) rhss
580 |>> xtG (D env1 e)
581 |>> xtList (xtBndr env1)
582 bndrs f }
583 xtE (D env (Case e b ty as)) f m
584 | null as = m { cm_ecase = cm_ecase m |> xtG (D env e)
585 |>> xtG (D env ty) f }
586 | otherwise = m { cm_case = cm_case m |> xtG (D env e)
587 |>> let env1 = extendCME env b
588 in xtList (xtA env1) as f }
589
590 -- TODO: this seems a bit dodgy, see 'eqTickish'
591 type TickishMap a = Map.Map (Tickish Id) a
592 lkTickish :: Tickish Id -> TickishMap a -> Maybe a
593 lkTickish = lookupTM
594
595 xtTickish :: Tickish Id -> XT a -> TickishMap a -> TickishMap a
596 xtTickish = alterTM
597
598 ------------------------
599 data AltMap a -- A single alternative
600 = AM { am_deflt :: CoreMapG a
601 , am_data :: NameEnv (CoreMapG a)
602 , am_lit :: LiteralMap (CoreMapG a) }
603
604 instance TrieMap AltMap where
605 type Key AltMap = CoreAlt
606 emptyTM = AM { am_deflt = emptyTM
607 , am_data = emptyNameEnv
608 , am_lit = emptyLiteralMap }
609 lookupTM = lkA emptyCME
610 alterTM = xtA emptyCME
611 foldTM = fdA
612 mapTM = mapA
613
614 instance Eq (DeBruijn CoreAlt) where
615 D env1 a1 == D env2 a2 = go a1 a2 where
616 go (DEFAULT, _, rhs1) (DEFAULT, _, rhs2)
617 = D env1 rhs1 == D env2 rhs2
618 go (LitAlt lit1, _, rhs1) (LitAlt lit2, _, rhs2)
619 = lit1 == lit2 && D env1 rhs1 == D env2 rhs2
620 go (DataAlt dc1, bs1, rhs1) (DataAlt dc2, bs2, rhs2)
621 = dc1 == dc2 &&
622 D (extendCMEs env1 bs1) rhs1 == D (extendCMEs env2 bs2) rhs2
623 go _ _ = False
624
625 mapA :: (a->b) -> AltMap a -> AltMap b
626 mapA f (AM { am_deflt = adeflt, am_data = adata, am_lit = alit })
627 = AM { am_deflt = mapTM f adeflt
628 , am_data = mapNameEnv (mapTM f) adata
629 , am_lit = mapTM (mapTM f) alit }
630
631 lkA :: CmEnv -> CoreAlt -> AltMap a -> Maybe a
632 lkA env (DEFAULT, _, rhs) = am_deflt >.> lkG (D env rhs)
633 lkA env (LitAlt lit, _, rhs) = am_lit >.> lkLit lit >=> lkG (D env rhs)
634 lkA env (DataAlt dc, bs, rhs) = am_data >.> lkNamed dc
635 >=> lkG (D (extendCMEs env bs) rhs)
636
637 xtA :: CmEnv -> CoreAlt -> XT a -> AltMap a -> AltMap a
638 xtA env (DEFAULT, _, rhs) f m =
639 m { am_deflt = am_deflt m |> xtG (D env rhs) f }
640 xtA env (LitAlt l, _, rhs) f m =
641 m { am_lit = am_lit m |> xtLit l |>> xtG (D env rhs) f }
642 xtA env (DataAlt d, bs, rhs) f m =
643 m { am_data = am_data m |> xtNamed d
644 |>> xtG (D (extendCMEs env bs) rhs) f }
645
646 fdA :: (a -> b -> b) -> AltMap a -> b -> b
647 fdA k m = foldTM k (am_deflt m)
648 . foldTM (foldTM k) (am_data m)
649 . foldTM (foldTM k) (am_lit m)
650
651 {-
652 ************************************************************************
653 * *
654 Coercions
655 * *
656 ************************************************************************
657 -}
658
659 -- We should really never care about the contents of a coercion. Instead,
660 -- just look up the coercion's type.
661 newtype CoercionMap a = CoercionMap (CoercionMapG a)
662
663 instance TrieMap CoercionMap where
664 type Key CoercionMap = Coercion
665 emptyTM = CoercionMap emptyTM
666 lookupTM k (CoercionMap m) = lookupTM (deBruijnize k) m
667 alterTM k f (CoercionMap m) = CoercionMap (alterTM (deBruijnize k) f m)
668 foldTM k (CoercionMap m) = foldTM k m
669 mapTM f (CoercionMap m) = CoercionMap (mapTM f m)
670
671 type CoercionMapG = GenMap CoercionMapX
672 newtype CoercionMapX a = CoercionMapX (TypeMapX a)
673
674 instance TrieMap CoercionMapX where
675 type Key CoercionMapX = DeBruijn Coercion
676 emptyTM = CoercionMapX emptyTM
677 lookupTM = lkC
678 alterTM = xtC
679 foldTM f (CoercionMapX core_tm) = foldTM f core_tm
680 mapTM f (CoercionMapX core_tm) = CoercionMapX (mapTM f core_tm)
681
682 instance Eq (DeBruijn Coercion) where
683 D env1 co1 == D env2 co2
684 = D env1 (coercionType co1) ==
685 D env2 (coercionType co2)
686
687 lkC :: DeBruijn Coercion -> CoercionMapX a -> Maybe a
688 lkC (D env co) (CoercionMapX core_tm) = lkT (D env $ coercionType co)
689 core_tm
690
691 xtC :: DeBruijn Coercion -> XT a -> CoercionMapX a -> CoercionMapX a
692 xtC (D env co) f (CoercionMapX m)
693 = CoercionMapX (xtT (D env $ coercionType co) f m)
694
695 {-
696 ************************************************************************
697 * *
698 Types
699 * *
700 ************************************************************************
701 -}
702
703 -- | @TypeMapG a@ is a map from @DeBruijn Type@ to @a@. The extended
704 -- key makes it suitable for recursive traversal, since it can track binders,
705 -- but it is strictly internal to this module. If you are including a 'TypeMap'
706 -- inside another 'TrieMap', this is the type you want. Note that this
707 -- lookup does not do a kind-check. Thus, all keys in this map must have
708 -- the same kind.
709 type TypeMapG = GenMap TypeMapX
710
711 -- | @TypeMapX a@ is the base map from @DeBruijn Type@ to @a@, but without the
712 -- 'GenMap' optimization.
713 data TypeMapX a
714 = TM { tm_var :: VarMap a
715 , tm_app :: TypeMapG (TypeMapG a)
716 , tm_tycon :: NameEnv a
717 , tm_forall :: TypeMapG (BndrMap a) -- See Note [Binders]
718 , tm_tylit :: TyLitMap a
719 , tm_coerce :: Maybe a
720 }
721 -- Note that there is no tyconapp case; see Note [Equality on AppTys] in Type
722
723 -- | squeeze out any synonyms, convert Constraint to *, and change TyConApps
724 -- to nested AppTys. Why the last one? See Note [Equality on AppTys] in Type
725 trieMapView :: Type -> Maybe Type
726 trieMapView ty | Just ty' <- coreViewOneStarKind ty = Just ty'
727 trieMapView (TyConApp tc tys@(_:_)) = Just $ foldl AppTy (TyConApp tc []) tys
728 trieMapView (ForAllTy (Anon arg) res)
729 = Just ((TyConApp funTyCon [] `AppTy` arg) `AppTy` res)
730 trieMapView _ = Nothing
731
732 instance TrieMap TypeMapX where
733 type Key TypeMapX = DeBruijn Type
734 emptyTM = emptyT
735 lookupTM = lkT
736 alterTM = xtT
737 foldTM = fdT
738 mapTM = mapT
739
740 instance Eq (DeBruijn Type) where
741 env_t@(D env t) == env_t'@(D env' t')
742 | Just new_t <- coreViewOneStarKind t = D env new_t == env_t'
743 | Just new_t' <- coreViewOneStarKind t' = env_t == D env' new_t'
744 | otherwise
745 = case (t, t') of
746 (CastTy t1 _, _) -> D env t1 == D env t'
747 (_, CastTy t1' _) -> D env t == D env t1'
748
749 (TyVarTy v, TyVarTy v')
750 -> case (lookupCME env v, lookupCME env' v') of
751 (Just bv, Just bv') -> bv == bv'
752 (Nothing, Nothing) -> v == v'
753 _ -> False
754 -- See Note [Equality on AppTys] in Type
755 (AppTy t1 t2, s) | Just (t1', t2') <- repSplitAppTy_maybe s
756 -> D env t1 == D env' t1' && D env t2 == D env' t2'
757 (s, AppTy t1' t2') | Just (t1, t2) <- repSplitAppTy_maybe s
758 -> D env t1 == D env' t1' && D env t2 == D env' t2'
759 (ForAllTy (Anon t1) t2, ForAllTy (Anon t1') t2')
760 -> D env t1 == D env' t1' && D env t2 == D env' t2'
761 (TyConApp tc tys, TyConApp tc' tys')
762 -> tc == tc' && D env tys == D env' tys'
763 (LitTy l, LitTy l')
764 -> l == l'
765 (ForAllTy (Named tv _) ty, ForAllTy (Named tv' _) ty')
766 -> D env (tyVarKind tv) == D env' (tyVarKind tv') &&
767 D (extendCME env tv) ty == D (extendCME env' tv') ty'
768 (CoercionTy {}, CoercionTy {})
769 -> True
770 _ -> False
771
772 instance Outputable a => Outputable (TypeMapG a) where
773 ppr m = text "TypeMap elts" <+> ppr (foldTM (:) m [])
774
775 emptyT :: TypeMapX a
776 emptyT = TM { tm_var = emptyTM
777 , tm_app = EmptyMap
778 , tm_tycon = emptyNameEnv
779 , tm_forall = EmptyMap
780 , tm_tylit = emptyTyLitMap
781 , tm_coerce = Nothing }
782
783 mapT :: (a->b) -> TypeMapX a -> TypeMapX b
784 mapT f (TM { tm_var = tvar, tm_app = tapp, tm_tycon = ttycon
785 , tm_forall = tforall, tm_tylit = tlit
786 , tm_coerce = tcoerce })
787 = TM { tm_var = mapTM f tvar
788 , tm_app = mapTM (mapTM f) tapp
789 , tm_tycon = mapNameEnv f ttycon
790 , tm_forall = mapTM (mapTM f) tforall
791 , tm_tylit = mapTM f tlit
792 , tm_coerce = fmap f tcoerce }
793
794 -----------------
795 lkT :: DeBruijn Type -> TypeMapX a -> Maybe a
796 lkT (D env ty) m = go ty m
797 where
798 go ty | Just ty' <- trieMapView ty = go ty'
799 go (TyVarTy v) = tm_var >.> lkVar env v
800 go (AppTy t1 t2) = tm_app >.> lkG (D env t1)
801 >=> lkG (D env t2)
802 go (TyConApp tc []) = tm_tycon >.> lkNamed tc
803 go ty@(TyConApp _ (_:_)) = pprPanic "lkT TyConApp" (ppr ty)
804 go (LitTy l) = tm_tylit >.> lkTyLit l
805 go (ForAllTy (Named tv _) ty) = tm_forall >.> lkG (D (extendCME env tv) ty)
806 >=> lkBndr env tv
807 go ty@(ForAllTy (Anon _) _) = pprPanic "lkT FunTy" (ppr ty)
808 go (CastTy t _) = go t
809 go (CoercionTy {}) = tm_coerce
810
811 -----------------
812 xtT :: DeBruijn Type -> XT a -> TypeMapX a -> TypeMapX a
813 xtT (D env ty) f m | Just ty' <- trieMapView ty = xtT (D env ty') f m
814
815 xtT (D env (TyVarTy v)) f m = m { tm_var = tm_var m |> xtVar env v f }
816 xtT (D env (AppTy t1 t2)) f m = m { tm_app = tm_app m |> xtG (D env t1)
817 |>> xtG (D env t2) f }
818 xtT (D _ (TyConApp tc [])) f m = m { tm_tycon = tm_tycon m |> xtNamed tc f }
819 xtT (D _ (LitTy l)) f m = m { tm_tylit = tm_tylit m |> xtTyLit l f }
820 xtT (D env (CastTy t _)) f m = xtT (D env t) f m
821 xtT (D _ (CoercionTy {})) f m = m { tm_coerce = tm_coerce m |> f }
822 xtT (D env (ForAllTy (Named tv _) ty)) f m
823 = m { tm_forall = tm_forall m |> xtG (D (extendCME env tv) ty)
824 |>> xtBndr env tv f }
825 xtT (D _ ty@(TyConApp _ (_:_))) _ _ = pprPanic "xtT TyConApp" (ppr ty)
826 xtT (D _ ty@(ForAllTy (Anon _) _)) _ _ = pprPanic "xtT FunTy" (ppr ty)
827
828 fdT :: (a -> b -> b) -> TypeMapX a -> b -> b
829 fdT k m = foldTM k (tm_var m)
830 . foldTM (foldTM k) (tm_app m)
831 . foldTM k (tm_tycon m)
832 . foldTM (foldTM k) (tm_forall m)
833 . foldTyLit k (tm_tylit m)
834 . foldMaybe k (tm_coerce m)
835
836 ------------------------
837 data TyLitMap a = TLM { tlm_number :: Map.Map Integer a
838 , tlm_string :: Map.Map FastString a
839 }
840
841 instance TrieMap TyLitMap where
842 type Key TyLitMap = TyLit
843 emptyTM = emptyTyLitMap
844 lookupTM = lkTyLit
845 alterTM = xtTyLit
846 foldTM = foldTyLit
847 mapTM = mapTyLit
848
849 emptyTyLitMap :: TyLitMap a
850 emptyTyLitMap = TLM { tlm_number = Map.empty, tlm_string = Map.empty }
851
852 mapTyLit :: (a->b) -> TyLitMap a -> TyLitMap b
853 mapTyLit f (TLM { tlm_number = tn, tlm_string = ts })
854 = TLM { tlm_number = Map.map f tn, tlm_string = Map.map f ts }
855
856 lkTyLit :: TyLit -> TyLitMap a -> Maybe a
857 lkTyLit l =
858 case l of
859 NumTyLit n -> tlm_number >.> Map.lookup n
860 StrTyLit n -> tlm_string >.> Map.lookup n
861
862 xtTyLit :: TyLit -> XT a -> TyLitMap a -> TyLitMap a
863 xtTyLit l f m =
864 case l of
865 NumTyLit n -> m { tlm_number = tlm_number m |> Map.alter f n }
866 StrTyLit n -> m { tlm_string = tlm_string m |> Map.alter f n }
867
868 foldTyLit :: (a -> b -> b) -> TyLitMap a -> b -> b
869 foldTyLit l m = flip (Map.fold l) (tlm_string m)
870 . flip (Map.fold l) (tlm_number m)
871
872 -------------------------------------------------
873 -- | @TypeMap a@ is a map from 'Type' to @a@. If you are a client, this
874 -- is the type you want. The keys in this map may have different kinds.
875 newtype TypeMap a = TypeMap (TypeMapG (TypeMapG a))
876
877 lkTT :: DeBruijn Type -> TypeMap a -> Maybe a
878 lkTT (D env ty) (TypeMap m) = lkG (D env $ typeKind ty) m
879 >>= lkG (D env ty)
880
881 xtTT :: DeBruijn Type -> XT a -> TypeMap a -> TypeMap a
882 xtTT (D env ty) f (TypeMap m)
883 = TypeMap (m |> xtG (D env $ typeKind ty)
884 |>> xtG (D env ty) f)
885
886 -- Below are some client-oriented functions which operate on 'TypeMap'.
887
888 instance TrieMap TypeMap where
889 type Key TypeMap = Type
890 emptyTM = TypeMap emptyTM
891 lookupTM k m = lkTT (deBruijnize k) m
892 alterTM k f m = xtTT (deBruijnize k) f m
893 foldTM k (TypeMap m) = foldTM (foldTM k) m
894 mapTM f (TypeMap m) = TypeMap (mapTM (mapTM f) m)
895
896 foldTypeMap :: (a -> b -> b) -> b -> TypeMap a -> b
897 foldTypeMap k z m = foldTM k m z
898
899 emptyTypeMap :: TypeMap a
900 emptyTypeMap = emptyTM
901
902 lookupTypeMap :: TypeMap a -> Type -> Maybe a
903 lookupTypeMap cm t = lookupTM t cm
904
905 extendTypeMap :: TypeMap a -> Type -> a -> TypeMap a
906 extendTypeMap m t v = alterTM t (const (Just v)) m
907
908 -- | A 'LooseTypeMap' doesn't do a kind-check. Thus, when lookup up (t |> g),
909 -- you'll find entries inserted under (t), even if (g) is non-reflexive.
910 newtype LooseTypeMap a
911 = LooseTypeMap (TypeMapG a)
912
913 instance TrieMap LooseTypeMap where
914 type Key LooseTypeMap = Type
915 emptyTM = LooseTypeMap emptyTM
916 lookupTM k (LooseTypeMap m) = lookupTM (deBruijnize k) m
917 alterTM k f (LooseTypeMap m) = LooseTypeMap (alterTM (deBruijnize k) f m)
918 foldTM f (LooseTypeMap m) = foldTM f m
919 mapTM f (LooseTypeMap m) = LooseTypeMap (mapTM f m)
920
921 {-
922 ************************************************************************
923 * *
924 Variables
925 * *
926 ************************************************************************
927 -}
928
929 type BoundVar = Int -- Bound variables are deBruijn numbered
930 type BoundVarMap a = IntMap.IntMap a
931
932 data CmEnv = CME { cme_next :: BoundVar
933 , cme_env :: VarEnv BoundVar }
934
935 emptyCME :: CmEnv
936 emptyCME = CME { cme_next = 0, cme_env = emptyVarEnv }
937
938 extendCME :: CmEnv -> Var -> CmEnv
939 extendCME (CME { cme_next = bv, cme_env = env }) v
940 = CME { cme_next = bv+1, cme_env = extendVarEnv env v bv }
941
942 extendCMEs :: CmEnv -> [Var] -> CmEnv
943 extendCMEs env vs = foldl extendCME env vs
944
945 lookupCME :: CmEnv -> Var -> Maybe BoundVar
946 lookupCME (CME { cme_env = env }) v = lookupVarEnv env v
947
948 -- | @DeBruijn a@ represents @a@ modulo alpha-renaming. This is achieved
949 -- by equipping the value with a 'CmEnv', which tracks an on-the-fly deBruijn
950 -- numbering. This allows us to define an 'Eq' instance for @DeBruijn a@, even
951 -- if this was not (easily) possible for @a@. Note: we purposely don't
952 -- export the constructor. Make a helper function if you find yourself
953 -- needing it.
954 data DeBruijn a = D CmEnv a
955
956 -- | Synthesizes a @DeBruijn a@ from an @a@, by assuming that there are no
957 -- bound binders (an empty 'CmEnv'). This is usually what you want if there
958 -- isn't already a 'CmEnv' in scope.
959 deBruijnize :: a -> DeBruijn a
960 deBruijnize = D emptyCME
961
962 instance Eq (DeBruijn a) => Eq (DeBruijn [a]) where
963 D _ [] == D _ [] = True
964 D env (x:xs) == D env' (x':xs') = D env x == D env' x' &&
965 D env xs == D env' xs'
966 _ == _ = False
967
968 --------- Variable binders -------------
969
970 -- | A 'BndrMap' is a 'TypeMapG' which allows us to distinguish between
971 -- binding forms whose binders have different types. For example,
972 -- if we are doing a 'TrieMap' lookup on @\(x :: Int) -> ()@, we should
973 -- not pick up an entry in the 'TrieMap' for @\(x :: Bool) -> ()@:
974 -- we can disambiguate this by matching on the type (or kind, if this
975 -- a binder in a type) of the binder.
976 type BndrMap = TypeMapG
977
978 -- Note [Binders]
979 -- ~~~~~~~~~~~~~~
980 -- We need to use 'BndrMap' for 'Coercion', 'CoreExpr' AND 'Type', since all
981 -- of these data types have binding forms.
982
983 lkBndr :: CmEnv -> Var -> BndrMap a -> Maybe a
984 lkBndr env v m = lkG (D env (varType v)) m
985
986 xtBndr :: CmEnv -> Var -> XT a -> BndrMap a -> BndrMap a
987 xtBndr env v f = xtG (D env (varType v)) f
988
989 --------- Variable occurrence -------------
990 data VarMap a = VM { vm_bvar :: BoundVarMap a -- Bound variable
991 , vm_fvar :: VarEnv a } -- Free variable
992
993 instance TrieMap VarMap where
994 type Key VarMap = Var
995 emptyTM = VM { vm_bvar = IntMap.empty, vm_fvar = emptyVarEnv }
996 lookupTM = lkVar emptyCME
997 alterTM = xtVar emptyCME
998 foldTM = fdVar
999 mapTM = mapVar
1000
1001 mapVar :: (a->b) -> VarMap a -> VarMap b
1002 mapVar f (VM { vm_bvar = bv, vm_fvar = fv })
1003 = VM { vm_bvar = mapTM f bv, vm_fvar = mapVarEnv f fv }
1004
1005 lkVar :: CmEnv -> Var -> VarMap a -> Maybe a
1006 lkVar env v
1007 | Just bv <- lookupCME env v = vm_bvar >.> lookupTM bv
1008 | otherwise = vm_fvar >.> lkFreeVar v
1009
1010 xtVar :: CmEnv -> Var -> XT a -> VarMap a -> VarMap a
1011 xtVar env v f m
1012 | Just bv <- lookupCME env v = m { vm_bvar = vm_bvar m |> xtInt bv f }
1013 | otherwise = m { vm_fvar = vm_fvar m |> xtFreeVar v f }
1014
1015 fdVar :: (a -> b -> b) -> VarMap a -> b -> b
1016 fdVar k m = foldTM k (vm_bvar m)
1017 . foldTM k (vm_fvar m)
1018
1019 lkFreeVar :: Var -> VarEnv a -> Maybe a
1020 lkFreeVar var env = lookupVarEnv env var
1021
1022 xtFreeVar :: Var -> XT a -> VarEnv a -> VarEnv a
1023 xtFreeVar v f m = alterVarEnv f m v