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