TcTypeable: Try to reuse KindReps
[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.
790 type TypeMapG = GenMap TypeMapX
791
792 -- | @TypeMapX a@ is the base map from @DeBruijn Type@ to @a@, but without the
793 -- 'GenMap' optimization.
794 data TypeMapX a
795 = TM { tm_var :: VarMap a
796 , tm_app :: TypeMapG (TypeMapG a)
797 , tm_tycon :: DNameEnv a
798 , tm_forall :: TypeMapG (BndrMap a) -- See Note [Binders]
799 , tm_tylit :: TyLitMap a
800 , tm_coerce :: Maybe a
801 }
802 -- Note that there is no tyconapp case; see Note [Equality on AppTys] in Type
803
804 -- | squeeze out any synonyms, convert Constraint to *, and change TyConApps
805 -- to nested AppTys. Why the last one? See Note [Equality on AppTys] in Type
806 trieMapView :: Type -> Maybe Type
807 trieMapView ty | Just ty' <- coreViewOneStarKind ty = Just ty'
808 trieMapView ty
809 | Just (tc, tys@(_:_)) <- splitTyConApp_maybe ty
810 = Just $ foldl AppTy (TyConApp tc []) tys
811 trieMapView _ = Nothing
812
813 instance TrieMap TypeMapX where
814 type Key TypeMapX = DeBruijn Type
815 emptyTM = emptyT
816 lookupTM = lkT
817 alterTM = xtT
818 foldTM = fdT
819 mapTM = mapT
820
821 instance Eq (DeBruijn Type) where
822 env_t@(D env t) == env_t'@(D env' t')
823 | Just new_t <- coreViewOneStarKind t = D env new_t == env_t'
824 | Just new_t' <- coreViewOneStarKind t' = env_t == D env' new_t'
825 | otherwise
826 = case (t, t') of
827 (CastTy t1 _, _) -> D env t1 == D env t'
828 (_, CastTy t1' _) -> D env t == D env t1'
829
830 (TyVarTy v, TyVarTy v')
831 -> case (lookupCME env v, lookupCME env' v') of
832 (Just bv, Just bv') -> bv == bv'
833 (Nothing, Nothing) -> v == v'
834 _ -> False
835 -- See Note [Equality on AppTys] in Type
836 (AppTy t1 t2, s) | Just (t1', t2') <- repSplitAppTy_maybe s
837 -> D env t1 == D env' t1' && D env t2 == D env' t2'
838 (s, AppTy t1' t2') | Just (t1, t2) <- repSplitAppTy_maybe s
839 -> D env t1 == D env' t1' && D env t2 == D env' t2'
840 (FunTy t1 t2, FunTy t1' t2')
841 -> D env t1 == D env' t1' && D env t2 == D env' t2'
842 (TyConApp tc tys, TyConApp tc' tys')
843 -> tc == tc' && D env tys == D env' tys'
844 (LitTy l, LitTy l')
845 -> l == l'
846 (ForAllTy (TvBndr tv _) ty, ForAllTy (TvBndr tv' _) ty')
847 -> D env (tyVarKind tv) == D env' (tyVarKind tv') &&
848 D (extendCME env tv) ty == D (extendCME env' tv') ty'
849 (CoercionTy {}, CoercionTy {})
850 -> True
851 _ -> False
852
853 instance {-# OVERLAPPING #-}
854 Outputable a => Outputable (TypeMapG a) where
855 ppr m = text "TypeMap elts" <+> ppr (foldTM (:) m [])
856
857 emptyT :: TypeMapX a
858 emptyT = TM { tm_var = emptyTM
859 , tm_app = EmptyMap
860 , tm_tycon = emptyDNameEnv
861 , tm_forall = EmptyMap
862 , tm_tylit = emptyTyLitMap
863 , tm_coerce = Nothing }
864
865 mapT :: (a->b) -> TypeMapX a -> TypeMapX b
866 mapT f (TM { tm_var = tvar, tm_app = tapp, tm_tycon = ttycon
867 , tm_forall = tforall, tm_tylit = tlit
868 , tm_coerce = tcoerce })
869 = TM { tm_var = mapTM f tvar
870 , tm_app = mapTM (mapTM f) tapp
871 , tm_tycon = mapTM f ttycon
872 , tm_forall = mapTM (mapTM f) tforall
873 , tm_tylit = mapTM f tlit
874 , tm_coerce = fmap f tcoerce }
875
876 -----------------
877 lkT :: DeBruijn Type -> TypeMapX a -> Maybe a
878 lkT (D env ty) m = go ty m
879 where
880 go ty | Just ty' <- trieMapView ty = go ty'
881 go (TyVarTy v) = tm_var >.> lkVar env v
882 go (AppTy t1 t2) = tm_app >.> lkG (D env t1)
883 >=> lkG (D env t2)
884 go (TyConApp tc []) = tm_tycon >.> lkDNamed tc
885 go ty@(TyConApp _ (_:_)) = pprPanic "lkT TyConApp" (ppr ty)
886 go (LitTy l) = tm_tylit >.> lkTyLit l
887 go (ForAllTy (TvBndr tv _) ty) = tm_forall >.> lkG (D (extendCME env tv) ty)
888 >=> lkBndr env tv
889 go ty@(FunTy {}) = pprPanic "lkT FunTy" (ppr ty)
890 go (CastTy t _) = go t
891 go (CoercionTy {}) = tm_coerce
892
893 -----------------
894 xtT :: DeBruijn Type -> XT a -> TypeMapX a -> TypeMapX a
895 xtT (D env ty) f m | Just ty' <- trieMapView ty = xtT (D env ty') f m
896
897 xtT (D env (TyVarTy v)) f m = m { tm_var = tm_var m |> xtVar env v f }
898 xtT (D env (AppTy t1 t2)) f m = m { tm_app = tm_app m |> xtG (D env t1)
899 |>> xtG (D env t2) f }
900 xtT (D _ (TyConApp tc [])) f m = m { tm_tycon = tm_tycon m |> xtDNamed tc f }
901 xtT (D _ (LitTy l)) f m = m { tm_tylit = tm_tylit m |> xtTyLit l f }
902 xtT (D env (CastTy t _)) f m = xtT (D env t) f m
903 xtT (D _ (CoercionTy {})) f m = m { tm_coerce = tm_coerce m |> f }
904 xtT (D env (ForAllTy (TvBndr tv _) ty)) f m
905 = m { tm_forall = tm_forall m |> xtG (D (extendCME env tv) ty)
906 |>> xtBndr env tv f }
907 xtT (D _ ty@(TyConApp _ (_:_))) _ _ = pprPanic "xtT TyConApp" (ppr ty)
908 xtT (D _ ty@(FunTy {})) _ _ = pprPanic "xtT FunTy" (ppr ty)
909
910 fdT :: (a -> b -> b) -> TypeMapX a -> b -> b
911 fdT k m = foldTM k (tm_var m)
912 . foldTM (foldTM k) (tm_app m)
913 . foldTM k (tm_tycon m)
914 . foldTM (foldTM k) (tm_forall m)
915 . foldTyLit k (tm_tylit m)
916 . foldMaybe k (tm_coerce m)
917
918 ------------------------
919 data TyLitMap a = TLM { tlm_number :: Map.Map Integer a
920 , tlm_string :: Map.Map FastString a
921 }
922
923 instance TrieMap TyLitMap where
924 type Key TyLitMap = TyLit
925 emptyTM = emptyTyLitMap
926 lookupTM = lkTyLit
927 alterTM = xtTyLit
928 foldTM = foldTyLit
929 mapTM = mapTyLit
930
931 emptyTyLitMap :: TyLitMap a
932 emptyTyLitMap = TLM { tlm_number = Map.empty, tlm_string = Map.empty }
933
934 mapTyLit :: (a->b) -> TyLitMap a -> TyLitMap b
935 mapTyLit f (TLM { tlm_number = tn, tlm_string = ts })
936 = TLM { tlm_number = Map.map f tn, tlm_string = Map.map f ts }
937
938 lkTyLit :: TyLit -> TyLitMap a -> Maybe a
939 lkTyLit l =
940 case l of
941 NumTyLit n -> tlm_number >.> Map.lookup n
942 StrTyLit n -> tlm_string >.> Map.lookup n
943
944 xtTyLit :: TyLit -> XT a -> TyLitMap a -> TyLitMap a
945 xtTyLit l f m =
946 case l of
947 NumTyLit n -> m { tlm_number = tlm_number m |> Map.alter f n }
948 StrTyLit n -> m { tlm_string = tlm_string m |> Map.alter f n }
949
950 foldTyLit :: (a -> b -> b) -> TyLitMap a -> b -> b
951 foldTyLit l m = flip (Map.foldr l) (tlm_string m)
952 . flip (Map.foldr l) (tlm_number m)
953
954 -------------------------------------------------
955 -- | @TypeMap a@ is a map from 'Type' to @a@. If you are a client, this
956 -- is the type you want. The keys in this map may have different kinds.
957 newtype TypeMap a = TypeMap (TypeMapG (TypeMapG a))
958
959 lkTT :: DeBruijn Type -> TypeMap a -> Maybe a
960 lkTT (D env ty) (TypeMap m) = lkG (D env $ typeKind ty) m
961 >>= lkG (D env ty)
962
963 xtTT :: DeBruijn Type -> XT a -> TypeMap a -> TypeMap a
964 xtTT (D env ty) f (TypeMap m)
965 = TypeMap (m |> xtG (D env $ typeKind ty)
966 |>> xtG (D env ty) f)
967
968 -- Below are some client-oriented functions which operate on 'TypeMap'.
969
970 instance TrieMap TypeMap where
971 type Key TypeMap = Type
972 emptyTM = TypeMap emptyTM
973 lookupTM k m = lkTT (deBruijnize k) m
974 alterTM k f m = xtTT (deBruijnize k) f m
975 foldTM k (TypeMap m) = foldTM (foldTM k) m
976 mapTM f (TypeMap m) = TypeMap (mapTM (mapTM f) m)
977
978 foldTypeMap :: (a -> b -> b) -> b -> TypeMap a -> b
979 foldTypeMap k z m = foldTM k m z
980
981 emptyTypeMap :: TypeMap a
982 emptyTypeMap = emptyTM
983
984 lookupTypeMap :: TypeMap a -> Type -> Maybe a
985 lookupTypeMap cm t = lookupTM t cm
986
987 extendTypeMap :: TypeMap a -> Type -> a -> TypeMap a
988 extendTypeMap m t v = alterTM t (const (Just v)) m
989
990 lookupTypeMapWithScope :: TypeMap a -> CmEnv -> Type -> Maybe a
991 lookupTypeMapWithScope m cm t = lkTT (D cm t) m
992
993 -- | Extend a 'TypeMap' with a type in the given context.
994 -- @extendTypeMapWithScope m (mkDeBruijnContext [a,b,c]) t v@ is equivalent to
995 -- @extendTypeMap m (forall a b c. t) v@, but allows reuse of the context over
996 -- multiple insertions.
997 extendTypeMapWithScope :: TypeMap a -> CmEnv -> Type -> a -> TypeMap a
998 extendTypeMapWithScope m cm t v = xtTT (D cm t) (const (Just v)) m
999
1000 -- | Construct a deBruijn environment with the given variables in scope.
1001 -- e.g. @mkDeBruijnEnv [a,b,c]@ constructs a context @forall a b c.@
1002 mkDeBruijnContext :: [Var] -> CmEnv
1003 mkDeBruijnContext = extendCMEs emptyCME
1004
1005 -- | A 'LooseTypeMap' doesn't do a kind-check. Thus, when lookup up (t |> g),
1006 -- you'll find entries inserted under (t), even if (g) is non-reflexive.
1007 newtype LooseTypeMap a
1008 = LooseTypeMap (TypeMapG a)
1009
1010 instance TrieMap LooseTypeMap where
1011 type Key LooseTypeMap = Type
1012 emptyTM = LooseTypeMap emptyTM
1013 lookupTM k (LooseTypeMap m) = lookupTM (deBruijnize k) m
1014 alterTM k f (LooseTypeMap m) = LooseTypeMap (alterTM (deBruijnize k) f m)
1015 foldTM f (LooseTypeMap m) = foldTM f m
1016 mapTM f (LooseTypeMap m) = LooseTypeMap (mapTM f m)
1017
1018 {-
1019 ************************************************************************
1020 * *
1021 Variables
1022 * *
1023 ************************************************************************
1024 -}
1025
1026 type BoundVar = Int -- Bound variables are deBruijn numbered
1027 type BoundVarMap a = IntMap.IntMap a
1028
1029 data CmEnv = CME { cme_next :: !BoundVar
1030 , cme_env :: VarEnv BoundVar }
1031
1032 emptyCME :: CmEnv
1033 emptyCME = CME { cme_next = 0, cme_env = emptyVarEnv }
1034
1035 extendCME :: CmEnv -> Var -> CmEnv
1036 extendCME (CME { cme_next = bv, cme_env = env }) v
1037 = CME { cme_next = bv+1, cme_env = extendVarEnv env v bv }
1038
1039 extendCMEs :: CmEnv -> [Var] -> CmEnv
1040 extendCMEs env vs = foldl extendCME env vs
1041
1042 lookupCME :: CmEnv -> Var -> Maybe BoundVar
1043 lookupCME (CME { cme_env = env }) v = lookupVarEnv env v
1044
1045 -- | @DeBruijn a@ represents @a@ modulo alpha-renaming. This is achieved
1046 -- by equipping the value with a 'CmEnv', which tracks an on-the-fly deBruijn
1047 -- numbering. This allows us to define an 'Eq' instance for @DeBruijn a@, even
1048 -- if this was not (easily) possible for @a@. Note: we purposely don't
1049 -- export the constructor. Make a helper function if you find yourself
1050 -- needing it.
1051 data DeBruijn a = D CmEnv a
1052
1053 -- | Synthesizes a @DeBruijn a@ from an @a@, by assuming that there are no
1054 -- bound binders (an empty 'CmEnv'). This is usually what you want if there
1055 -- isn't already a 'CmEnv' in scope.
1056 deBruijnize :: a -> DeBruijn a
1057 deBruijnize = D emptyCME
1058
1059 instance Eq (DeBruijn a) => Eq (DeBruijn [a]) where
1060 D _ [] == D _ [] = True
1061 D env (x:xs) == D env' (x':xs') = D env x == D env' x' &&
1062 D env xs == D env' xs'
1063 _ == _ = False
1064
1065 --------- Variable binders -------------
1066
1067 -- | A 'BndrMap' is a 'TypeMapG' which allows us to distinguish between
1068 -- binding forms whose binders have different types. For example,
1069 -- if we are doing a 'TrieMap' lookup on @\(x :: Int) -> ()@, we should
1070 -- not pick up an entry in the 'TrieMap' for @\(x :: Bool) -> ()@:
1071 -- we can disambiguate this by matching on the type (or kind, if this
1072 -- a binder in a type) of the binder.
1073 type BndrMap = TypeMapG
1074
1075 -- Note [Binders]
1076 -- ~~~~~~~~~~~~~~
1077 -- We need to use 'BndrMap' for 'Coercion', 'CoreExpr' AND 'Type', since all
1078 -- of these data types have binding forms.
1079
1080 lkBndr :: CmEnv -> Var -> BndrMap a -> Maybe a
1081 lkBndr env v m = lkG (D env (varType v)) m
1082
1083 xtBndr :: CmEnv -> Var -> XT a -> BndrMap a -> BndrMap a
1084 xtBndr env v f = xtG (D env (varType v)) f
1085
1086 --------- Variable occurrence -------------
1087 data VarMap a = VM { vm_bvar :: BoundVarMap a -- Bound variable
1088 , vm_fvar :: DVarEnv a } -- Free variable
1089
1090 instance TrieMap VarMap where
1091 type Key VarMap = Var
1092 emptyTM = VM { vm_bvar = IntMap.empty, vm_fvar = emptyDVarEnv }
1093 lookupTM = lkVar emptyCME
1094 alterTM = xtVar emptyCME
1095 foldTM = fdVar
1096 mapTM = mapVar
1097
1098 mapVar :: (a->b) -> VarMap a -> VarMap b
1099 mapVar f (VM { vm_bvar = bv, vm_fvar = fv })
1100 = VM { vm_bvar = mapTM f bv, vm_fvar = mapTM f fv }
1101
1102 lkVar :: CmEnv -> Var -> VarMap a -> Maybe a
1103 lkVar env v
1104 | Just bv <- lookupCME env v = vm_bvar >.> lookupTM bv
1105 | otherwise = vm_fvar >.> lkDFreeVar v
1106
1107 xtVar :: CmEnv -> Var -> XT a -> VarMap a -> VarMap a
1108 xtVar env v f m
1109 | Just bv <- lookupCME env v = m { vm_bvar = vm_bvar m |> alterTM bv f }
1110 | otherwise = m { vm_fvar = vm_fvar m |> xtDFreeVar v f }
1111
1112 fdVar :: (a -> b -> b) -> VarMap a -> b -> b
1113 fdVar k m = foldTM k (vm_bvar m)
1114 . foldTM k (vm_fvar m)
1115
1116 lkDFreeVar :: Var -> DVarEnv a -> Maybe a
1117 lkDFreeVar var env = lookupDVarEnv env var
1118
1119 xtDFreeVar :: Var -> XT a -> DVarEnv a -> DVarEnv a
1120 xtDFreeVar v f m = alterDVarEnv f m v