3d0693466a3023e3be9a3bf64e6a53d6ef328ec3
[ghc.git] / compiler / coreSyn / CoreMap.hs
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 -}
5
6 {-# LANGUAGE CPP #-}
7 {-# LANGUAGE RankNTypes #-}
8 {-# LANGUAGE TypeFamilies #-}
9 {-# LANGUAGE FlexibleContexts #-}
10 {-# LANGUAGE TypeSynonymInstances #-}
11 {-# LANGUAGE FlexibleInstances #-}
12 {-# LANGUAGE UndecidableInstances #-}
13
14 module CoreMap(
15 -- * Maps over Core expressions
16 CoreMap, emptyCoreMap, extendCoreMap, lookupCoreMap, foldCoreMap,
17 -- * Maps over 'Type's
18 TypeMap, emptyTypeMap, extendTypeMap, lookupTypeMap, foldTypeMap,
19 LooseTypeMap,
20 -- ** With explicit scoping
21 CmEnv, lookupCME, extendTypeMapWithScope, lookupTypeMapWithScope,
22 mkDeBruijnContext,
23 -- * Maps over 'Maybe' values
24 MaybeMap,
25 -- * Maps over 'List' values
26 ListMap,
27 -- * Maps over 'Literal's
28 LiteralMap,
29 -- * Map for compressing leaves. See Note [Compressed TrieMap]
30 GenMap,
31 -- * 'TrieMap' class
32 TrieMap(..), insertTM, deleteTM,
33 lkDFreeVar, xtDFreeVar,
34 lkDNamed, xtDNamed,
35 (>.>), (|>), (|>>),
36 ) where
37
38 #include "HsVersions.h"
39
40 import GhcPrelude
41
42 import TrieMap
43 import CoreSyn
44 import Coercion
45 import Name
46 import Type
47 import TyCoRep
48 import Var
49 import FastString(FastString)
50 import Util
51
52 import qualified Data.Map as Map
53 import qualified Data.IntMap as IntMap
54 import VarEnv
55 import NameEnv
56 import Outputable
57 import Control.Monad( (>=>) )
58
59 {-
60 This module implements TrieMaps over Core related data structures
61 like CoreExpr or Type. It is built on the Tries from the TrieMap
62 module.
63
64 The code is very regular and boilerplate-like, but there is
65 some neat handling of *binders*. In effect they are deBruijn
66 numbered on the fly.
67
68
69 -}
70
71 ----------------------
72 -- Recall that
73 -- Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c
74
75 -- NB: Be careful about RULES and type families (#5821). So we should make sure
76 -- to specify @Key TypeMapX@ (and not @DeBruijn Type@, the reduced form)
77
78 -- The CoreMap makes heavy use of GenMap. However the CoreMap Types are not
79 -- known when defining GenMap so we can only specialize them here.
80
81 {-# SPECIALIZE lkG :: Key TypeMapX -> TypeMapG a -> Maybe a #-}
82 {-# SPECIALIZE lkG :: Key CoercionMapX -> CoercionMapG a -> Maybe a #-}
83 {-# SPECIALIZE lkG :: Key CoreMapX -> CoreMapG a -> Maybe a #-}
84
85
86 {-# SPECIALIZE xtG :: Key TypeMapX -> XT a -> TypeMapG a -> TypeMapG a #-}
87 {-# SPECIALIZE xtG :: Key CoercionMapX -> XT a -> CoercionMapG a -> CoercionMapG a #-}
88 {-# SPECIALIZE xtG :: Key CoreMapX -> XT a -> CoreMapG a -> CoreMapG a #-}
89
90 {-# SPECIALIZE mapG :: (a -> b) -> TypeMapG a -> TypeMapG b #-}
91 {-# SPECIALIZE mapG :: (a -> b) -> CoercionMapG a -> CoercionMapG b #-}
92 {-# SPECIALIZE mapG :: (a -> b) -> CoreMapG a -> CoreMapG b #-}
93
94 {-# SPECIALIZE fdG :: (a -> b -> b) -> TypeMapG a -> b -> b #-}
95 {-# SPECIALIZE fdG :: (a -> b -> b) -> CoercionMapG a -> b -> b #-}
96 {-# SPECIALIZE fdG :: (a -> b -> b) -> CoreMapG a -> b -> b #-}
97
98
99 {-
100 ************************************************************************
101 * *
102 CoreMap
103 * *
104 ************************************************************************
105 -}
106
107 lkDNamed :: NamedThing n => n -> DNameEnv a -> Maybe a
108 lkDNamed n env = lookupDNameEnv env (getName n)
109
110 xtDNamed :: NamedThing n => n -> XT a -> DNameEnv a -> DNameEnv a
111 xtDNamed tc f m = alterDNameEnv f m (getName tc)
112
113
114 {-
115 Note [Binders]
116 ~~~~~~~~~~~~~~
117 * In general we check binders as late as possible because types are
118 less likely to differ than expression structure. That's why
119 cm_lam :: CoreMapG (TypeMapG a)
120 rather than
121 cm_lam :: TypeMapG (CoreMapG a)
122
123 * We don't need to look at the type of some binders, notably
124 - the case binder in (Case _ b _ _)
125 - the binders in an alternative
126 because they are totally fixed by the context
127
128 Note [Empty case alternatives]
129 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
130 * For a key (Case e b ty (alt:alts)) we don't need to look the return type
131 'ty', because every alternative has that type.
132
133 * For a key (Case e b ty []) we MUST look at the return type 'ty', because
134 otherwise (Case (error () "urk") _ Int []) would compare equal to
135 (Case (error () "urk") _ Bool [])
136 which is utterly wrong (Trac #6097)
137
138 We could compare the return type regardless, but the wildly common case
139 is that it's unnecessary, so we have two fields (cm_case and cm_ecase)
140 for the two possibilities. Only cm_ecase looks at the type.
141
142 See also Note [Empty case alternatives] in CoreSyn.
143 -}
144
145 -- | @CoreMap a@ is a map from 'CoreExpr' to @a@. If you are a client, this
146 -- is the type you want.
147 newtype CoreMap a = CoreMap (CoreMapG a)
148
149 instance TrieMap CoreMap where
150 type Key CoreMap = CoreExpr
151 emptyTM = CoreMap emptyTM
152 lookupTM k (CoreMap m) = lookupTM (deBruijnize k) m
153 alterTM k f (CoreMap m) = CoreMap (alterTM (deBruijnize k) f m)
154 foldTM k (CoreMap m) = foldTM k m
155 mapTM f (CoreMap m) = CoreMap (mapTM f m)
156
157 -- | @CoreMapG a@ is a map from @DeBruijn CoreExpr@ to @a@. The extended
158 -- key makes it suitable for recursive traversal, since it can track binders,
159 -- but it is strictly internal to this module. If you are including a 'CoreMap'
160 -- inside another 'TrieMap', this is the type you want.
161 type CoreMapG = GenMap CoreMapX
162
163 -- | @CoreMapX a@ is the base map from @DeBruijn CoreExpr@ to @a@, but without
164 -- the 'GenMap' optimization.
165 data CoreMapX a
166 = CM { cm_var :: VarMap a
167 , cm_lit :: LiteralMap a
168 , cm_co :: CoercionMapG a
169 , cm_type :: TypeMapG a
170 , cm_cast :: CoreMapG (CoercionMapG a)
171 , cm_tick :: CoreMapG (TickishMap a)
172 , cm_app :: CoreMapG (CoreMapG a)
173 , cm_lam :: CoreMapG (BndrMap a) -- Note [Binders]
174 , cm_letn :: CoreMapG (CoreMapG (BndrMap a))
175 , cm_letr :: ListMap CoreMapG (CoreMapG (ListMap BndrMap a))
176 , cm_case :: CoreMapG (ListMap AltMap a)
177 , cm_ecase :: CoreMapG (TypeMapG a) -- Note [Empty case alternatives]
178 }
179
180 instance Eq (DeBruijn CoreExpr) where
181 D env1 e1 == D env2 e2 = go e1 e2 where
182 go (Var v1) (Var v2) = case (lookupCME env1 v1, lookupCME env2 v2) of
183 (Just b1, Just b2) -> b1 == b2
184 (Nothing, Nothing) -> v1 == v2
185 _ -> False
186 go (Lit lit1) (Lit lit2) = lit1 == lit2
187 go (Type t1) (Type t2) = D env1 t1 == D env2 t2
188 go (Coercion co1) (Coercion co2) = D env1 co1 == D env2 co2
189 go (Cast e1 co1) (Cast e2 co2) = D env1 co1 == D env2 co2 && go e1 e2
190 go (App f1 a1) (App f2 a2) = go f1 f2 && go a1 a2
191 -- This seems a bit dodgy, see 'eqTickish'
192 go (Tick n1 e1) (Tick n2 e2) = n1 == n2 && go e1 e2
193
194 go (Lam b1 e1) (Lam b2 e2)
195 = D env1 (varType b1) == D env2 (varType b2)
196 && D (extendCME env1 b1) e1 == D (extendCME env2 b2) e2
197
198 go (Let (NonRec v1 r1) e1) (Let (NonRec v2 r2) e2)
199 = go r1 r2
200 && D (extendCME env1 v1) e1 == D (extendCME env2 v2) e2
201
202 go (Let (Rec ps1) e1) (Let (Rec ps2) e2)
203 = equalLength ps1 ps2
204 && D env1' rs1 == D env2' rs2
205 && D env1' e1 == D env2' e2
206 where
207 (bs1,rs1) = unzip ps1
208 (bs2,rs2) = unzip ps2
209 env1' = extendCMEs env1 bs1
210 env2' = extendCMEs env2 bs2
211
212 go (Case e1 b1 t1 a1) (Case e2 b2 t2 a2)
213 | null a1 -- See Note [Empty case alternatives]
214 = null a2 && go e1 e2 && D env1 t1 == D env2 t2
215 | otherwise
216 = go e1 e2 && D (extendCME env1 b1) a1 == D (extendCME env2 b2) a2
217
218 go _ _ = False
219
220 emptyE :: CoreMapX a
221 emptyE = CM { cm_var = emptyTM, cm_lit = emptyTM
222 , cm_co = emptyTM, cm_type = emptyTM
223 , cm_cast = emptyTM, cm_app = emptyTM
224 , cm_lam = emptyTM, cm_letn = emptyTM
225 , cm_letr = emptyTM, cm_case = emptyTM
226 , cm_ecase = emptyTM, cm_tick = emptyTM }
227
228 instance TrieMap CoreMapX where
229 type Key CoreMapX = DeBruijn CoreExpr
230 emptyTM = emptyE
231 lookupTM = lkE
232 alterTM = xtE
233 foldTM = fdE
234 mapTM = mapE
235
236 --------------------------
237 mapE :: (a->b) -> CoreMapX a -> CoreMapX b
238 mapE f (CM { cm_var = cvar, cm_lit = clit
239 , cm_co = cco, cm_type = ctype
240 , cm_cast = ccast , cm_app = capp
241 , cm_lam = clam, cm_letn = cletn
242 , cm_letr = cletr, cm_case = ccase
243 , cm_ecase = cecase, cm_tick = ctick })
244 = CM { cm_var = mapTM f cvar, cm_lit = mapTM f clit
245 , cm_co = mapTM f cco, cm_type = mapTM f ctype
246 , cm_cast = mapTM (mapTM f) ccast, cm_app = mapTM (mapTM f) capp
247 , cm_lam = mapTM (mapTM f) clam, cm_letn = mapTM (mapTM (mapTM f)) cletn
248 , cm_letr = mapTM (mapTM (mapTM f)) cletr, cm_case = mapTM (mapTM f) ccase
249 , cm_ecase = mapTM (mapTM f) cecase, cm_tick = mapTM (mapTM f) ctick }
250
251 --------------------------
252 lookupCoreMap :: CoreMap a -> CoreExpr -> Maybe a
253 lookupCoreMap cm e = lookupTM e cm
254
255 extendCoreMap :: CoreMap a -> CoreExpr -> a -> CoreMap a
256 extendCoreMap m e v = alterTM e (\_ -> Just v) m
257
258 foldCoreMap :: (a -> b -> b) -> b -> CoreMap a -> b
259 foldCoreMap k z m = foldTM k m z
260
261 emptyCoreMap :: CoreMap a
262 emptyCoreMap = emptyTM
263
264 instance Outputable a => Outputable (CoreMap a) where
265 ppr m = text "CoreMap elts" <+> ppr (foldTM (:) m [])
266
267 -------------------------
268 fdE :: (a -> b -> b) -> CoreMapX a -> b -> b
269 fdE k m
270 = foldTM k (cm_var m)
271 . foldTM k (cm_lit m)
272 . foldTM k (cm_co m)
273 . foldTM k (cm_type m)
274 . foldTM (foldTM k) (cm_cast m)
275 . foldTM (foldTM k) (cm_tick m)
276 . foldTM (foldTM k) (cm_app m)
277 . foldTM (foldTM k) (cm_lam m)
278 . foldTM (foldTM (foldTM k)) (cm_letn m)
279 . foldTM (foldTM (foldTM k)) (cm_letr m)
280 . foldTM (foldTM k) (cm_case m)
281 . foldTM (foldTM k) (cm_ecase m)
282
283 -- lkE: lookup in trie for expressions
284 lkE :: DeBruijn CoreExpr -> CoreMapX a -> Maybe a
285 lkE (D env expr) cm = go expr cm
286 where
287 go (Var v) = cm_var >.> lkVar env v
288 go (Lit l) = cm_lit >.> lookupTM l
289 go (Type t) = cm_type >.> lkG (D env t)
290 go (Coercion c) = cm_co >.> lkG (D env c)
291 go (Cast e c) = cm_cast >.> lkG (D env e) >=> lkG (D env c)
292 go (Tick tickish e) = cm_tick >.> lkG (D env e) >=> lkTickish tickish
293 go (App e1 e2) = cm_app >.> lkG (D env e2) >=> lkG (D env e1)
294 go (Lam v e) = cm_lam >.> lkG (D (extendCME env v) e)
295 >=> lkBndr env v
296 go (Let (NonRec b r) e) = cm_letn >.> lkG (D env r)
297 >=> lkG (D (extendCME env b) e) >=> lkBndr env b
298 go (Let (Rec prs) e) = let (bndrs,rhss) = unzip prs
299 env1 = extendCMEs env bndrs
300 in cm_letr
301 >.> lkList (lkG . D env1) rhss
302 >=> lkG (D env1 e)
303 >=> lkList (lkBndr env1) bndrs
304 go (Case e b ty as) -- See Note [Empty case alternatives]
305 | null as = cm_ecase >.> lkG (D env e) >=> lkG (D env ty)
306 | otherwise = cm_case >.> lkG (D env e)
307 >=> lkList (lkA (extendCME env b)) as
308
309 xtE :: DeBruijn CoreExpr -> XT a -> CoreMapX a -> CoreMapX a
310 xtE (D env (Var v)) f m = m { cm_var = cm_var m
311 |> xtVar env v f }
312 xtE (D env (Type t)) f m = m { cm_type = cm_type m
313 |> xtG (D env t) f }
314 xtE (D env (Coercion c)) f m = m { cm_co = cm_co m
315 |> xtG (D env c) f }
316 xtE (D _ (Lit l)) f m = m { cm_lit = cm_lit m |> alterTM l f }
317 xtE (D env (Cast e c)) f m = m { cm_cast = cm_cast m |> xtG (D env e)
318 |>> xtG (D env c) f }
319 xtE (D env (Tick t e)) f m = m { cm_tick = cm_tick m |> xtG (D env e)
320 |>> xtTickish t f }
321 xtE (D env (App e1 e2)) f m = m { cm_app = cm_app m |> xtG (D env e2)
322 |>> xtG (D env e1) f }
323 xtE (D env (Lam v e)) f m = m { cm_lam = cm_lam m
324 |> xtG (D (extendCME env v) e)
325 |>> xtBndr env v f }
326 xtE (D env (Let (NonRec b r) e)) f m = m { cm_letn = cm_letn m
327 |> xtG (D (extendCME env b) e)
328 |>> xtG (D env r)
329 |>> xtBndr env b f }
330 xtE (D env (Let (Rec prs) e)) f m = m { cm_letr =
331 let (bndrs,rhss) = unzip prs
332 env1 = extendCMEs env bndrs
333 in cm_letr m
334 |> xtList (xtG . D env1) rhss
335 |>> xtG (D env1 e)
336 |>> xtList (xtBndr env1)
337 bndrs f }
338 xtE (D env (Case e b ty as)) f m
339 | null as = m { cm_ecase = cm_ecase m |> xtG (D env e)
340 |>> xtG (D env ty) f }
341 | otherwise = m { cm_case = cm_case m |> xtG (D env e)
342 |>> let env1 = extendCME env b
343 in xtList (xtA env1) as f }
344
345 -- TODO: this seems a bit dodgy, see 'eqTickish'
346 type TickishMap a = Map.Map (Tickish Id) a
347 lkTickish :: Tickish Id -> TickishMap a -> Maybe a
348 lkTickish = lookupTM
349
350 xtTickish :: Tickish Id -> XT a -> TickishMap a -> TickishMap a
351 xtTickish = alterTM
352
353 ------------------------
354 data AltMap a -- A single alternative
355 = AM { am_deflt :: CoreMapG a
356 , am_data :: DNameEnv (CoreMapG a)
357 , am_lit :: LiteralMap (CoreMapG a) }
358
359 instance TrieMap AltMap where
360 type Key AltMap = CoreAlt
361 emptyTM = AM { am_deflt = emptyTM
362 , am_data = emptyDNameEnv
363 , am_lit = emptyTM }
364 lookupTM = lkA emptyCME
365 alterTM = xtA emptyCME
366 foldTM = fdA
367 mapTM = mapA
368
369 instance Eq (DeBruijn CoreAlt) where
370 D env1 a1 == D env2 a2 = go a1 a2 where
371 go (DEFAULT, _, rhs1) (DEFAULT, _, rhs2)
372 = D env1 rhs1 == D env2 rhs2
373 go (LitAlt lit1, _, rhs1) (LitAlt lit2, _, rhs2)
374 = lit1 == lit2 && D env1 rhs1 == D env2 rhs2
375 go (DataAlt dc1, bs1, rhs1) (DataAlt dc2, bs2, rhs2)
376 = dc1 == dc2 &&
377 D (extendCMEs env1 bs1) rhs1 == D (extendCMEs env2 bs2) rhs2
378 go _ _ = False
379
380 mapA :: (a->b) -> AltMap a -> AltMap b
381 mapA f (AM { am_deflt = adeflt, am_data = adata, am_lit = alit })
382 = AM { am_deflt = mapTM f adeflt
383 , am_data = mapTM (mapTM f) adata
384 , am_lit = mapTM (mapTM f) alit }
385
386 lkA :: CmEnv -> CoreAlt -> AltMap a -> Maybe a
387 lkA env (DEFAULT, _, rhs) = am_deflt >.> lkG (D env rhs)
388 lkA env (LitAlt lit, _, rhs) = am_lit >.> lookupTM lit >=> lkG (D env rhs)
389 lkA env (DataAlt dc, bs, rhs) = am_data >.> lkDNamed dc
390 >=> lkG (D (extendCMEs env bs) rhs)
391
392 xtA :: CmEnv -> CoreAlt -> XT a -> AltMap a -> AltMap a
393 xtA env (DEFAULT, _, rhs) f m =
394 m { am_deflt = am_deflt m |> xtG (D env rhs) f }
395 xtA env (LitAlt l, _, rhs) f m =
396 m { am_lit = am_lit m |> alterTM l |>> xtG (D env rhs) f }
397 xtA env (DataAlt d, bs, rhs) f m =
398 m { am_data = am_data m |> xtDNamed d
399 |>> xtG (D (extendCMEs env bs) rhs) f }
400
401 fdA :: (a -> b -> b) -> AltMap a -> b -> b
402 fdA k m = foldTM k (am_deflt m)
403 . foldTM (foldTM k) (am_data m)
404 . foldTM (foldTM k) (am_lit m)
405
406 {-
407 ************************************************************************
408 * *
409 Coercions
410 * *
411 ************************************************************************
412 -}
413
414 -- We should really never care about the contents of a coercion. Instead,
415 -- just look up the coercion's type.
416 newtype CoercionMap a = CoercionMap (CoercionMapG a)
417
418 instance TrieMap CoercionMap where
419 type Key CoercionMap = Coercion
420 emptyTM = CoercionMap emptyTM
421 lookupTM k (CoercionMap m) = lookupTM (deBruijnize k) m
422 alterTM k f (CoercionMap m) = CoercionMap (alterTM (deBruijnize k) f m)
423 foldTM k (CoercionMap m) = foldTM k m
424 mapTM f (CoercionMap m) = CoercionMap (mapTM f m)
425
426 type CoercionMapG = GenMap CoercionMapX
427 newtype CoercionMapX a = CoercionMapX (TypeMapX a)
428
429 instance TrieMap CoercionMapX where
430 type Key CoercionMapX = DeBruijn Coercion
431 emptyTM = CoercionMapX emptyTM
432 lookupTM = lkC
433 alterTM = xtC
434 foldTM f (CoercionMapX core_tm) = foldTM f core_tm
435 mapTM f (CoercionMapX core_tm) = CoercionMapX (mapTM f core_tm)
436
437 instance Eq (DeBruijn Coercion) where
438 D env1 co1 == D env2 co2
439 = D env1 (coercionType co1) ==
440 D env2 (coercionType co2)
441
442 lkC :: DeBruijn Coercion -> CoercionMapX a -> Maybe a
443 lkC (D env co) (CoercionMapX core_tm) = lkT (D env $ coercionType co)
444 core_tm
445
446 xtC :: DeBruijn Coercion -> XT a -> CoercionMapX a -> CoercionMapX a
447 xtC (D env co) f (CoercionMapX m)
448 = CoercionMapX (xtT (D env $ coercionType co) f m)
449
450 {-
451 ************************************************************************
452 * *
453 Types
454 * *
455 ************************************************************************
456 -}
457
458 -- | @TypeMapG a@ is a map from @DeBruijn Type@ to @a@. The extended
459 -- key makes it suitable for recursive traversal, since it can track binders,
460 -- but it is strictly internal to this module. If you are including a 'TypeMap'
461 -- inside another 'TrieMap', this is the type you want. Note that this
462 -- lookup does not do a kind-check. Thus, all keys in this map must have
463 -- the same kind. Also note that this map respects the distinction between
464 -- @Type@ and @Constraint@, despite the fact that they are equivalent type
465 -- synonyms in Core.
466 type TypeMapG = GenMap TypeMapX
467
468 -- | @TypeMapX a@ is the base map from @DeBruijn Type@ to @a@, but without the
469 -- 'GenMap' optimization.
470 data TypeMapX a
471 = TM { tm_var :: VarMap a
472 , tm_app :: TypeMapG (TypeMapG a)
473 , tm_tycon :: DNameEnv a
474 , tm_forall :: TypeMapG (BndrMap a) -- See Note [Binders]
475 , tm_tylit :: TyLitMap a
476 , tm_coerce :: Maybe a
477 }
478 -- Note that there is no tyconapp case; see Note [Equality on AppTys] in Type
479
480 -- | Squeeze out any synonyms, and change TyConApps to nested AppTys. Why the
481 -- last one? See Note [Equality on AppTys] in Type
482 --
483 -- Note, however, that we keep Constraint and Type apart here, despite the fact
484 -- that they are both synonyms of TYPE 'LiftedRep (see #11715).
485 trieMapView :: Type -> Maybe Type
486 trieMapView ty
487 -- First check for TyConApps that need to be expanded to
488 -- AppTy chains.
489 | Just (tc, tys@(_:_)) <- tcSplitTyConApp_maybe ty
490 = Just $ foldl' AppTy (TyConApp tc []) tys
491
492 -- Then resolve any remaining nullary synonyms.
493 | Just ty' <- tcView ty = Just ty'
494 trieMapView _ = Nothing
495
496 instance TrieMap TypeMapX where
497 type Key TypeMapX = DeBruijn Type
498 emptyTM = emptyT
499 lookupTM = lkT
500 alterTM = xtT
501 foldTM = fdT
502 mapTM = mapT
503
504 instance Eq (DeBruijn Type) where
505 env_t@(D env t) == env_t'@(D env' t')
506 | Just new_t <- tcView t = D env new_t == env_t'
507 | Just new_t' <- tcView t' = env_t == D env' new_t'
508 | otherwise
509 = case (t, t') of
510 (CastTy t1 _, _) -> D env t1 == D env t'
511 (_, CastTy t1' _) -> D env t == D env t1'
512
513 (TyVarTy v, TyVarTy v')
514 -> case (lookupCME env v, lookupCME env' v') of
515 (Just bv, Just bv') -> bv == bv'
516 (Nothing, Nothing) -> v == v'
517 _ -> False
518 -- See Note [Equality on AppTys] in Type
519 (AppTy t1 t2, s) | Just (t1', t2') <- repSplitAppTy_maybe s
520 -> D env t1 == D env' t1' && D env t2 == D env' t2'
521 (s, AppTy t1' t2') | Just (t1, t2) <- repSplitAppTy_maybe s
522 -> D env t1 == D env' t1' && D env t2 == D env' t2'
523 (FunTy _ t1 t2, FunTy _ t1' t2')
524 -> D env t1 == D env' t1' && D env t2 == D env' t2'
525 (TyConApp tc tys, TyConApp tc' tys')
526 -> tc == tc' && D env tys == D env' tys'
527 (LitTy l, LitTy l')
528 -> l == l'
529 (ForAllTy (Bndr tv _) ty, ForAllTy (Bndr tv' _) ty')
530 -> D env (varType tv) == D env' (varType tv') &&
531 D (extendCME env tv) ty == D (extendCME env' tv') ty'
532 (CoercionTy {}, CoercionTy {})
533 -> True
534 _ -> False
535
536 instance {-# OVERLAPPING #-}
537 Outputable a => Outputable (TypeMapG a) where
538 ppr m = text "TypeMap elts" <+> ppr (foldTM (:) m [])
539
540 emptyT :: TypeMapX a
541 emptyT = TM { tm_var = emptyTM
542 , tm_app = emptyTM
543 , tm_tycon = emptyDNameEnv
544 , tm_forall = emptyTM
545 , tm_tylit = emptyTyLitMap
546 , tm_coerce = Nothing }
547
548 mapT :: (a->b) -> TypeMapX a -> TypeMapX b
549 mapT f (TM { tm_var = tvar, tm_app = tapp, tm_tycon = ttycon
550 , tm_forall = tforall, tm_tylit = tlit
551 , tm_coerce = tcoerce })
552 = TM { tm_var = mapTM f tvar
553 , tm_app = mapTM (mapTM f) tapp
554 , tm_tycon = mapTM f ttycon
555 , tm_forall = mapTM (mapTM f) tforall
556 , tm_tylit = mapTM f tlit
557 , tm_coerce = fmap f tcoerce }
558
559 -----------------
560 lkT :: DeBruijn Type -> TypeMapX a -> Maybe a
561 lkT (D env ty) m = go ty m
562 where
563 go ty | Just ty' <- trieMapView ty = go ty'
564 go (TyVarTy v) = tm_var >.> lkVar env v
565 go (AppTy t1 t2) = tm_app >.> lkG (D env t1)
566 >=> lkG (D env t2)
567 go (TyConApp tc []) = tm_tycon >.> lkDNamed tc
568 go ty@(TyConApp _ (_:_)) = pprPanic "lkT TyConApp" (ppr ty)
569 go (LitTy l) = tm_tylit >.> lkTyLit l
570 go (ForAllTy (Bndr tv _) ty) = tm_forall >.> lkG (D (extendCME env tv) ty)
571 >=> lkBndr env tv
572 go ty@(FunTy {}) = pprPanic "lkT FunTy" (ppr ty)
573 go (CastTy t _) = go t
574 go (CoercionTy {}) = tm_coerce
575
576 -----------------
577 xtT :: DeBruijn Type -> XT a -> TypeMapX a -> TypeMapX a
578 xtT (D env ty) f m | Just ty' <- trieMapView ty = xtT (D env ty') f m
579
580 xtT (D env (TyVarTy v)) f m = m { tm_var = tm_var m |> xtVar env v f }
581 xtT (D env (AppTy t1 t2)) f m = m { tm_app = tm_app m |> xtG (D env t1)
582 |>> xtG (D env t2) f }
583 xtT (D _ (TyConApp tc [])) f m = m { tm_tycon = tm_tycon m |> xtDNamed tc f }
584 xtT (D _ (LitTy l)) f m = m { tm_tylit = tm_tylit m |> xtTyLit l f }
585 xtT (D env (CastTy t _)) f m = xtT (D env t) f m
586 xtT (D _ (CoercionTy {})) f m = m { tm_coerce = tm_coerce m |> f }
587 xtT (D env (ForAllTy (Bndr tv _) ty)) f m
588 = m { tm_forall = tm_forall m |> xtG (D (extendCME env tv) ty)
589 |>> xtBndr env tv f }
590 xtT (D _ ty@(TyConApp _ (_:_))) _ _ = pprPanic "xtT TyConApp" (ppr ty)
591 xtT (D _ ty@(FunTy {})) _ _ = pprPanic "xtT FunTy" (ppr ty)
592
593 fdT :: (a -> b -> b) -> TypeMapX a -> b -> b
594 fdT k m = foldTM k (tm_var m)
595 . foldTM (foldTM k) (tm_app m)
596 . foldTM k (tm_tycon m)
597 . foldTM (foldTM k) (tm_forall m)
598 . foldTyLit k (tm_tylit m)
599 . foldMaybe k (tm_coerce m)
600
601 ------------------------
602 data TyLitMap a = TLM { tlm_number :: Map.Map Integer a
603 , tlm_string :: Map.Map FastString a
604 }
605
606 instance TrieMap TyLitMap where
607 type Key TyLitMap = TyLit
608 emptyTM = emptyTyLitMap
609 lookupTM = lkTyLit
610 alterTM = xtTyLit
611 foldTM = foldTyLit
612 mapTM = mapTyLit
613
614 emptyTyLitMap :: TyLitMap a
615 emptyTyLitMap = TLM { tlm_number = Map.empty, tlm_string = Map.empty }
616
617 mapTyLit :: (a->b) -> TyLitMap a -> TyLitMap b
618 mapTyLit f (TLM { tlm_number = tn, tlm_string = ts })
619 = TLM { tlm_number = Map.map f tn, tlm_string = Map.map f ts }
620
621 lkTyLit :: TyLit -> TyLitMap a -> Maybe a
622 lkTyLit l =
623 case l of
624 NumTyLit n -> tlm_number >.> Map.lookup n
625 StrTyLit n -> tlm_string >.> Map.lookup n
626
627 xtTyLit :: TyLit -> XT a -> TyLitMap a -> TyLitMap a
628 xtTyLit l f m =
629 case l of
630 NumTyLit n -> m { tlm_number = tlm_number m |> Map.alter f n }
631 StrTyLit n -> m { tlm_string = tlm_string m |> Map.alter f n }
632
633 foldTyLit :: (a -> b -> b) -> TyLitMap a -> b -> b
634 foldTyLit l m = flip (Map.foldr l) (tlm_string m)
635 . flip (Map.foldr l) (tlm_number m)
636
637 -------------------------------------------------
638 -- | @TypeMap a@ is a map from 'Type' to @a@. If you are a client, this
639 -- is the type you want. The keys in this map may have different kinds.
640 newtype TypeMap a = TypeMap (TypeMapG (TypeMapG a))
641
642 lkTT :: DeBruijn Type -> TypeMap a -> Maybe a
643 lkTT (D env ty) (TypeMap m) = lkG (D env $ typeKind ty) m
644 >>= lkG (D env ty)
645
646 xtTT :: DeBruijn Type -> XT a -> TypeMap a -> TypeMap a
647 xtTT (D env ty) f (TypeMap m)
648 = TypeMap (m |> xtG (D env $ typeKind ty)
649 |>> xtG (D env ty) f)
650
651 -- Below are some client-oriented functions which operate on 'TypeMap'.
652
653 instance TrieMap TypeMap where
654 type Key TypeMap = Type
655 emptyTM = TypeMap emptyTM
656 lookupTM k m = lkTT (deBruijnize k) m
657 alterTM k f m = xtTT (deBruijnize k) f m
658 foldTM k (TypeMap m) = foldTM (foldTM k) m
659 mapTM f (TypeMap m) = TypeMap (mapTM (mapTM f) m)
660
661 foldTypeMap :: (a -> b -> b) -> b -> TypeMap a -> b
662 foldTypeMap k z m = foldTM k m z
663
664 emptyTypeMap :: TypeMap a
665 emptyTypeMap = emptyTM
666
667 lookupTypeMap :: TypeMap a -> Type -> Maybe a
668 lookupTypeMap cm t = lookupTM t cm
669
670 extendTypeMap :: TypeMap a -> Type -> a -> TypeMap a
671 extendTypeMap m t v = alterTM t (const (Just v)) m
672
673 lookupTypeMapWithScope :: TypeMap a -> CmEnv -> Type -> Maybe a
674 lookupTypeMapWithScope m cm t = lkTT (D cm t) m
675
676 -- | Extend a 'TypeMap' with a type in the given context.
677 -- @extendTypeMapWithScope m (mkDeBruijnContext [a,b,c]) t v@ is equivalent to
678 -- @extendTypeMap m (forall a b c. t) v@, but allows reuse of the context over
679 -- multiple insertions.
680 extendTypeMapWithScope :: TypeMap a -> CmEnv -> Type -> a -> TypeMap a
681 extendTypeMapWithScope m cm t v = xtTT (D cm t) (const (Just v)) m
682
683 -- | Construct a deBruijn environment with the given variables in scope.
684 -- e.g. @mkDeBruijnEnv [a,b,c]@ constructs a context @forall a b c.@
685 mkDeBruijnContext :: [Var] -> CmEnv
686 mkDeBruijnContext = extendCMEs emptyCME
687
688 -- | A 'LooseTypeMap' doesn't do a kind-check. Thus, when lookup up (t |> g),
689 -- you'll find entries inserted under (t), even if (g) is non-reflexive.
690 newtype LooseTypeMap a
691 = LooseTypeMap (TypeMapG a)
692
693 instance TrieMap LooseTypeMap where
694 type Key LooseTypeMap = Type
695 emptyTM = LooseTypeMap emptyTM
696 lookupTM k (LooseTypeMap m) = lookupTM (deBruijnize k) m
697 alterTM k f (LooseTypeMap m) = LooseTypeMap (alterTM (deBruijnize k) f m)
698 foldTM f (LooseTypeMap m) = foldTM f m
699 mapTM f (LooseTypeMap m) = LooseTypeMap (mapTM f m)
700
701 {-
702 ************************************************************************
703 * *
704 Variables
705 * *
706 ************************************************************************
707 -}
708
709 type BoundVar = Int -- Bound variables are deBruijn numbered
710 type BoundVarMap a = IntMap.IntMap a
711
712 data CmEnv = CME { cme_next :: !BoundVar
713 , cme_env :: VarEnv BoundVar }
714
715 emptyCME :: CmEnv
716 emptyCME = CME { cme_next = 0, cme_env = emptyVarEnv }
717
718 extendCME :: CmEnv -> Var -> CmEnv
719 extendCME (CME { cme_next = bv, cme_env = env }) v
720 = CME { cme_next = bv+1, cme_env = extendVarEnv env v bv }
721
722 extendCMEs :: CmEnv -> [Var] -> CmEnv
723 extendCMEs env vs = foldl' extendCME env vs
724
725 lookupCME :: CmEnv -> Var -> Maybe BoundVar
726 lookupCME (CME { cme_env = env }) v = lookupVarEnv env v
727
728 -- | @DeBruijn a@ represents @a@ modulo alpha-renaming. This is achieved
729 -- by equipping the value with a 'CmEnv', which tracks an on-the-fly deBruijn
730 -- numbering. This allows us to define an 'Eq' instance for @DeBruijn a@, even
731 -- if this was not (easily) possible for @a@. Note: we purposely don't
732 -- export the constructor. Make a helper function if you find yourself
733 -- needing it.
734 data DeBruijn a = D CmEnv a
735
736 -- | Synthesizes a @DeBruijn a@ from an @a@, by assuming that there are no
737 -- bound binders (an empty 'CmEnv'). This is usually what you want if there
738 -- isn't already a 'CmEnv' in scope.
739 deBruijnize :: a -> DeBruijn a
740 deBruijnize = D emptyCME
741
742 instance Eq (DeBruijn a) => Eq (DeBruijn [a]) where
743 D _ [] == D _ [] = True
744 D env (x:xs) == D env' (x':xs') = D env x == D env' x' &&
745 D env xs == D env' xs'
746 _ == _ = False
747
748 --------- Variable binders -------------
749
750 -- | A 'BndrMap' is a 'TypeMapG' which allows us to distinguish between
751 -- binding forms whose binders have different types. For example,
752 -- if we are doing a 'TrieMap' lookup on @\(x :: Int) -> ()@, we should
753 -- not pick up an entry in the 'TrieMap' for @\(x :: Bool) -> ()@:
754 -- we can disambiguate this by matching on the type (or kind, if this
755 -- a binder in a type) of the binder.
756 type BndrMap = TypeMapG
757
758 -- Note [Binders]
759 -- ~~~~~~~~~~~~~~
760 -- We need to use 'BndrMap' for 'Coercion', 'CoreExpr' AND 'Type', since all
761 -- of these data types have binding forms.
762
763 lkBndr :: CmEnv -> Var -> BndrMap a -> Maybe a
764 lkBndr env v m = lkG (D env (varType v)) m
765
766 xtBndr :: CmEnv -> Var -> XT a -> BndrMap a -> BndrMap a
767 xtBndr env v f = xtG (D env (varType v)) f
768
769 --------- Variable occurrence -------------
770 data VarMap a = VM { vm_bvar :: BoundVarMap a -- Bound variable
771 , vm_fvar :: DVarEnv a } -- Free variable
772
773 instance TrieMap VarMap where
774 type Key VarMap = Var
775 emptyTM = VM { vm_bvar = IntMap.empty, vm_fvar = emptyDVarEnv }
776 lookupTM = lkVar emptyCME
777 alterTM = xtVar emptyCME
778 foldTM = fdVar
779 mapTM = mapVar
780
781 mapVar :: (a->b) -> VarMap a -> VarMap b
782 mapVar f (VM { vm_bvar = bv, vm_fvar = fv })
783 = VM { vm_bvar = mapTM f bv, vm_fvar = mapTM f fv }
784
785 lkVar :: CmEnv -> Var -> VarMap a -> Maybe a
786 lkVar env v
787 | Just bv <- lookupCME env v = vm_bvar >.> lookupTM bv
788 | otherwise = vm_fvar >.> lkDFreeVar v
789
790 xtVar :: CmEnv -> Var -> XT a -> VarMap a -> VarMap a
791 xtVar env v f m
792 | Just bv <- lookupCME env v = m { vm_bvar = vm_bvar m |> alterTM bv f }
793 | otherwise = m { vm_fvar = vm_fvar m |> xtDFreeVar v f }
794
795 fdVar :: (a -> b -> b) -> VarMap a -> b -> b
796 fdVar k m = foldTM k (vm_bvar m)
797 . foldTM k (vm_fvar m)
798
799 lkDFreeVar :: Var -> DVarEnv a -> Maybe a
800 lkDFreeVar var env = lookupDVarEnv env var
801
802 xtDFreeVar :: Var -> XT a -> DVarEnv a -> DVarEnv a
803 xtDFreeVar v f m = alterDVarEnv f m v