PageRenderTime 67ms CodeModel.GetById 56ms app.highlight 4ms RepoModel.GetById 1ms app.codeStats 0ms

/compiler/coreSyn/TrieMap.lhs

https://bitbucket.org/carter/ghc
Haskell | 803 lines | 646 code | 127 blank | 30 comment | 1 complexity | 97b8951896454e09060f2cc5fe010c16 MD5 | raw file
  1%
  2% (c) The University of Glasgow 2006
  3% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
  4%
  5
  6\begin{code}
  7{-# OPTIONS -fno-warn-tabs #-}
  8-- The above warning supression flag is a temporary kludge.
  9-- While working on this module you are encouraged to remove it and
 10-- detab the module (please do the detabbing in a separate patch). See
 11--     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces
 12-- for details
 13
 14{-# LANGUAGE TypeFamilies #-}
 15module TrieMap(
 16   CoreMap, emptyCoreMap, extendCoreMap, lookupCoreMap, foldCoreMap,
 17   TypeMap, foldTypeMap, lookupTypeMap_mod,
 18   CoercionMap, 
 19   MaybeMap, 
 20   ListMap,
 21   TrieMap(..)
 22 ) where
 23
 24import CoreSyn
 25import Coercion
 26import Literal
 27import Name
 28import Type
 29import TypeRep
 30import Var
 31import UniqFM
 32import Unique( Unique )
 33import FastString(FastString)
 34
 35import Unify ( niFixTvSubst )
 36
 37import qualified Data.Map    as Map
 38import qualified Data.IntMap as IntMap
 39import VarEnv
 40import NameEnv
 41import Outputable
 42import Control.Monad( (>=>) )
 43\end{code}
 44
 45This module implements TrieMaps, which are finite mappings
 46whose key is a structured value like a CoreExpr or Type.
 47
 48The code is very regular and boilerplate-like, but there is
 49some neat handling of *binders*.  In effect they are deBruijn 
 50numbered on the fly.
 51
 52%************************************************************************
 53%*									*
 54                   The TrieMap class
 55%*									*
 56%************************************************************************
 57
 58\begin{code}
 59type XT a = Maybe a -> Maybe a	-- How to alter a non-existent elt (Nothing)
 60     	    	       		--               or an existing elt (Just)
 61
 62class TrieMap m where
 63   type Key m :: *
 64   emptyTM  :: m a
 65   lookupTM :: forall b. Key m -> m b -> Maybe b
 66   alterTM  :: forall b. Key m -> XT b -> m b -> m b
 67   mapTM    :: (a->b) -> m a -> m b
 68
 69   foldTM   :: (a -> b -> b) -> m a -> b -> b
 70      -- The unusual argument order here makes 
 71      -- it easy to compose calls to foldTM; 
 72      -- see for example fdE below
 73
 74----------------------
 75-- Recall that 
 76--   Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c
 77
 78(>.>) :: (a -> b) -> (b -> c) -> a -> c
 79-- Reverse function composition (do f first, then g)
 80infixr 1 >.>
 81(f >.> g) x = g (f x)
 82infixr 1 |>, |>>
 83
 84(|>) :: a -> (a->b) -> b     -- Reverse application
 85x |> f = f x
 86
 87----------------------
 88(|>>) :: TrieMap m2 
 89      => (XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
 90      -> (m2 a -> m2 a)
 91      -> m1 (m2 a) -> m1 (m2 a)
 92(|>>) f g = f (Just . g . deMaybe)
 93
 94deMaybe :: TrieMap m => Maybe (m a) -> m a
 95deMaybe Nothing  = emptyTM
 96deMaybe (Just m) = m
 97\end{code}
 98
 99%************************************************************************
100%*									*
101                   IntMaps
102%*									*
103%************************************************************************
104
105\begin{code}
106instance TrieMap IntMap.IntMap where
107  type Key IntMap.IntMap = Int
108  emptyTM = IntMap.empty
109  lookupTM k m = IntMap.lookup k m
110  alterTM = xtInt
111  foldTM k m z = IntMap.fold k z m
112  mapTM f m = IntMap.map f m
113
114xtInt :: Int -> XT a -> IntMap.IntMap a -> IntMap.IntMap a
115xtInt k f m = IntMap.alter f k m
116
117instance Ord k => TrieMap (Map.Map k) where
118  type Key (Map.Map k) = k
119  emptyTM = Map.empty
120  lookupTM = Map.lookup
121  alterTM k f m = Map.alter f k m
122  foldTM k m z = Map.fold k z m
123  mapTM f m = Map.map f m
124
125instance TrieMap UniqFM where
126  type Key UniqFM = Unique
127  emptyTM = emptyUFM
128  lookupTM k m = lookupUFM m k
129  alterTM k f m = alterUFM f m k
130  foldTM k m z = foldUFM k z m
131  mapTM f m = mapUFM f m
132\end{code}
133
134
135%************************************************************************
136%*									*
137                   Lists
138%*									*
139%************************************************************************
140
141If              m is a map from k -> val
142then (MaybeMap m) is a map from (Maybe k) -> val
143
144\begin{code}
145data MaybeMap m a = MM { mm_nothing  :: Maybe a, mm_just :: m a }
146
147instance TrieMap m => TrieMap (MaybeMap m) where
148   type Key (MaybeMap m) = Maybe (Key m)
149   emptyTM  = MM { mm_nothing = Nothing, mm_just = emptyTM }
150   lookupTM = lkMaybe lookupTM
151   alterTM  = xtMaybe alterTM
152   foldTM   = fdMaybe 
153   mapTM    = mapMb
154
155mapMb :: TrieMap m => (a->b) -> MaybeMap m a -> MaybeMap m b
156mapMb f (MM { mm_nothing = mn, mm_just = mj }) 
157  = MM { mm_nothing = fmap f mn, mm_just = mapTM f mj }
158
159lkMaybe :: TrieMap m => (forall b. k -> m b -> Maybe b)
160        -> Maybe k -> MaybeMap m a -> Maybe a
161lkMaybe _  Nothing  = mm_nothing
162lkMaybe lk (Just x) = mm_just >.> lk x
163
164xtMaybe :: TrieMap m => (forall b. k -> XT b -> m b -> m b)
165        -> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a
166xtMaybe _  Nothing  f m = m { mm_nothing  = f (mm_nothing m) }
167xtMaybe tr (Just x) f m = m { mm_just = mm_just m |> tr x f }
168
169fdMaybe :: TrieMap m => (a -> b -> b) -> MaybeMap m a -> b -> b
170fdMaybe k m = foldMaybe k (mm_nothing m)
171            . foldTM k (mm_just m)
172
173--------------------
174data ListMap m a
175  = LM { lm_nil  :: Maybe a
176       , lm_cons :: m (ListMap m a) }
177
178instance TrieMap m => TrieMap (ListMap m) where
179   type Key (ListMap m) = [Key m]
180   emptyTM  = LM { lm_nil = Nothing, lm_cons = emptyTM }
181   lookupTM = lkList lookupTM
182   alterTM  = xtList alterTM
183   foldTM   = fdList 
184   mapTM    = mapList
185
186mapList :: TrieMap m => (a->b) -> ListMap m a -> ListMap m b
187mapList f (LM { lm_nil = mnil, lm_cons = mcons })
188  = LM { lm_nil = fmap f mnil, lm_cons = mapTM (mapTM f) mcons }
189
190lkList :: TrieMap m => (forall b. k -> m b -> Maybe b)
191        -> [k] -> ListMap m a -> Maybe a
192lkList _  []     = lm_nil
193lkList lk (x:xs) = lm_cons >.> lk x >=> lkList lk xs
194
195xtList :: TrieMap m => (forall b. k -> XT b -> m b -> m b)
196        -> [k] -> XT a -> ListMap m a -> ListMap m a
197xtList _  []     f m = m { lm_nil  = f (lm_nil m) }
198xtList tr (x:xs) f m = m { lm_cons = lm_cons m |> tr x |>> xtList tr xs f }
199
200fdList :: forall m a b. TrieMap m 
201       => (a -> b -> b) -> ListMap m a -> b -> b
202fdList k m = foldMaybe k          (lm_nil m)
203           . foldTM    (fdList k) (lm_cons m)
204
205foldMaybe :: (a -> b -> b) -> Maybe a -> b -> b
206foldMaybe _ Nothing  b = b
207foldMaybe k (Just a) b = k a b
208\end{code}
209
210
211%************************************************************************
212%*									*
213                   Basic maps
214%*									*
215%************************************************************************
216
217\begin{code}
218lkNamed :: NamedThing n => n -> NameEnv a -> Maybe a
219lkNamed n env = lookupNameEnv env (getName n)
220
221xtNamed :: NamedThing n => n -> XT a -> NameEnv a -> NameEnv a
222xtNamed tc f m = alterNameEnv f m (getName tc)
223
224------------------------
225type LiteralMap  a = Map.Map Literal a
226
227emptyLiteralMap :: LiteralMap a
228emptyLiteralMap = emptyTM
229
230lkLit :: Literal -> LiteralMap a -> Maybe a
231lkLit = lookupTM
232
233xtLit :: Literal -> XT a -> LiteralMap a -> LiteralMap a
234xtLit = alterTM
235\end{code}
236
237%************************************************************************
238%*									*
239                   CoreMap
240%*									*
241%************************************************************************
242
243Note [Binders]
244~~~~~~~~~~~~~~
245 * In general we check binders as late as possible because types are
246   less likely to differ than expression structure.  That's why
247      cm_lam :: CoreMap (TypeMap a)
248   rather than
249      cm_lam :: TypeMap (CoreMap a)
250
251 * We don't need to look at the type of some binders, notalby
252     - the case binder in (Case _ b _ _)
253     - the binders in an alternative
254   because they are totally fixed by the context
255
256Note [Empty case alternatives]
257~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
258* For a key (Case e b ty (alt:alts))  we don't need to look the return type
259  'ty', because every alternative has that type.
260
261* For a key (Case e b ty []) we MUST look at the return type 'ty', because
262  otherwise (Case (error () "urk") _ Int  []) would compare equal to 
263            (Case (error () "urk") _ Bool [])
264  which is utterly wrong (Trac #6097)
265
266We could compare the return type regardless, but the wildly common case
267is that it's unnecesary, so we have two fields (cm_case and cm_ecase)
268for the two possibilities.  Only cm_ecase looks at the type.
269
270See also Note [Empty case alternatives] in CoreSyn.
271
272\begin{code}
273data CoreMap a
274  = EmptyCM
275  | CM { cm_var   :: VarMap a
276       , cm_lit   :: LiteralMap a
277       , cm_co    :: CoercionMap a
278       , cm_type  :: TypeMap a
279       , cm_cast  :: CoreMap (CoercionMap a)
280       , cm_tick  :: CoreMap (TickishMap a)
281       , cm_app   :: CoreMap (CoreMap a)
282       , cm_lam   :: CoreMap (TypeMap a)    -- Note [Binders]
283       , cm_letn  :: CoreMap (CoreMap (BndrMap a))
284       , cm_letr  :: ListMap CoreMap (CoreMap (ListMap BndrMap a))
285       , cm_case  :: CoreMap (ListMap AltMap a)
286       , cm_ecase :: CoreMap (TypeMap a)    -- Note [Empty case alternatives]
287     }
288
289
290wrapEmptyCM :: CoreMap a
291wrapEmptyCM = CM { cm_var = emptyTM, cm_lit = emptyLiteralMap
292 		 , cm_co = emptyTM, cm_type = emptyTM
293 		 , cm_cast = emptyTM, cm_app = emptyTM 
294 		 , cm_lam = emptyTM, cm_letn = emptyTM 
295 		 , cm_letr = emptyTM, cm_case = emptyTM
296                 , cm_ecase = emptyTM, cm_tick = emptyTM }
297
298instance TrieMap CoreMap where
299   type Key CoreMap = CoreExpr
300   emptyTM  = EmptyCM
301   lookupTM = lkE emptyCME
302   alterTM  = xtE emptyCME
303   foldTM   = fdE
304   mapTM    = mapE
305
306--------------------------
307mapE :: (a->b) -> CoreMap a -> CoreMap b
308mapE _ EmptyCM = EmptyCM
309mapE f (CM { cm_var = cvar, cm_lit = clit
310           , cm_co = cco, cm_type = ctype
311 	   , cm_cast = ccast , cm_app = capp
312 	   , cm_lam = clam, cm_letn = cletn 
313 	   , cm_letr = cletr, cm_case = ccase
314           , cm_ecase = cecase, cm_tick = ctick })
315  = CM { cm_var = mapTM f cvar, cm_lit = mapTM f clit 
316       , cm_co = mapTM f cco, cm_type = mapTM f ctype
317       , cm_cast = mapTM (mapTM f) ccast, cm_app = mapTM (mapTM f) capp
318       , cm_lam = mapTM (mapTM f) clam, cm_letn = mapTM (mapTM (mapTM f)) cletn 
319       , cm_letr = mapTM (mapTM (mapTM f)) cletr, cm_case = mapTM (mapTM f) ccase
320       , cm_ecase = mapTM (mapTM f) cecase, cm_tick = mapTM (mapTM f) ctick }
321
322--------------------------
323lookupCoreMap :: CoreMap a -> CoreExpr -> Maybe a
324lookupCoreMap cm e = lkE emptyCME e cm
325
326extendCoreMap :: CoreMap a -> CoreExpr -> a -> CoreMap a
327extendCoreMap m e v = xtE emptyCME e (\_ -> Just v) m
328
329foldCoreMap :: (a -> b -> b) -> b -> CoreMap a -> b
330foldCoreMap k z m = fdE k m z
331
332emptyCoreMap :: CoreMap a
333emptyCoreMap = EmptyCM
334
335instance Outputable a => Outputable (CoreMap a) where
336  ppr m = text "CoreMap elts" <+> ppr (foldCoreMap (:) [] m)
337
338-------------------------
339fdE :: (a -> b -> b) -> CoreMap a -> b -> b
340fdE _ EmptyCM = \z -> z
341fdE k m 
342  = foldTM k (cm_var m) 
343  . foldTM k (cm_lit m)
344  . foldTM k (cm_co m)
345  . foldTM k (cm_type m)
346  . foldTM (foldTM k) (cm_cast m)
347  . foldTM (foldTM k) (cm_tick m)
348  . foldTM (foldTM k) (cm_app m)
349  . foldTM (foldTM k) (cm_lam m)
350  . foldTM (foldTM (foldTM k)) (cm_letn m)
351  . foldTM (foldTM (foldTM k)) (cm_letr m)
352  . foldTM (foldTM k) (cm_case m)
353  . foldTM (foldTM k) (cm_ecase m)
354
355lkE :: CmEnv -> CoreExpr -> CoreMap a -> Maybe a
356-- lkE: lookup in trie for expressions
357lkE env expr cm
358  | EmptyCM <- cm = Nothing
359  | otherwise     = go expr cm
360  where 
361    go (Var v)  	    = cm_var  >.> lkVar env v
362    go (Lit l)              = cm_lit  >.> lkLit l
363    go (Type t) 	    = cm_type >.> lkT env t
364    go (Coercion c)         = cm_co   >.> lkC env c
365    go (Cast e c)           = cm_cast >.> lkE env e >=> lkC env c
366    go (Tick tickish e)     = cm_tick >.> lkE env e >=> lkTickish tickish
367    go (App e1 e2)          = cm_app  >.> lkE env e2 >=> lkE env e1
368    go (Lam v e)            = cm_lam  >.> lkE (extendCME env v) e >=> lkBndr env v
369    go (Let (NonRec b r) e) = cm_letn >.> lkE env r 
370                              >=> lkE (extendCME env b) e >=> lkBndr env b
371    go (Let (Rec prs) e)    = let (bndrs,rhss) = unzip prs
372                                  env1 = extendCMEs env bndrs
373                              in cm_letr
374                                 >.> lkList (lkE env1) rhss >=> lkE env1 e
375                                 >=> lkList (lkBndr env1) bndrs
376    go (Case e b ty as)     -- See Note [Empty case alternatives]
377               | null as    = cm_ecase >.> lkE env e >=> lkT env ty
378               | otherwise  = cm_case >.> lkE env e 
379                              >=> lkList (lkA (extendCME env b)) as
380
381xtE :: CmEnv -> CoreExpr -> XT a -> CoreMap a -> CoreMap a
382xtE env e              f EmptyCM = xtE env e f wrapEmptyCM
383xtE env (Var v)              f m = m { cm_var  = cm_var m  |> xtVar env v f }
384xtE env (Type t) 	     f m = m { cm_type = cm_type m |> xtT env t f }
385xtE env (Coercion c)         f m = m { cm_co   = cm_co m   |> xtC env c f }
386xtE _   (Lit l)              f m = m { cm_lit  = cm_lit m  |> xtLit l f }
387xtE env (Cast e c)           f m = m { cm_cast = cm_cast m |> xtE env e |>>
388                                                 xtC env c f }
389xtE env (Tick t e)           f m = m { cm_tick = cm_tick m |> xtE env e |>> xtTickish t f }
390xtE env (App e1 e2)          f m = m { cm_app = cm_app m |> xtE env e2 |>> xtE env e1 f }
391xtE env (Lam v e)            f m = m { cm_lam = cm_lam m |> xtE (extendCME env v) e
392                                                 |>> xtBndr env v f }
393xtE env (Let (NonRec b r) e) f m = m { cm_letn = cm_letn m 
394                                                 |> xtE (extendCME env b) e 
395                                                 |>> xtE env r |>> xtBndr env b f }
396xtE env (Let (Rec prs) e)    f m = m { cm_letr = let (bndrs,rhss) = unzip prs
397                                                     env1 = extendCMEs env bndrs
398                                                 in cm_letr m 
399                                                    |>  xtList (xtE env1) rhss 
400                                                    |>> xtE env1 e 
401                                                    |>> xtList (xtBndr env1) bndrs f }
402xtE env (Case e b ty as)     f m 
403                     | null as   = m { cm_ecase = cm_ecase m |> xtE env e |>> xtT env ty f }
404                     | otherwise = m { cm_case = cm_case m |> xtE env e 
405                                                 |>> let env1 = extendCME env b
406                                                     in xtList (xtA env1) as f }
407
408type TickishMap a = Map.Map (Tickish Id) a
409lkTickish :: Tickish Id -> TickishMap a -> Maybe a
410lkTickish = lookupTM
411
412xtTickish :: Tickish Id -> XT a -> TickishMap a -> TickishMap a
413xtTickish = alterTM
414
415------------------------
416data AltMap a	-- A single alternative
417  = AM { am_deflt :: CoreMap a
418       , am_data  :: NameEnv (CoreMap a)
419       , am_lit   :: LiteralMap (CoreMap a) }
420
421instance TrieMap AltMap where
422   type Key AltMap = CoreAlt
423   emptyTM  = AM { am_deflt = emptyTM
424                 , am_data = emptyNameEnv
425                 , am_lit  = emptyLiteralMap }
426   lookupTM = lkA emptyCME
427   alterTM  = xtA emptyCME
428   foldTM   = fdA
429   mapTM    = mapA
430
431mapA :: (a->b) -> AltMap a -> AltMap b
432mapA f (AM { am_deflt = adeflt, am_data = adata, am_lit = alit })
433  = AM { am_deflt = mapTM f adeflt
434       , am_data = mapNameEnv (mapTM f) adata
435       , am_lit = mapTM (mapTM f) alit }
436 
437lkA :: CmEnv -> CoreAlt -> AltMap a -> Maybe a
438lkA env (DEFAULT,    _, rhs)  = am_deflt >.> lkE env rhs
439lkA env (LitAlt lit, _, rhs)  = am_lit >.> lkLit lit >=> lkE env rhs
440lkA env (DataAlt dc, bs, rhs) = am_data >.> lkNamed dc >=> lkE (extendCMEs env bs) rhs
441
442xtA :: CmEnv -> CoreAlt -> XT a -> AltMap a -> AltMap a
443xtA env (DEFAULT, _, rhs)    f m = m { am_deflt = am_deflt m |> xtE env rhs f }
444xtA env (LitAlt l, _, rhs)   f m = m { am_lit   = am_lit m   |> xtLit l |>> xtE env rhs f }
445xtA env (DataAlt d, bs, rhs) f m = m { am_data  = am_data m  |> xtNamed d 
446                                                             |>> xtE (extendCMEs env bs) rhs f }
447
448fdA :: (a -> b -> b) -> AltMap a -> b -> b
449fdA k m = foldTM k (am_deflt m)
450        . foldTM (foldTM k) (am_data m)
451        . foldTM (foldTM k) (am_lit m)
452\end{code}
453
454%************************************************************************
455%*									*
456                   Coercions
457%*									*
458%************************************************************************
459
460\begin{code}
461data CoercionMap a 
462  = EmptyKM
463  | KM { km_refl :: TypeMap a
464       , km_tc_app :: NameEnv (ListMap CoercionMap a)
465       , km_app    :: CoercionMap (CoercionMap a)
466       , km_forall :: CoercionMap (TypeMap a)
467       , km_var    :: VarMap a
468       , km_axiom  :: NameEnv (ListMap CoercionMap a)
469       , km_unsafe :: TypeMap (TypeMap a)
470       , km_sym    :: CoercionMap a
471       , km_trans  :: CoercionMap (CoercionMap a)
472       , km_nth    :: IntMap.IntMap (CoercionMap a)
473       , km_left   :: CoercionMap a
474       , km_right  :: CoercionMap a
475       , km_inst   :: CoercionMap (TypeMap a) }
476
477wrapEmptyKM :: CoercionMap a
478wrapEmptyKM = KM { km_refl = emptyTM, km_tc_app = emptyNameEnv
479                 , km_app = emptyTM, km_forall = emptyTM
480                 , km_var = emptyTM, km_axiom = emptyNameEnv
481                 , km_unsafe = emptyTM, km_sym = emptyTM, km_trans = emptyTM
482                 , km_nth = emptyTM, km_left = emptyTM, km_right = emptyTM
483                 , km_inst = emptyTM }
484
485instance TrieMap CoercionMap where
486   type Key CoercionMap = Coercion
487   emptyTM  = EmptyKM
488   lookupTM = lkC emptyCME
489   alterTM  = xtC emptyCME
490   foldTM   = fdC
491   mapTM    = mapC
492
493mapC :: (a->b) -> CoercionMap a -> CoercionMap b
494mapC _ EmptyKM = EmptyKM
495mapC f (KM { km_refl = krefl, km_tc_app = ktc
496           , km_app = kapp, km_forall = kforall
497           , km_var = kvar, km_axiom = kax
498           , km_unsafe = kunsafe, km_sym = ksym, km_trans = ktrans
499           , km_nth = knth, km_left = kml, km_right = kmr
500           , km_inst = kinst })
501  = KM { km_refl   = mapTM f krefl
502       , km_tc_app = mapNameEnv (mapTM f) ktc
503       , km_app    = mapTM (mapTM f) kapp
504       , km_forall = mapTM (mapTM f) kforall
505       , km_var    = mapTM f kvar
506       , km_axiom  = mapNameEnv (mapTM f) kax
507       , km_unsafe = mapTM (mapTM f) kunsafe
508       , km_sym    = mapTM f ksym
509       , km_trans  = mapTM (mapTM f) ktrans
510       , km_nth    = IntMap.map (mapTM f) knth
511       , km_left   = mapTM f kml
512       , km_right  = mapTM f kmr
513       , km_inst   = mapTM (mapTM f) kinst }
514
515lkC :: CmEnv -> Coercion -> CoercionMap a -> Maybe a
516lkC env co m 
517  | EmptyKM <- m = Nothing
518  | otherwise    = go co m
519  where
520    go (Refl ty)           = km_refl   >.> lkT env ty
521    go (TyConAppCo tc cs)  = km_tc_app >.> lkNamed tc >=> lkList (lkC env) cs
522    go (AxiomInstCo ax cs) = km_axiom  >.> lkNamed ax >=> lkList (lkC env) cs
523    go (AppCo c1 c2)       = km_app    >.> lkC env c1 >=> lkC env c2
524    go (TransCo c1 c2)     = km_trans  >.> lkC env c1 >=> lkC env c2
525    go (UnsafeCo t1 t2)    = km_unsafe >.> lkT env t1 >=> lkT env t2
526    go (InstCo c t)        = km_inst   >.> lkC env c  >=> lkT env t
527    go (ForAllCo v c)      = km_forall >.> lkC (extendCME env v) c >=> lkBndr env v
528    go (CoVarCo v)         = km_var    >.> lkVar env v
529    go (SymCo c)           = km_sym    >.> lkC env c
530    go (NthCo n c)         = km_nth    >.> lookupTM n >=> lkC env c
531    go (LRCo CLeft  c)     = km_left   >.> lkC env c
532    go (LRCo CRight c)     = km_right  >.> lkC env c
533
534xtC :: CmEnv -> Coercion -> XT a -> CoercionMap a -> CoercionMap a
535xtC env co f EmptyKM = xtC env co f wrapEmptyKM
536xtC env (Refl ty)           f m = m { km_refl   = km_refl m   |> xtT env ty f }
537xtC env (TyConAppCo tc cs)  f m = m { km_tc_app = km_tc_app m |> xtNamed tc |>> xtList (xtC env) cs f }
538xtC env (AxiomInstCo ax cs) f m = m { km_axiom  = km_axiom m  |> xtNamed ax |>> xtList (xtC env) cs f }
539xtC env (AppCo c1 c2)       f m = m { km_app    = km_app m    |> xtC env c1 |>> xtC env c2 f }
540xtC env (TransCo c1 c2)     f m = m { km_trans  = km_trans m  |> xtC env c1 |>> xtC env c2 f }
541xtC env (UnsafeCo t1 t2)    f m = m { km_unsafe = km_unsafe m |> xtT env t1 |>> xtT env t2 f }
542xtC env (InstCo c t)        f m = m { km_inst   = km_inst m   |> xtC env c  |>> xtT env t  f }
543xtC env (ForAllCo v c)      f m = m { km_forall = km_forall m |> xtC (extendCME env v) c 
544                                                  |>> xtBndr env v f }
545xtC env (CoVarCo v)         f m = m { km_var 	= km_var m   |> xtVar env v f }
546xtC env (SymCo c)           f m = m { km_sym 	= km_sym m   |> xtC env   c f }
547xtC env (NthCo n c)         f m = m { km_nth 	= km_nth m   |> xtInt n |>> xtC env c f } 
548xtC env (LRCo CLeft  c)     f m = m { km_left 	= km_left  m |> xtC env c f } 
549xtC env (LRCo CRight c)     f m = m { km_right 	= km_right m |> xtC env c f } 
550
551fdC :: (a -> b -> b) -> CoercionMap a -> b -> b
552fdC _ EmptyKM = \z -> z
553fdC k m = foldTM k (km_refl m)
554        . foldTM (foldTM k) (km_tc_app m)
555        . foldTM (foldTM k) (km_app m)
556        . foldTM (foldTM k) (km_forall m)
557        . foldTM k (km_var m)
558        . foldTM (foldTM k) (km_axiom m)
559        . foldTM (foldTM k) (km_unsafe m)
560        . foldTM k (km_sym m)
561        . foldTM (foldTM k) (km_trans m)
562        . foldTM (foldTM k) (km_nth m)
563        . foldTM k          (km_left m)
564        . foldTM k          (km_right m)
565        . foldTM (foldTM k) (km_inst m)
566\end{code}
567
568
569%************************************************************************
570%*									*
571                   Types
572%*									*
573%************************************************************************
574
575\begin{code}
576data TypeMap a
577  = EmptyTM
578  | TM { tm_var   :: VarMap a
579       , tm_app    :: TypeMap (TypeMap a)
580       , tm_fun    :: TypeMap (TypeMap a)
581       , tm_tc_app :: NameEnv (ListMap TypeMap a)
582       , tm_forall :: TypeMap (BndrMap a)
583       , tm_tylit  :: TyLitMap a
584       }
585
586
587instance Outputable a => Outputable (TypeMap a) where
588  ppr m = text "TypeMap elts" <+> ppr (foldTypeMap (:) [] m)
589
590foldTypeMap :: (a -> b -> b) -> b -> TypeMap a -> b
591foldTypeMap k z m = fdT k m z
592
593wrapEmptyTypeMap :: TypeMap a
594wrapEmptyTypeMap = TM { tm_var  = emptyTM
595                      , tm_app  = EmptyTM
596                      , tm_fun  = EmptyTM
597                      , tm_tc_app = emptyNameEnv
598                      , tm_forall = EmptyTM
599                      , tm_tylit  = emptyTyLitMap }
600
601instance TrieMap TypeMap where
602   type Key TypeMap = Type
603   emptyTM  = EmptyTM
604   lookupTM = lkT emptyCME
605   alterTM  = xtT emptyCME
606   foldTM   = fdT
607   mapTM    = mapT
608
609mapT :: (a->b) -> TypeMap a -> TypeMap b
610mapT _ EmptyTM = EmptyTM
611mapT f (TM { tm_var  = tvar, tm_app = tapp, tm_fun = tfun
612           , tm_tc_app = ttcapp, tm_forall = tforall, tm_tylit = tlit })
613  = TM { tm_var    = mapTM f tvar
614       , tm_app    = mapTM (mapTM f) tapp
615       , tm_fun    = mapTM (mapTM f) tfun
616       , tm_tc_app = mapNameEnv (mapTM f) ttcapp
617       , tm_forall = mapTM (mapTM f) tforall
618       , tm_tylit  = mapTM f tlit }
619
620-----------------
621lkT :: CmEnv -> Type -> TypeMap a -> Maybe a
622lkT env ty m
623  | EmptyTM <- m = Nothing
624  | otherwise    = go ty m
625  where
626    go ty | Just ty' <- coreView ty = go ty'
627    go (TyVarTy v)       = tm_var    >.> lkVar env v
628    go (AppTy t1 t2)     = tm_app    >.> lkT env t1 >=> lkT env t2
629    go (FunTy t1 t2)     = tm_fun    >.> lkT env t1 >=> lkT env t2
630    go (TyConApp tc tys) = tm_tc_app >.> lkNamed tc >=> lkList (lkT env) tys
631    go (LitTy l)         = tm_tylit  >.> lkTyLit l
632    go (ForAllTy tv ty)  = tm_forall >.> lkT (extendCME env tv) ty >=> lkBndr env tv
633
634
635lkT_mod :: CmEnv  
636        -> TyVarEnv Type -- TvSubstEnv 
637        -> Type
638        -> TypeMap b -> Maybe b 
639lkT_mod env s ty m
640  | EmptyTM <- m = Nothing
641  | Just ty' <- coreView ty
642  = lkT_mod env s ty' m
643  | [] <- candidates 
644  = go env s ty m
645  | otherwise
646  = Just $ snd (head candidates) -- Yikes!
647  where
648     -- Hopefully intersects is much smaller than traversing the whole vm_fvar
649    intersects = eltsUFM $
650                 intersectUFM_C (,) s (vm_fvar $ tm_var m)
651    candidates = [ (u,ct) | (u,ct) <- intersects
652                          , Type.substTy (niFixTvSubst s) u `eqType` ty ]
653                  
654    go env _s (TyVarTy v)      = tm_var    >.> lkVar env v
655    go env s (AppTy t1 t2)     = tm_app    >.> lkT_mod env s t1 >=> lkT_mod env s t2
656    go env s (FunTy t1 t2)     = tm_fun    >.> lkT_mod env s t1 >=> lkT_mod env s t2
657    go env s (TyConApp tc tys) = tm_tc_app >.> lkNamed tc >=> lkList (lkT_mod env s) tys
658    go _env _s (LitTy l)       = tm_tylit  >.> lkTyLit l
659    go _env _s (ForAllTy _tv _ty) = const Nothing
660    
661    {- DV TODO: Add proper lookup for ForAll -}
662
663lookupTypeMap_mod :: TyVarEnv a -- A substitution to be applied to the /keys/ of type map 
664                  -> (a -> Type)
665                  -> Type 
666                  -> TypeMap b -> Maybe b
667lookupTypeMap_mod s f = lkT_mod emptyCME (mapVarEnv f s)
668
669-----------------
670xtT :: CmEnv -> Type -> XT a -> TypeMap a -> TypeMap a
671xtT env ty f m
672  | EmptyTM <- m            = xtT env ty  f wrapEmptyTypeMap 
673  | Just ty' <- coreView ty = xtT env ty' f m                
674
675xtT env (TyVarTy v)       f  m = m { tm_var    = tm_var m |> xtVar env v f }
676xtT env (AppTy t1 t2)     f  m = m { tm_app    = tm_app m |> xtT env t1 |>> xtT env t2 f }
677xtT env (FunTy t1 t2)     f  m = m { tm_fun    = tm_fun m |> xtT env t1 |>> xtT env t2 f }
678xtT env (ForAllTy tv ty)  f  m = m { tm_forall = tm_forall m |> xtT (extendCME env tv) ty 
679                                                 |>> xtBndr env tv f }
680xtT env (TyConApp tc tys) f  m = m { tm_tc_app = tm_tc_app m |> xtNamed tc 
681                                                 |>> xtList (xtT env) tys f }
682xtT _   (LitTy l)         f  m = m { tm_tylit  = tm_tylit m |> xtTyLit l f }
683
684fdT :: (a -> b -> b) -> TypeMap a -> b -> b
685fdT _ EmptyTM = \z -> z
686fdT k m = foldTM k (tm_var m)
687        . foldTM (foldTM k) (tm_app m)
688        . foldTM (foldTM k) (tm_fun m)
689        . foldTM (foldTM k) (tm_tc_app m)
690        . foldTM (foldTM k) (tm_forall m)
691        . foldTyLit k (tm_tylit m)
692
693
694
695------------------------
696data TyLitMap a = TLM { tlm_number :: Map.Map Integer a
697                      , tlm_string :: Map.Map FastString a
698                      }
699
700instance TrieMap TyLitMap where
701   type Key TyLitMap = TyLit
702   emptyTM  = emptyTyLitMap
703   lookupTM = lkTyLit
704   alterTM  = xtTyLit
705   foldTM   = foldTyLit
706   mapTM    = mapTyLit
707   
708emptyTyLitMap :: TyLitMap a
709emptyTyLitMap = TLM { tlm_number = Map.empty, tlm_string = Map.empty }
710
711mapTyLit :: (a->b) -> TyLitMap a -> TyLitMap b
712mapTyLit f (TLM { tlm_number = tn, tlm_string = ts })
713  = TLM { tlm_number = Map.map f tn, tlm_string = Map.map f ts }
714
715lkTyLit :: TyLit -> TyLitMap a -> Maybe a
716lkTyLit l =
717  case l of
718    NumTyLit n -> tlm_number >.> Map.lookup n
719    StrTyLit n -> tlm_string >.> Map.lookup n
720
721xtTyLit :: TyLit -> XT a -> TyLitMap a -> TyLitMap a
722xtTyLit l f m =
723  case l of
724    NumTyLit n -> m { tlm_number = tlm_number m |> Map.alter f n }
725    StrTyLit n -> m { tlm_string = tlm_string m |> Map.alter f n }
726
727foldTyLit :: (a -> b -> b) -> TyLitMap a -> b -> b
728foldTyLit l m = flip (Map.fold l) (tlm_string m)
729              . flip (Map.fold l) (tlm_number m)
730\end{code}
731
732
733%************************************************************************
734%*									*
735                   Variables
736%*									*
737%************************************************************************
738
739\begin{code}
740type BoundVar = Int  -- Bound variables are deBruijn numbered
741type BoundVarMap a = IntMap.IntMap a
742
743data CmEnv = CME { cme_next :: BoundVar
744                 , cme_env  :: VarEnv BoundVar } 
745
746emptyCME :: CmEnv
747emptyCME = CME { cme_next = 0, cme_env = emptyVarEnv }
748
749extendCME :: CmEnv -> Var -> CmEnv
750extendCME (CME { cme_next = bv, cme_env = env }) v
751  = CME { cme_next = bv+1, cme_env = extendVarEnv env v bv }
752
753extendCMEs :: CmEnv -> [Var] -> CmEnv
754extendCMEs env vs = foldl extendCME env vs
755
756lookupCME :: CmEnv -> Var -> Maybe BoundVar
757lookupCME (CME { cme_env = env }) v = lookupVarEnv env v
758
759--------- Variable binders -------------
760type BndrMap = TypeMap 
761
762lkBndr :: CmEnv -> Var -> BndrMap a -> Maybe a
763lkBndr env v m = lkT env (varType v) m
764
765xtBndr :: CmEnv -> Var -> XT a -> BndrMap a -> BndrMap a
766xtBndr env v f = xtT env (varType v) f
767
768--------- Variable occurrence -------------
769data VarMap a = VM { vm_bvar   :: BoundVarMap a  -- Bound variable
770                   , vm_fvar   :: VarEnv a }  	  -- Free variable
771
772instance TrieMap VarMap where
773   type Key VarMap = Var
774   emptyTM  = VM { vm_bvar = IntMap.empty, vm_fvar = emptyVarEnv }
775   lookupTM = lkVar emptyCME
776   alterTM  = xtVar emptyCME
777   foldTM   = fdVar
778   mapTM    = mapVar
779
780mapVar :: (a->b) -> VarMap a -> VarMap b
781mapVar f (VM { vm_bvar = bv, vm_fvar = fv })
782  = VM { vm_bvar = mapTM f bv, vm_fvar = mapVarEnv f fv }
783
784lkVar :: CmEnv -> Var -> VarMap a -> Maybe a
785lkVar env v 
786  | Just bv <- lookupCME env v = vm_bvar >.> lookupTM bv
787  | otherwise                  = vm_fvar >.> lkFreeVar v
788
789xtVar :: CmEnv -> Var -> XT a -> VarMap a -> VarMap a
790xtVar env v f m
791  | Just bv <- lookupCME env v = m { vm_bvar = vm_bvar m |> xtInt bv f }
792  | otherwise                  = m { vm_fvar = vm_fvar m |> xtFreeVar v f }
793
794fdVar :: (a -> b -> b) -> VarMap a -> b -> b
795fdVar k m = foldTM k (vm_bvar m)
796          . foldTM k (vm_fvar m)
797
798lkFreeVar :: Var -> VarEnv a -> Maybe a
799lkFreeVar var env = lookupVarEnv env var
800
801xtFreeVar :: Var -> XT a -> VarEnv a -> VarEnv a
802xtFreeVar v f m = alterVarEnv f m v
803\end{code}