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, 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 rnBndrs2 :: RnEnv2 -> [Var] -> [Var] -> RnEnv2
 200 -- Arg lists must be of equal length
 201 rnBndrs2 env bsL bsR = foldl2 rnBndr2 env bsL bsR
 202
 203 rnBndr2 :: RnEnv2 -> Var -> Var -> RnEnv2
 204 -- (rnBndr2 env bL bR) go under a binder bL in the Left term 1,
 205 --                     and binder bR in the Right term
 206 -- It finds a new binder, new_b,
 207 -- and returns an environment mapping bL->new_b and bR->new_b resp.
 208 rnBndr2 (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL bR
 209   = RV2 { envL     = extendVarEnv envL bL new_b   -- See Note
 210         , envR     = extendVarEnv envR bR new_b   -- [Rebinding]
 211         , in_scope = extendInScopeSet in_scope new_b }
 212   where
 213         -- Find a new binder not in scope in either term
 214     new_b | not (bL `elemInScopeSet` in_scope) = bL
 215           | not (bR `elemInScopeSet` in_scope) = bR
 216           | otherwise                          = uniqAway' in_scope bL
 217
 218         -- Note [Rebinding]
 219         -- If the new var is the same as the old one, note that
 220         -- the extendVarEnv *deletes* any current renaming
 221         -- E.g.   (\x. \x. ...)  ~  (\y. \z. ...)
 222         --
 223         --   Inside \x  \y      { [x->y], [y->y],       {y} }
 224         --       \x  \z         { [x->x], [y->y, z->x], {y,x} }
 225
 226 rnBndrL, rnBndrR :: RnEnv2 -> Var -> (RnEnv2, Var)
 227 -- Used when there's a binder on one side or the other only
 228 -- Useful when eta-expanding
 229 --
 230 rnBndrL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
 231   = (RV2 { envL     = extendVarEnv envL bL new_b
 232          , envR     = extendVarEnv envR new_b new_b     -- Note [rnBndrLR]
 233          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 234   where
 235     new_b = uniqAway in_scope bL
 236
 237 rnBndrR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
 238   = (RV2 { envL     = extendVarEnv envL new_b new_b     -- Note [rnBndrLR]
 239          , envR     = extendVarEnv envR bR new_b
 240          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 241   where
 242     new_b = uniqAway in_scope bR
 243
 244 -- Note [rnBndrLR]
 245 -- ~~~~~~~~~~~~~~~
 246 -- Notice that in rnBndrL, rnBndrR, we extend envR, envL respectively
 247 -- with a binding [new_b -> new_b], where new_b is the new binder.
 248 -- This is important when doing eta expansion; e.g. matching (\x.M) ~ N
 249 -- In effect we switch to (\x'.M) ~ (\x'.N x'), where x' is new_b
 250 -- So we must add x' to the env of both L and R.  (x' is fresh, so it
 251 -- can't capture anything in N.)
 252 --
 253 -- If we don't do this, we can get silly matches like
 254 --      forall a.  \y.a  ~   v
 255 -- succeeding with [x -> v y], which is bogus of course
 256
 257 rnOccL, rnOccR :: RnEnv2 -> Var -> Var
 258 -- Look up the renaming of an occurrence in the left or right term
 259 rnOccL (RV2 { envL = env }) v = lookupVarEnv env v `orElse` v
 260 rnOccR (RV2 { envR = env }) v = lookupVarEnv env v `orElse` v
 261
 262 inRnEnvL, inRnEnvR :: RnEnv2 -> Var -> Bool
 263 -- Tells whether a variable is locally bound
 264 inRnEnvL (RV2 { envL = env }) v = v `elemVarEnv` env
 265 inRnEnvR (RV2 { envR = env }) v = v `elemVarEnv` env
 266
 267 lookupRnInScope :: RnEnv2 -> Var -> Var
 268 lookupRnInScope env v = lookupInScope (in_scope env) v `orElse` v
 269
 270 nukeRnEnvL, nukeRnEnvR :: RnEnv2 -> RnEnv2
 271 nukeRnEnvL env = env { envL = emptyVarEnv }
 272 nukeRnEnvR env = env { envR = emptyVarEnv }
 273 \end{code}
 274
 275
 276 %************************************************************************
 277 %*                                                                      *
 278                 Tidying
 279 %*                                                                      *
 280 %************************************************************************
 281
 282 When tidying up print names, we keep a mapping of in-scope occ-names
 283 (the TidyOccEnv) and a Var-to-Var of the current renamings.
 284
 285 \begin{code}
 286 type TidyEnv = (TidyOccEnv, VarEnv Var)
 287
 288 emptyTidyEnv :: TidyEnv
 289 emptyTidyEnv = (emptyTidyOccEnv, emptyVarEnv)
 290 \end{code}
 291
 292
 293 %************************************************************************
 294 %*                                                                      *
 295 \subsection{@VarEnv@s}
 296 %*                                                                      *
 297 %************************************************************************
 298
 299 \begin{code}
 300 type VarEnv elt   = UniqFM elt
 301 type IdEnv elt    = VarEnv elt
 302 type TyVarEnv elt = VarEnv elt
 303
 304 emptyVarEnv       :: VarEnv a
 305 mkVarEnv          :: [(Var, a)] -> VarEnv a
 306 zipVarEnv         :: [Var] -> [a] -> VarEnv a
 307 unitVarEnv        :: Var -> a -> VarEnv a
 308 extendVarEnv      :: VarEnv a -> Var -> a -> VarEnv a
 309 extendVarEnv_C    :: (a->a->a) -> VarEnv a -> Var -> a -> VarEnv a
 310 plusVarEnv        :: VarEnv a -> VarEnv a -> VarEnv a
 311 extendVarEnvList  :: VarEnv a -> [(Var, a)] -> VarEnv a
 312
 313 lookupVarEnv_Directly :: VarEnv a -> Unique -> Maybe a
 314 filterVarEnv_Directly :: (Unique -> a -> Bool) -> VarEnv a -> VarEnv a
 315 delVarEnvList     :: VarEnv a -> [Var] -> VarEnv a
 316 delVarEnv         :: VarEnv a -> Var -> VarEnv a
 317 plusVarEnv_C      :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
 318 mapVarEnv         :: (a -> b) -> VarEnv a -> VarEnv b
 319 modifyVarEnv      :: (a -> a) -> VarEnv a -> Var -> VarEnv a
 320 varEnvElts        :: VarEnv a -> [a]
 321 varEnvKeys        :: VarEnv a -> [Unique]
 322
 323 isEmptyVarEnv     :: VarEnv a -> Bool
 324 lookupVarEnv      :: VarEnv a -> Var -> Maybe a
 325 lookupVarEnv_NF   :: VarEnv a -> Var -> a
 326 lookupWithDefaultVarEnv :: VarEnv a -> a -> Var -> a
 327 elemVarEnv        :: Var -> VarEnv a -> Bool
 328 elemVarEnvByKey   :: Unique -> VarEnv a -> Bool
 329 foldVarEnv        :: (a -> b -> b) -> b -> VarEnv a -> b
 330 \end{code}
 331
 332 \begin{code}
 333 elemVarEnv       = elemUFM
 334 elemVarEnvByKey  = elemUFM_Directly
 335 extendVarEnv     = addToUFM
 336 extendVarEnv_C   = addToUFM_C
 337 extendVarEnvList = addListToUFM
 338 plusVarEnv_C     = plusUFM_C
 339 delVarEnvList    = delListFromUFM
 340 delVarEnv        = delFromUFM
 341 plusVarEnv       = plusUFM
 342 lookupVarEnv     = lookupUFM
 343 lookupWithDefaultVarEnv = lookupWithDefaultUFM
 344 mapVarEnv        = mapUFM
 345 mkVarEnv         = listToUFM
 346 emptyVarEnv      = emptyUFM
 347 varEnvElts       = eltsUFM
 348 varEnvKeys       = keysUFM
 349 unitVarEnv       = unitUFM
 350 isEmptyVarEnv    = isNullUFM
 351 foldVarEnv       = foldUFM
 352 lookupVarEnv_Directly = lookupUFM_Directly
 353 filterVarEnv_Directly = filterUFM_Directly
 354
 355 zipVarEnv tyvars tys   = mkVarEnv (zipEqual "zipVarEnv" tyvars tys)
 356 lookupVarEnv_NF env id = case (lookupVarEnv env id) of { Just xx -> xx }
 357 \end{code}
 358
 359 @modifyVarEnv@: Look up a thing in the VarEnv,
 360 then mash it with the modify function, and put it back.
 361
 362 \begin{code}
 363 modifyVarEnv mangle_fn env key
 364   = case (lookupVarEnv env key) of
 365       Nothing -> env
 366       Just xx -> extendVarEnv env key (mangle_fn xx)
 367
 368 modifyVarEnv_Directly mangle_fn env key
 369   = case (lookupUFM_Directly env key) of
 370       Nothing -> env
 371       Just xx -> addToUFM_Directly env key (mangle_fn xx)
 372 \end{code}