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