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