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