never executed always true always false
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 -}
5
6 {-# LANGUAGE FlexibleContexts #-}
7 {-# LANGUAGE FlexibleInstances #-}
8 {-# LANGUAGE TypeFamilies #-}
9
10 module GHC.Core.Map.Type (
11 -- * Re-export generic interface
12 TrieMap(..), XT,
13
14 -- * Maps over 'Type's
15 TypeMap, emptyTypeMap, extendTypeMap, lookupTypeMap, foldTypeMap,
16 LooseTypeMap,
17 -- ** With explicit scoping
18 CmEnv, lookupCME, extendTypeMapWithScope, lookupTypeMapWithScope,
19 mkDeBruijnContext, extendCME, extendCMEs, emptyCME,
20
21 -- * Utilities for use by friends only
22 TypeMapG, CoercionMapG,
23
24 DeBruijn(..), deBruijnize,
25
26 BndrMap, xtBndr, lkBndr,
27 VarMap, xtVar, lkVar, lkDFreeVar, xtDFreeVar,
28
29 xtDNamed, lkDNamed
30
31 ) where
32
33 -- This module is separate from GHC.Core.Map.Expr to avoid a module loop
34 -- between GHC.Core.Unify (which depends on this module) and GHC.Core
35
36 import GHC.Prelude
37
38 import GHC.Core.Type
39 import GHC.Core.Coercion
40 import GHC.Core.TyCo.Rep
41 import GHC.Data.TrieMap
42
43 import GHC.Data.FastString
44 import GHC.Types.Name
45 import GHC.Types.Name.Env
46 import GHC.Types.Var
47 import GHC.Types.Var.Env
48 import GHC.Types.Unique.FM
49 import GHC.Utils.Outputable
50
51 import GHC.Utils.Panic
52
53 import qualified Data.Map as Map
54 import qualified Data.IntMap as IntMap
55
56 import Control.Monad ( (>=>) )
57 import GHC.Data.Maybe
58
59 -- NB: Be careful about RULES and type families (#5821). So we should make sure
60 -- to specify @Key TypeMapX@ (and not @DeBruijn Type@, the reduced form)
61
62 {-# SPECIALIZE lkG :: Key TypeMapX -> TypeMapG a -> Maybe a #-}
63 {-# SPECIALIZE lkG :: Key CoercionMapX -> CoercionMapG a -> Maybe a #-}
64
65 {-# SPECIALIZE xtG :: Key TypeMapX -> XT a -> TypeMapG a -> TypeMapG a #-}
66 {-# SPECIALIZE xtG :: Key CoercionMapX -> XT a -> CoercionMapG a -> CoercionMapG a #-}
67
68 {-# SPECIALIZE mapG :: (a -> b) -> TypeMapG a -> TypeMapG b #-}
69 {-# SPECIALIZE mapG :: (a -> b) -> CoercionMapG a -> CoercionMapG b #-}
70
71 {-# SPECIALIZE fdG :: (a -> b -> b) -> TypeMapG a -> b -> b #-}
72 {-# SPECIALIZE fdG :: (a -> b -> b) -> CoercionMapG a -> b -> b #-}
73
74 {-
75 ************************************************************************
76 * *
77 Coercions
78 * *
79 ************************************************************************
80 -}
81
82 -- We should really never care about the contents of a coercion. Instead,
83 -- just look up the coercion's type.
84 newtype CoercionMap a = CoercionMap (CoercionMapG a)
85
86 instance TrieMap CoercionMap where
87 type Key CoercionMap = Coercion
88 emptyTM = CoercionMap emptyTM
89 lookupTM k (CoercionMap m) = lookupTM (deBruijnize k) m
90 alterTM k f (CoercionMap m) = CoercionMap (alterTM (deBruijnize k) f m)
91 foldTM k (CoercionMap m) = foldTM k m
92 mapTM f (CoercionMap m) = CoercionMap (mapTM f m)
93 filterTM f (CoercionMap m) = CoercionMap (filterTM f m)
94
95 type CoercionMapG = GenMap CoercionMapX
96 newtype CoercionMapX a = CoercionMapX (TypeMapX a)
97
98 instance TrieMap CoercionMapX where
99 type Key CoercionMapX = DeBruijn Coercion
100 emptyTM = CoercionMapX emptyTM
101 lookupTM = lkC
102 alterTM = xtC
103 foldTM f (CoercionMapX core_tm) = foldTM f core_tm
104 mapTM f (CoercionMapX core_tm) = CoercionMapX (mapTM f core_tm)
105 filterTM f (CoercionMapX core_tm) = CoercionMapX (filterTM f core_tm)
106
107 instance Eq (DeBruijn Coercion) where
108 D env1 co1 == D env2 co2
109 = D env1 (coercionType co1) ==
110 D env2 (coercionType co2)
111
112 lkC :: DeBruijn Coercion -> CoercionMapX a -> Maybe a
113 lkC (D env co) (CoercionMapX core_tm) = lkT (D env $ coercionType co)
114 core_tm
115
116 xtC :: DeBruijn Coercion -> XT a -> CoercionMapX a -> CoercionMapX a
117 xtC (D env co) f (CoercionMapX m)
118 = CoercionMapX (xtT (D env $ coercionType co) f m)
119
120 {-
121 ************************************************************************
122 * *
123 Types
124 * *
125 ************************************************************************
126 -}
127
128 -- | @TypeMapG a@ is a map from @DeBruijn Type@ to @a@. The extended
129 -- key makes it suitable for recursive traversal, since it can track binders,
130 -- but it is strictly internal to this module. If you are including a 'TypeMap'
131 -- inside another 'TrieMap', this is the type you want. Note that this
132 -- lookup does not do a kind-check. Thus, all keys in this map must have
133 -- the same kind. Also note that this map respects the distinction between
134 -- @Type@ and @Constraint@, despite the fact that they are equivalent type
135 -- synonyms in Core.
136 type TypeMapG = GenMap TypeMapX
137
138 -- | @TypeMapX a@ is the base map from @DeBruijn Type@ to @a@, but without the
139 -- 'GenMap' optimization. See Note [Computing equality on types] in GHC.Core.Type.
140 data TypeMapX a
141 = TM { tm_var :: VarMap a
142 , tm_app :: TypeMapG (TypeMapG a) -- Note [Equality on AppTys] in GHC.Core.Type
143 , tm_tycon :: DNameEnv a
144
145 -- only InvisArg arrows here
146 , tm_funty :: TypeMapG (TypeMapG (TypeMapG a))
147 -- keyed on the argument, result rep, and result
148 -- constraints are never linear-restricted and are always lifted
149 -- See also Note [Equality on FunTys] in GHC.Core.TyCo.Rep
150
151 , tm_forall :: TypeMapG (BndrMap a) -- See Note [Binders] in GHC.Core.Map.Expr
152 , tm_tylit :: TyLitMap a
153 , tm_coerce :: Maybe a
154 }
155 -- Note that there is no tyconapp case; see Note [Equality on AppTys] in GHC.Core.Type
156
157 -- | Squeeze out any synonyms, and change TyConApps to nested AppTys. Why the
158 -- last one? See Note [Equality on AppTys] in GHC.Core.Type
159 --
160 -- Note, however, that we keep Constraint and Type apart here, despite the fact
161 -- that they are both synonyms of TYPE 'LiftedRep (see #11715).
162 --
163 -- We also keep (Eq a => a) as a FunTy, distinct from ((->) (Eq a) a).
164 trieMapView :: Type -> Maybe Type
165 trieMapView ty
166 -- First check for TyConApps that need to be expanded to
167 -- AppTy chains.
168 | Just (tc, tys@(_:_)) <- tcSplitTyConApp_maybe ty
169 = Just $ foldl' AppTy (mkTyConTy tc) tys
170
171 -- Then resolve any remaining nullary synonyms.
172 | Just ty' <- tcView ty = Just ty'
173 trieMapView _ = Nothing
174
175 instance TrieMap TypeMapX where
176 type Key TypeMapX = DeBruijn Type
177 emptyTM = emptyT
178 lookupTM = lkT
179 alterTM = xtT
180 foldTM = fdT
181 mapTM = mapT
182 filterTM = filterT
183
184 instance Eq (DeBruijn Type) where
185 env_t@(D env t) == env_t'@(D env' t')
186 | Just new_t <- tcView t = D env new_t == env_t'
187 | Just new_t' <- tcView t' = env_t == D env' new_t'
188 | otherwise
189 = case (t, t') of
190 (CastTy t1 _, _) -> D env t1 == D env t'
191 (_, CastTy t1' _) -> D env t == D env t1'
192
193 (TyVarTy v, TyVarTy v')
194 -> case (lookupCME env v, lookupCME env' v') of
195 (Just bv, Just bv') -> bv == bv'
196 (Nothing, Nothing) -> v == v'
197 _ -> False
198 -- See Note [Equality on AppTys] in GHC.Core.Type
199 (AppTy t1 t2, s) | Just (t1', t2') <- repSplitAppTy_maybe s
200 -> D env t1 == D env' t1' && D env t2 == D env' t2'
201 (s, AppTy t1' t2') | Just (t1, t2) <- repSplitAppTy_maybe s
202 -> D env t1 == D env' t1' && D env t2 == D env' t2'
203 (FunTy v1 w1 t1 t2, FunTy v1' w1' t1' t2')
204 -> v1 == v1' &&
205 D env w1 == D env w1' &&
206 D env t1 == D env' t1' &&
207 D env t2 == D env' t2'
208 (TyConApp tc tys, TyConApp tc' tys')
209 -> tc == tc' && D env tys == D env' tys'
210 (LitTy l, LitTy l')
211 -> l == l'
212 (ForAllTy (Bndr tv _) ty, ForAllTy (Bndr tv' _) ty')
213 -> D env (varType tv) == D env' (varType tv') &&
214 D (extendCME env tv) ty == D (extendCME env' tv') ty'
215 (CoercionTy {}, CoercionTy {})
216 -> True
217 _ -> False
218
219 instance {-# OVERLAPPING #-}
220 Outputable a => Outputable (TypeMapG a) where
221 ppr m = text "TypeMap elts" <+> ppr (foldTM (:) m [])
222
223 emptyT :: TypeMapX a
224 emptyT = TM { tm_var = emptyTM
225 , tm_app = emptyTM
226 , tm_tycon = emptyDNameEnv
227 , tm_funty = emptyTM
228 , tm_forall = emptyTM
229 , tm_tylit = emptyTyLitMap
230 , tm_coerce = Nothing }
231
232 mapT :: (a->b) -> TypeMapX a -> TypeMapX b
233 mapT f (TM { tm_var = tvar, tm_app = tapp, tm_tycon = ttycon
234 , tm_funty = tfunty, tm_forall = tforall, tm_tylit = tlit
235 , tm_coerce = tcoerce })
236 = TM { tm_var = mapTM f tvar
237 , tm_app = mapTM (mapTM f) tapp
238 , tm_tycon = mapTM f ttycon
239 , tm_funty = mapTM (mapTM (mapTM f)) tfunty
240 , tm_forall = mapTM (mapTM f) tforall
241 , tm_tylit = mapTM f tlit
242 , tm_coerce = fmap f tcoerce }
243
244 -----------------
245 lkT :: DeBruijn Type -> TypeMapX a -> Maybe a
246 lkT (D env ty) m = go ty m
247 where
248 go ty | Just ty' <- trieMapView ty = go ty'
249 go (TyVarTy v) = tm_var >.> lkVar env v
250 go (AppTy t1 t2) = tm_app >.> lkG (D env t1)
251 >=> lkG (D env t2)
252 go (TyConApp tc []) = tm_tycon >.> lkDNamed tc
253 go ty@(TyConApp _ (_:_)) = pprPanic "lkT TyConApp" (ppr ty)
254 go (LitTy l) = tm_tylit >.> lkTyLit l
255 go (ForAllTy (Bndr tv _) ty) = tm_forall >.> lkG (D (extendCME env tv) ty)
256 >=> lkBndr env tv
257 go (FunTy InvisArg _ arg res)
258 | Just res_rep <- getRuntimeRep_maybe res
259 = tm_funty >.> lkG (D env arg)
260 >=> lkG (D env res_rep)
261 >=> lkG (D env res)
262 go ty@(FunTy {}) = pprPanic "lkT FunTy" (ppr ty)
263 go (CastTy t _) = go t
264 go (CoercionTy {}) = tm_coerce
265
266 -----------------
267 xtT :: DeBruijn Type -> XT a -> TypeMapX a -> TypeMapX a
268 xtT (D env ty) f m | Just ty' <- trieMapView ty = xtT (D env ty') f m
269
270 xtT (D env (TyVarTy v)) f m = m { tm_var = tm_var m |> xtVar env v f }
271 xtT (D env (AppTy t1 t2)) f m = m { tm_app = tm_app m |> xtG (D env t1)
272 |>> xtG (D env t2) f }
273 xtT (D _ (TyConApp tc [])) f m = m { tm_tycon = tm_tycon m |> xtDNamed tc f }
274 xtT (D env (FunTy InvisArg _ t1 t2)) f m = m { tm_funty = tm_funty m |> xtG (D env t1)
275 |>> xtG (D env t2_rep)
276 |>> xtG (D env t2) f }
277 where t2_rep = expectJust "xtT FunTy InvisArg" (getRuntimeRep_maybe t2)
278 xtT (D _ (LitTy l)) f m = m { tm_tylit = tm_tylit m |> xtTyLit l f }
279 xtT (D env (CastTy t _)) f m = xtT (D env t) f m
280 xtT (D _ (CoercionTy {})) f m = m { tm_coerce = tm_coerce m |> f }
281 xtT (D env (ForAllTy (Bndr tv _) ty)) f m
282 = m { tm_forall = tm_forall m |> xtG (D (extendCME env tv) ty)
283 |>> xtBndr env tv f }
284 xtT (D _ ty@(TyConApp _ (_:_))) _ _ = pprPanic "xtT TyConApp" (ppr ty)
285 xtT (D _ ty@(FunTy {})) _ _ = pprPanic "xtT FunTy" (ppr ty)
286
287 fdT :: (a -> b -> b) -> TypeMapX a -> b -> b
288 fdT k m = foldTM k (tm_var m)
289 . foldTM (foldTM k) (tm_app m)
290 . foldTM k (tm_tycon m)
291 . foldTM (foldTM (foldTM k)) (tm_funty m)
292 . foldTM (foldTM k) (tm_forall m)
293 . foldTyLit k (tm_tylit m)
294 . foldMaybe k (tm_coerce m)
295
296 filterT :: (a -> Bool) -> TypeMapX a -> TypeMapX a
297 filterT f (TM { tm_var = tvar, tm_app = tapp, tm_tycon = ttycon
298 , tm_funty = tfunty, tm_forall = tforall, tm_tylit = tlit
299 , tm_coerce = tcoerce })
300 = TM { tm_var = filterTM f tvar
301 , tm_app = mapTM (filterTM f) tapp
302 , tm_tycon = filterTM f ttycon
303 , tm_funty = mapTM (mapTM (filterTM f)) tfunty
304 , tm_forall = mapTM (filterTM f) tforall
305 , tm_tylit = filterTM f tlit
306 , tm_coerce = filterMaybe f tcoerce }
307
308 ------------------------
309 data TyLitMap a = TLM { tlm_number :: Map.Map Integer a
310 , tlm_string :: UniqFM FastString a
311 , tlm_char :: Map.Map Char a
312 }
313
314 instance TrieMap TyLitMap where
315 type Key TyLitMap = TyLit
316 emptyTM = emptyTyLitMap
317 lookupTM = lkTyLit
318 alterTM = xtTyLit
319 foldTM = foldTyLit
320 mapTM = mapTyLit
321 filterTM = filterTyLit
322
323 emptyTyLitMap :: TyLitMap a
324 emptyTyLitMap = TLM { tlm_number = Map.empty, tlm_string = emptyUFM, tlm_char = Map.empty }
325
326 mapTyLit :: (a->b) -> TyLitMap a -> TyLitMap b
327 mapTyLit f (TLM { tlm_number = tn, tlm_string = ts, tlm_char = tc })
328 = TLM { tlm_number = Map.map f tn, tlm_string = mapUFM f ts, tlm_char = Map.map f tc }
329
330 lkTyLit :: TyLit -> TyLitMap a -> Maybe a
331 lkTyLit l =
332 case l of
333 NumTyLit n -> tlm_number >.> Map.lookup n
334 StrTyLit n -> tlm_string >.> (`lookupUFM` n)
335 CharTyLit n -> tlm_char >.> Map.lookup n
336
337 xtTyLit :: TyLit -> XT a -> TyLitMap a -> TyLitMap a
338 xtTyLit l f m =
339 case l of
340 NumTyLit n -> m { tlm_number = Map.alter f n (tlm_number m) }
341 StrTyLit n -> m { tlm_string = alterUFM f (tlm_string m) n }
342 CharTyLit n -> m { tlm_char = Map.alter f n (tlm_char m) }
343
344 foldTyLit :: (a -> b -> b) -> TyLitMap a -> b -> b
345 foldTyLit l m = flip (foldUFM l) (tlm_string m)
346 . flip (Map.foldr l) (tlm_number m)
347 . flip (Map.foldr l) (tlm_char m)
348
349 filterTyLit :: (a -> Bool) -> TyLitMap a -> TyLitMap a
350 filterTyLit f (TLM { tlm_number = tn, tlm_string = ts, tlm_char = tc })
351 = TLM { tlm_number = Map.filter f tn, tlm_string = filterUFM f ts, tlm_char = Map.filter f tc }
352
353 -------------------------------------------------
354 -- | @TypeMap a@ is a map from 'Type' to @a@. If you are a client, this
355 -- is the type you want. The keys in this map may have different kinds.
356 newtype TypeMap a = TypeMap (TypeMapG (TypeMapG a))
357
358 lkTT :: DeBruijn Type -> TypeMap a -> Maybe a
359 lkTT (D env ty) (TypeMap m) = lkG (D env $ typeKind ty) m
360 >>= lkG (D env ty)
361
362 xtTT :: DeBruijn Type -> XT a -> TypeMap a -> TypeMap a
363 xtTT (D env ty) f (TypeMap m)
364 = TypeMap (m |> xtG (D env $ typeKind ty)
365 |>> xtG (D env ty) f)
366
367 -- Below are some client-oriented functions which operate on 'TypeMap'.
368
369 instance TrieMap TypeMap where
370 type Key TypeMap = Type
371 emptyTM = TypeMap emptyTM
372 lookupTM k m = lkTT (deBruijnize k) m
373 alterTM k f m = xtTT (deBruijnize k) f m
374 foldTM k (TypeMap m) = foldTM (foldTM k) m
375 mapTM f (TypeMap m) = TypeMap (mapTM (mapTM f) m)
376 filterTM f (TypeMap m) = TypeMap (mapTM (filterTM f) m)
377
378 foldTypeMap :: (a -> b -> b) -> b -> TypeMap a -> b
379 foldTypeMap k z m = foldTM k m z
380
381 emptyTypeMap :: TypeMap a
382 emptyTypeMap = emptyTM
383
384 lookupTypeMap :: TypeMap a -> Type -> Maybe a
385 lookupTypeMap cm t = lookupTM t cm
386
387 extendTypeMap :: TypeMap a -> Type -> a -> TypeMap a
388 extendTypeMap m t v = alterTM t (const (Just v)) m
389
390 lookupTypeMapWithScope :: TypeMap a -> CmEnv -> Type -> Maybe a
391 lookupTypeMapWithScope m cm t = lkTT (D cm t) m
392
393 -- | Extend a 'TypeMap' with a type in the given context.
394 -- @extendTypeMapWithScope m (mkDeBruijnContext [a,b,c]) t v@ is equivalent to
395 -- @extendTypeMap m (forall a b c. t) v@, but allows reuse of the context over
396 -- multiple insertions.
397 extendTypeMapWithScope :: TypeMap a -> CmEnv -> Type -> a -> TypeMap a
398 extendTypeMapWithScope m cm t v = xtTT (D cm t) (const (Just v)) m
399
400 -- | Construct a deBruijn environment with the given variables in scope.
401 -- e.g. @mkDeBruijnEnv [a,b,c]@ constructs a context @forall a b c.@
402 mkDeBruijnContext :: [Var] -> CmEnv
403 mkDeBruijnContext = extendCMEs emptyCME
404
405 -- | A 'LooseTypeMap' doesn't do a kind-check. Thus, when lookup up (t |> g),
406 -- you'll find entries inserted under (t), even if (g) is non-reflexive.
407 newtype LooseTypeMap a
408 = LooseTypeMap (TypeMapG a)
409
410 instance TrieMap LooseTypeMap where
411 type Key LooseTypeMap = Type
412 emptyTM = LooseTypeMap emptyTM
413 lookupTM k (LooseTypeMap m) = lookupTM (deBruijnize k) m
414 alterTM k f (LooseTypeMap m) = LooseTypeMap (alterTM (deBruijnize k) f m)
415 foldTM f (LooseTypeMap m) = foldTM f m
416 mapTM f (LooseTypeMap m) = LooseTypeMap (mapTM f m)
417 filterTM f (LooseTypeMap m) = LooseTypeMap (filterTM f m)
418
419 {-
420 ************************************************************************
421 * *
422 Variables
423 * *
424 ************************************************************************
425 -}
426
427 type BoundVar = Int -- Bound variables are deBruijn numbered
428 type BoundVarMap a = IntMap.IntMap a
429
430 data CmEnv = CME { cme_next :: !BoundVar
431 , cme_env :: VarEnv BoundVar }
432
433 emptyCME :: CmEnv
434 emptyCME = CME { cme_next = 0, cme_env = emptyVarEnv }
435
436 extendCME :: CmEnv -> Var -> CmEnv
437 extendCME (CME { cme_next = bv, cme_env = env }) v
438 = CME { cme_next = bv+1, cme_env = extendVarEnv env v bv }
439
440 extendCMEs :: CmEnv -> [Var] -> CmEnv
441 extendCMEs env vs = foldl' extendCME env vs
442
443 lookupCME :: CmEnv -> Var -> Maybe BoundVar
444 lookupCME (CME { cme_env = env }) v = lookupVarEnv env v
445
446 -- | @DeBruijn a@ represents @a@ modulo alpha-renaming. This is achieved
447 -- by equipping the value with a 'CmEnv', which tracks an on-the-fly deBruijn
448 -- numbering. This allows us to define an 'Eq' instance for @DeBruijn a@, even
449 -- if this was not (easily) possible for @a@. Note: we purposely don't
450 -- export the constructor. Make a helper function if you find yourself
451 -- needing it.
452 data DeBruijn a = D CmEnv a
453
454 -- | Synthesizes a @DeBruijn a@ from an @a@, by assuming that there are no
455 -- bound binders (an empty 'CmEnv'). This is usually what you want if there
456 -- isn't already a 'CmEnv' in scope.
457 deBruijnize :: a -> DeBruijn a
458 deBruijnize = D emptyCME
459
460 instance Eq (DeBruijn a) => Eq (DeBruijn [a]) where
461 D _ [] == D _ [] = True
462 D env (x:xs) == D env' (x':xs') = D env x == D env' x' &&
463 D env xs == D env' xs'
464 _ == _ = False
465
466 instance Eq (DeBruijn a) => Eq (DeBruijn (Maybe a)) where
467 D _ Nothing == D _ Nothing = True
468 D env (Just x) == D env' (Just x') = D env x == D env' x'
469 _ == _ = False
470
471 --------- Variable binders -------------
472
473 -- | A 'BndrMap' is a 'TypeMapG' which allows us to distinguish between
474 -- binding forms whose binders have different types. For example,
475 -- if we are doing a 'TrieMap' lookup on @\(x :: Int) -> ()@, we should
476 -- not pick up an entry in the 'TrieMap' for @\(x :: Bool) -> ()@:
477 -- we can disambiguate this by matching on the type (or kind, if this
478 -- a binder in a type) of the binder.
479 --
480 -- We also need to do the same for multiplicity! Which, since multiplicities are
481 -- encoded simply as a 'Type', amounts to have a Trie for a pair of types. Tries
482 -- of pairs are composition.
483 data BndrMap a = BndrMap (TypeMapG (MaybeMap TypeMapG a))
484
485 instance TrieMap BndrMap where
486 type Key BndrMap = Var
487 emptyTM = BndrMap emptyTM
488 lookupTM = lkBndr emptyCME
489 alterTM = xtBndr emptyCME
490 foldTM = fdBndrMap
491 mapTM = mapBndrMap
492 filterTM = ftBndrMap
493
494 mapBndrMap :: (a -> b) -> BndrMap a -> BndrMap b
495 mapBndrMap f (BndrMap tm) = BndrMap (mapTM (mapTM f) tm)
496
497 fdBndrMap :: (a -> b -> b) -> BndrMap a -> b -> b
498 fdBndrMap f (BndrMap tm) = foldTM (foldTM f) tm
499
500
501 -- We need to use 'BndrMap' for 'Coercion', 'CoreExpr' AND 'Type', since all
502 -- of these data types have binding forms.
503
504 lkBndr :: CmEnv -> Var -> BndrMap a -> Maybe a
505 lkBndr env v (BndrMap tymap) = do
506 multmap <- lkG (D env (varType v)) tymap
507 lookupTM (D env <$> varMultMaybe v) multmap
508
509
510 xtBndr :: forall a . CmEnv -> Var -> XT a -> BndrMap a -> BndrMap a
511 xtBndr env v xt (BndrMap tymap) =
512 BndrMap (tymap |> xtG (D env (varType v)) |>> (alterTM (D env <$> varMultMaybe v) xt))
513
514 ftBndrMap :: (a -> Bool) -> BndrMap a -> BndrMap a
515 ftBndrMap f (BndrMap tm) = BndrMap (mapTM (filterTM f) tm)
516
517 --------- Variable occurrence -------------
518 data VarMap a = VM { vm_bvar :: BoundVarMap a -- Bound variable
519 , vm_fvar :: DVarEnv a } -- Free variable
520
521 instance TrieMap VarMap where
522 type Key VarMap = Var
523 emptyTM = VM { vm_bvar = IntMap.empty, vm_fvar = emptyDVarEnv }
524 lookupTM = lkVar emptyCME
525 alterTM = xtVar emptyCME
526 foldTM = fdVar
527 mapTM = mapVar
528 filterTM = ftVar
529
530 mapVar :: (a->b) -> VarMap a -> VarMap b
531 mapVar f (VM { vm_bvar = bv, vm_fvar = fv })
532 = VM { vm_bvar = mapTM f bv, vm_fvar = mapTM f fv }
533
534 lkVar :: CmEnv -> Var -> VarMap a -> Maybe a
535 lkVar env v
536 | Just bv <- lookupCME env v = vm_bvar >.> lookupTM bv
537 | otherwise = vm_fvar >.> lkDFreeVar v
538
539 xtVar :: CmEnv -> Var -> XT a -> VarMap a -> VarMap a
540 xtVar env v f m
541 | Just bv <- lookupCME env v = m { vm_bvar = vm_bvar m |> alterTM bv f }
542 | otherwise = m { vm_fvar = vm_fvar m |> xtDFreeVar v f }
543
544 fdVar :: (a -> b -> b) -> VarMap a -> b -> b
545 fdVar k m = foldTM k (vm_bvar m)
546 . foldTM k (vm_fvar m)
547
548 lkDFreeVar :: Var -> DVarEnv a -> Maybe a
549 lkDFreeVar var env = lookupDVarEnv env var
550
551 xtDFreeVar :: Var -> XT a -> DVarEnv a -> DVarEnv a
552 xtDFreeVar v f m = alterDVarEnv f m v
553
554 ftVar :: (a -> Bool) -> VarMap a -> VarMap a
555 ftVar f (VM { vm_bvar = bv, vm_fvar = fv })
556 = VM { vm_bvar = filterTM f bv, vm_fvar = filterTM f fv }
557
558 -------------------------------------------------
559 lkDNamed :: NamedThing n => n -> DNameEnv a -> Maybe a
560 lkDNamed n env = lookupDNameEnv env (getName n)
561
562 xtDNamed :: NamedThing n => n -> XT a -> DNameEnv a -> DNameEnv a
563 xtDNamed tc f m = alterDNameEnv f m (getName tc)