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