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