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