compiler/basicTypes/VarEnv.lhs

   1 %
   2 % (c) The University of Glasgow 2006
   3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   4 %
   5
   6 \begin{code}
   7 module VarEnv (
   8         VarEnv, IdEnv, TyVarEnv,
   9         emptyVarEnv, unitVarEnv, mkVarEnv,
  10         elemVarEnv, varEnvElts, varEnvKeys,
  11         extendVarEnv, extendVarEnv_C, extendVarEnvList,
  12         plusVarEnv, plusVarEnv_C,
  13         delVarEnvList, delVarEnv,
  14         lookupVarEnv, lookupVarEnv_NF, lookupWithDefaultVarEnv,
  15         mapVarEnv, zipVarEnv,
  16         modifyVarEnv, modifyVarEnv_Directly,
  17         isEmptyVarEnv, foldVarEnv,
  18         elemVarEnvByKey, lookupVarEnv_Directly,
  19         filterVarEnv_Directly,
  20
  21         -- InScopeSet
  22         InScopeSet, emptyInScopeSet, mkInScopeSet, delInScopeSet,
  23         extendInScopeSet, extendInScopeSetList, extendInScopeSetSet,
  24         modifyInScopeSet,
  25         getInScopeVars, lookupInScope, elemInScopeSet, uniqAway,
  26         mapInScopeSet,
  27
  28         -- RnEnv2 and its operations
  29         RnEnv2, mkRnEnv2, rnBndr2, rnBndrs2, rnOccL, rnOccR, inRnEnvL, inRnEnvR,
  30                 rnBndrL, rnBndrR, nukeRnEnvL, nukeRnEnvR, extendRnInScopeList,
  31                 rnInScope, rnInScopeSet, lookupRnInScope,
  32
  33         -- TidyEnvs
  34         TidyEnv, emptyTidyEnv
  35     ) where
  36
  37 #include "HsVersions.h"
  38
  39 import OccName
  40 import Var
  41 import VarSet
  42 import UniqFM
  43 import Unique
  44 import Util
  45 import Maybes
  46 import StaticFlags
  47 import Outputable
  48 import FastTypes
  49 \end{code}
  50
  51
  52 %************************************************************************
  53 %*                                                                      *
  54                 In-scope sets
  55 %*                                                                      *
  56 %************************************************************************
  57
  58 \begin{code}
  59 data InScopeSet = InScope (VarEnv Var) FastInt
  60         -- The Int# is a kind of hash-value used by uniqAway
  61         -- For example, it might be the size of the set
  62         -- INVARIANT: it's not zero; we use it as a multiplier in uniqAway
  63
  64 instance Outputable InScopeSet where
  65   ppr (InScope s i) = ptext SLIT("InScope") <+> ppr s
  66
  67 emptyInScopeSet :: InScopeSet
  68 emptyInScopeSet = InScope emptyVarSet 1#
  69
  70 getInScopeVars ::  InScopeSet -> VarEnv Var
  71 getInScopeVars (InScope vs _) = vs
  72
  73 mkInScopeSet :: VarEnv Var -> InScopeSet
  74 mkInScopeSet in_scope = InScope in_scope 1#
  75
  76 extendInScopeSet :: InScopeSet -> Var -> InScopeSet
  77 extendInScopeSet (InScope in_scope n) v = InScope (extendVarEnv in_scope v v) (n +# 1#)
  78
  79 extendInScopeSetList :: InScopeSet -> [Var] -> InScopeSet
  80 extendInScopeSetList (InScope in_scope n) vs
  81    = InScope (foldl (\s v -> extendVarEnv s v v) in_scope vs)
  82                     (n +# iUnbox (length vs))
  83
  84 extendInScopeSetSet :: InScopeSet -> VarEnv Var -> InScopeSet
  85 extendInScopeSetSet (InScope in_scope n) vs
  86    = InScope (in_scope `plusVarEnv` vs) (n +# iUnbox (sizeUFM vs))
  87
  88 modifyInScopeSet :: InScopeSet -> Var -> Var -> InScopeSet
  89 -- Exploit the fact that the in-scope "set" is really a map
  90 --      Make old_v map to new_v
  91 modifyInScopeSet (InScope in_scope n) old_v new_v = InScope (extendVarEnv in_scope old_v new_v) (n +# 1#)
  92
  93 delInScopeSet :: InScopeSet -> Var -> InScopeSet
  94 delInScopeSet (InScope in_scope n) v = InScope (in_scope `delVarEnv` v) n
  95
  96 mapInScopeSet :: (Var -> Var) -> InScopeSet -> InScopeSet
  97 mapInScopeSet f (InScope in_scope n) = InScope (mapVarEnv f in_scope) n
  98
  99 elemInScopeSet :: Var -> InScopeSet -> Bool
 100 elemInScopeSet v (InScope in_scope n) = v `elemVarEnv` in_scope
 101
 102 lookupInScope :: InScopeSet -> Var -> Maybe Var
 103 -- It's important to look for a fixed point
 104 -- When we see (case x of y { I# v -> ... })
 105 -- we add  [x -> y] to the in-scope set (Simplify.simplCaseBinder).
 106 -- When we lookup up an occurrence of x, we map to y, but then
 107 -- we want to look up y in case it has acquired more evaluation information by now.
 108 lookupInScope (InScope in_scope n) v
 109   = go v
 110   where
 111     go v = case lookupVarEnv in_scope v of
 112                 Just v' | v == v'   -> Just v'  -- Reached a fixed point
 113                         | otherwise -> go v'
 114                 Nothing             -> Nothing
 115 \end{code}
 116
 117 \begin{code}
 118 uniqAway :: InScopeSet -> Var -> Var
 119 -- (uniqAway in_scope v) finds a unique that is not used in the
 120 -- in-scope set, and gives that to v.  It starts with v's current unique, of course,
 121 -- in the hope that it won't have to change it, and thereafter uses a combination
 122 -- of that and the hash-code found in the in-scope set
 123 uniqAway in_scope var
 124   | var `elemInScopeSet` in_scope = uniqAway' in_scope var      -- Make a new one
 125   | otherwise                     = var                         -- Nothing to do
 126
 127 uniqAway' :: InScopeSet -> Var -> Var
 128 -- This one *always* makes up a new variable
 129 uniqAway' (InScope set n) var
 130   = try 1#
 131   where
 132     orig_unique = getUnique var
 133     try k
 134 #ifdef DEBUG
 135           | k ># 1000#
 136           = pprPanic "uniqAway loop:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n))
 137 #endif
 138           | uniq `elemVarSetByKey` set = try (k +# 1#)
 139 #ifdef DEBUG
 140           | opt_PprStyle_Debug && k ># 3#
 141           = pprTrace "uniqAway:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n))
 142             setVarUnique var uniq
 143 #endif
 144           | otherwise = setVarUnique var uniq
 145           where
 146             uniq = deriveUnique orig_unique (iBox (n *# k))
 147 \end{code}
 148
 149
 150 %************************************************************************
 151 %*                                                                      *
 152                 Dual renaming
 153 %*                                                                      *
 154 %************************************************************************
 155
 156 When we are comparing (or matching) types or terms, we are faced with
 157 "going under" corresponding binders.  E.g. when comparing
 158         \x. e1  ~   \y. e2
 159
 160 Basically we want to rename [x->y] or [y->x], but there are lots of
 161 things we must be careful of.  In particular, x might be free in e2, or
 162 y in e1.  So the idea is that we come up with a fresh binder that is free
 163 in neither, and rename x and y respectively.  That means we must maintain
 164         a) a renaming for the left-hand expression
 165         b) a renaming for the right-hand expressions
 166         c) an in-scope set
 167
 168 Furthermore, when matching, we want to be able to have an 'occurs check',
 169 to prevent
 170         \x. f   ~   \y. y
 171 matching with f->y.  So for each expression we want to know that set of
 172 locally-bound variables. That is precisely the domain of the mappings (a)
 173 and (b), but we must ensure that we always extend the mappings as we go in.
 174
 175
 176 \begin{code}
 177 data RnEnv2
 178   = RV2 { envL     :: VarEnv Var        -- Renaming for Left term
 179         , envR     :: VarEnv Var        -- Renaming for Right term
 180         , in_scope :: InScopeSet }      -- In scope in left or right terms
 181
 182 -- The renamings envL and envR are *guaranteed* to contain a binding
 183 -- for every variable bound as we go into the term, even if it is not
 184 -- renamed.  That way we can ask what variables are locally bound
 185 -- (inRnEnvL, inRnEnvR)
 186
 187 mkRnEnv2 :: InScopeSet -> RnEnv2
 188 mkRnEnv2 vars = RV2     { envL     = emptyVarEnv
 189                         , envR     = emptyVarEnv
 190                         , in_scope = vars }
 191
 192 extendRnInScopeList :: RnEnv2 -> [Var] -> RnEnv2
 193 extendRnInScopeList env vs
 194   = env { in_scope = extendInScopeSetList (in_scope env) vs }
 195
 196 rnInScope :: Var -> RnEnv2 -> Bool
 197 rnInScope x env = x `elemInScopeSet` in_scope env
 198
 199 rnInScopeSet :: RnEnv2 -> InScopeSet
 200 rnInScopeSet = in_scope
 201
 202 rnBndrs2 :: RnEnv2 -> [Var] -> [Var] -> RnEnv2
 203 -- Arg lists must be of equal length
 204 rnBndrs2 env bsL bsR = foldl2 rnBndr2 env bsL bsR
 205
 206 rnBndr2 :: RnEnv2 -> Var -> Var -> RnEnv2
 207 -- (rnBndr2 env bL bR) go under a binder bL in the Left term 1,
 208 --                     and binder bR in the Right term
 209 -- It finds a new binder, new_b,
 210 -- and returns an environment mapping bL->new_b and bR->new_b resp.
 211 rnBndr2 (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL bR
 212   = RV2 { envL     = extendVarEnv envL bL new_b   -- See Note
 213         , envR     = extendVarEnv envR bR new_b   -- [Rebinding]
 214         , in_scope = extendInScopeSet in_scope new_b }
 215   where
 216         -- Find a new binder not in scope in either term
 217     new_b | not (bL `elemInScopeSet` in_scope) = bL
 218           | not (bR `elemInScopeSet` in_scope) = bR
 219           | otherwise                          = uniqAway' in_scope bL
 220
 221         -- Note [Rebinding]
 222         -- If the new var is the same as the old one, note that
 223         -- the extendVarEnv *deletes* any current renaming
 224         -- E.g.   (\x. \x. ...)  ~  (\y. \z. ...)
 225         --
 226         --   Inside \x  \y      { [x->y], [y->y],       {y} }
 227         --       \x  \z         { [x->x], [y->y, z->x], {y,x} }
 228
 229 rnBndrL, rnBndrR :: RnEnv2 -> Var -> (RnEnv2, Var)
 230 -- Used when there's a binder on one side or the other only
 231 -- Useful when eta-expanding
 232 --
 233 rnBndrL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
 234   = (RV2 { envL     = extendVarEnv envL bL new_b
 235          , envR     = extendVarEnv envR new_b new_b     -- Note [rnBndrLR]
 236          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 237   where
 238     new_b = uniqAway in_scope bL
 239
 240 rnBndrR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
 241   = (RV2 { envL     = extendVarEnv envL new_b new_b     -- Note [rnBndrLR]
 242          , envR     = extendVarEnv envR bR new_b
 243          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 244   where
 245     new_b = uniqAway in_scope bR
 246
 247 -- Note [rnBndrLR]
 248 -- ~~~~~~~~~~~~~~~
 249 -- Notice that in rnBndrL, rnBndrR, we extend envR, envL respectively
 250 -- with a binding [new_b -> new_b], where new_b is the new binder.
 251 -- This is important when doing eta expansion; e.g. matching (\x.M) ~ N
 252 -- In effect we switch to (\x'.M) ~ (\x'.N x'), where x' is new_b
 253 -- So we must add x' to the env of both L and R.  (x' is fresh, so it
 254 -- can't capture anything in N.)
 255 --
 256 -- If we don't do this, we can get silly matches like
 257 --      forall a.  \y.a  ~   v
 258 -- succeeding with [x -> v y], which is bogus of course
 259
 260 rnOccL, rnOccR :: RnEnv2 -> Var -> Var
 261 -- Look up the renaming of an occurrence in the left or right term
 262 rnOccL (RV2 { envL = env }) v = lookupVarEnv env v `orElse` v
 263 rnOccR (RV2 { envR = env }) v = lookupVarEnv env v `orElse` v
 264
 265 inRnEnvL, inRnEnvR :: RnEnv2 -> Var -> Bool
 266 -- Tells whether a variable is locally bound
 267 inRnEnvL (RV2 { envL = env }) v = v `elemVarEnv` env
 268 inRnEnvR (RV2 { envR = env }) v = v `elemVarEnv` env
 269
 270 lookupRnInScope :: RnEnv2 -> Var -> Var
 271 lookupRnInScope env v = lookupInScope (in_scope env) v `orElse` v
 272
 273 nukeRnEnvL, nukeRnEnvR :: RnEnv2 -> RnEnv2
 274 nukeRnEnvL env = env { envL = emptyVarEnv }
 275 nukeRnEnvR env = env { envR = emptyVarEnv }
 276 \end{code}
 277
 278
 279 %************************************************************************
 280 %*                                                                      *
 281                 Tidying
 282 %*                                                                      *
 283 %************************************************************************
 284
 285 When tidying up print names, we keep a mapping of in-scope occ-names
 286 (the TidyOccEnv) and a Var-to-Var of the current renamings.
 287
 288 \begin{code}
 289 type TidyEnv = (TidyOccEnv, VarEnv Var)
 290
 291 emptyTidyEnv :: TidyEnv
 292 emptyTidyEnv = (emptyTidyOccEnv, emptyVarEnv)
 293 \end{code}
 294
 295
 296 %************************************************************************
 297 %*                                                                      *
 298 \subsection{@VarEnv@s}
 299 %*                                                                      *
 300 %************************************************************************
 301
 302 \begin{code}
 303 type VarEnv elt   = UniqFM elt
 304 type IdEnv elt    = VarEnv elt
 305 type TyVarEnv elt = VarEnv elt
 306
 307 emptyVarEnv       :: VarEnv a
 308 mkVarEnv          :: [(Var, a)] -> VarEnv a
 309 zipVarEnv         :: [Var] -> [a] -> VarEnv a
 310 unitVarEnv        :: Var -> a -> VarEnv a
 311 extendVarEnv      :: VarEnv a -> Var -> a -> VarEnv a
 312 extendVarEnv_C    :: (a->a->a) -> VarEnv a -> Var -> a -> VarEnv a
 313 plusVarEnv        :: VarEnv a -> VarEnv a -> VarEnv a
 314 extendVarEnvList  :: VarEnv a -> [(Var, a)] -> VarEnv a
 315
 316 lookupVarEnv_Directly :: VarEnv a -> Unique -> Maybe a
 317 filterVarEnv_Directly :: (Unique -> a -> Bool) -> VarEnv a -> VarEnv a
 318 delVarEnvList     :: VarEnv a -> [Var] -> VarEnv a
 319 delVarEnv         :: VarEnv a -> Var -> VarEnv a
 320 plusVarEnv_C      :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
 321 mapVarEnv         :: (a -> b) -> VarEnv a -> VarEnv b
 322 modifyVarEnv      :: (a -> a) -> VarEnv a -> Var -> VarEnv a
 323 varEnvElts        :: VarEnv a -> [a]
 324 varEnvKeys        :: VarEnv a -> [Unique]
 325
 326 isEmptyVarEnv     :: VarEnv a -> Bool
 327 lookupVarEnv      :: VarEnv a -> Var -> Maybe a
 328 lookupVarEnv_NF   :: VarEnv a -> Var -> a
 329 lookupWithDefaultVarEnv :: VarEnv a -> a -> Var -> a
 330 elemVarEnv        :: Var -> VarEnv a -> Bool
 331 elemVarEnvByKey   :: Unique -> VarEnv a -> Bool
 332 foldVarEnv        :: (a -> b -> b) -> b -> VarEnv a -> b
 333 \end{code}
 334
 335 \begin{code}
 336 elemVarEnv       = elemUFM
 337 elemVarEnvByKey  = elemUFM_Directly
 338 extendVarEnv     = addToUFM
 339 extendVarEnv_C   = addToUFM_C
 340 extendVarEnvList = addListToUFM
 341 plusVarEnv_C     = plusUFM_C
 342 delVarEnvList    = delListFromUFM
 343 delVarEnv        = delFromUFM
 344 plusVarEnv       = plusUFM
 345 lookupVarEnv     = lookupUFM
 346 lookupWithDefaultVarEnv = lookupWithDefaultUFM
 347 mapVarEnv        = mapUFM
 348 mkVarEnv         = listToUFM
 349 emptyVarEnv      = emptyUFM
 350 varEnvElts       = eltsUFM
 351 varEnvKeys       = keysUFM
 352 unitVarEnv       = unitUFM
 353 isEmptyVarEnv    = isNullUFM
 354 foldVarEnv       = foldUFM
 355 lookupVarEnv_Directly = lookupUFM_Directly
 356 filterVarEnv_Directly = filterUFM_Directly
 357
 358 zipVarEnv tyvars tys   = mkVarEnv (zipEqual "zipVarEnv" tyvars tys)
 359 lookupVarEnv_NF env id = case (lookupVarEnv env id) of { Just xx -> xx }
 360 \end{code}
 361
 362 @modifyVarEnv@: Look up a thing in the VarEnv,
 363 then mash it with the modify function, and put it back.
 364
 365 \begin{code}
 366 modifyVarEnv mangle_fn env key
 367   = case (lookupVarEnv env key) of
 368       Nothing -> env
 369       Just xx -> extendVarEnv env key (mangle_fn xx)
 370
 371 modifyVarEnv_Directly mangle_fn env key
 372   = case (lookupUFM_Directly env key) of
 373       Nothing -> env
 374       Just xx -> addToUFM_Directly env key (mangle_fn xx)
 375 \end{code}