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