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