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         minusVarEnv, intersectsVarEnv,
  18         lookupVarEnv, lookupVarEnv_NF, lookupWithDefaultVarEnv,
  19         mapVarEnv, zipVarEnv,
  20         modifyVarEnv, modifyVarEnv_Directly,
  21         isEmptyVarEnv, foldVarEnv,
  22         elemVarEnvByKey, lookupVarEnv_Directly,
  23         filterVarEnv_Directly, restrictVarEnv,
  24
  25         -- * The InScopeSet type
  26         InScopeSet,
  27
  28         -- ** Operations on InScopeSets
  29         emptyInScopeSet, mkInScopeSet, delInScopeSet,
  30         extendInScopeSet, extendInScopeSetList, extendInScopeSetSet,
  31         getInScopeVars, lookupInScope, lookupInScope_Directly,
  32         unionInScope, elemInScopeSet, uniqAway,
  33
  34         -- * The RnEnv2 type
  35         RnEnv2,
  36
  37         -- ** Operations on RnEnv2s
  38         mkRnEnv2, rnBndr2, rnBndrs2, rnOccL, rnOccR, inRnEnvL, inRnEnvR,
  39         rnBndrL, rnBndrR, nukeRnEnvL, nukeRnEnvR, extendRnInScopeList,
  40         rnEtaL, rnEtaR,
  41         rnInScope, rnInScopeSet, lookupRnInScope,
  42
  43         -- * TidyEnv and its operation
  44         TidyEnv,
  45         emptyTidyEnv
  46     ) where
  47
  48 import OccName
  49 import Var
  50 import VarSet
  51 import UniqFM
  52 import Unique
  53 import Util
  54 import Maybes
  55 import Outputable
  56 import FastTypes
  57 import StaticFlags
  58 import FastString
  59 \end{code}
  60
  61
  62 %************************************************************************
  63 %*                                                                      *
  64                 In-scope sets
  65 %*                                                                      *
  66 %************************************************************************
  67
  68 \begin{code}
  69 -- | A set of variables that are in scope at some point
  70 data InScopeSet = InScope (VarEnv Var) FastInt
  71         -- The (VarEnv Var) is just a VarSet.  But we write it like
  72         -- this to remind ourselves that you can look up a Var in
  73         -- the InScopeSet. Typically the InScopeSet contains the
  74         -- canonical version of the variable (e.g. with an informative
  75         -- unfolding), so this lookup is useful.
  76         --
  77         -- INVARIANT: the VarEnv maps (the Unique of) a variable to
  78         --            a variable with the same Uniqua.  (This was not
  79         --            the case in the past, when we had a grevious hack
  80         --            mapping var1 to var2.
  81         --
  82         -- The FastInt is a kind of hash-value used by uniqAway
  83         -- For example, it might be the size of the set
  84         -- INVARIANT: it's not zero; we use it as a multiplier in uniqAway
  85
  86 instance Outputable InScopeSet where
  87   ppr (InScope s _) = ptext (sLit "InScope") <+> ppr s
  88
  89 emptyInScopeSet :: InScopeSet
  90 emptyInScopeSet = InScope emptyVarSet (_ILIT(1))
  91
  92 getInScopeVars ::  InScopeSet -> VarEnv Var
  93 getInScopeVars (InScope vs _) = vs
  94
  95 mkInScopeSet :: VarEnv Var -> InScopeSet
  96 mkInScopeSet in_scope = InScope in_scope (_ILIT(1))
  97
  98 extendInScopeSet :: InScopeSet -> Var -> InScopeSet
  99 extendInScopeSet (InScope in_scope n) v = InScope (extendVarEnv in_scope v v) (n +# _ILIT(1))
 100
 101 extendInScopeSetList :: InScopeSet -> [Var] -> InScopeSet
 102 extendInScopeSetList (InScope in_scope n) vs
 103    = InScope (foldl (\s v -> extendVarEnv s v v) in_scope vs)
 104                     (n +# iUnbox (length vs))
 105
 106 extendInScopeSetSet :: InScopeSet -> VarEnv Var -> InScopeSet
 107 extendInScopeSetSet (InScope in_scope n) vs
 108    = InScope (in_scope `plusVarEnv` vs) (n +# iUnbox (sizeUFM vs))
 109
 110 delInScopeSet :: InScopeSet -> Var -> InScopeSet
 111 delInScopeSet (InScope in_scope n) v = InScope (in_scope `delVarEnv` v) n
 112
 113 elemInScopeSet :: Var -> InScopeSet -> Bool
 114 elemInScopeSet v (InScope in_scope _) = v `elemVarEnv` in_scope
 115
 116 -- | Look up a variable the 'InScopeSet'.  This lets you map from
 117 -- the variable's identity (unique) to its full value.
 118 lookupInScope :: InScopeSet -> Var -> Maybe Var
 119 lookupInScope (InScope in_scope _) v  = lookupVarEnv in_scope v
 120
 121 lookupInScope_Directly :: InScopeSet -> Unique -> Maybe Var
 122 lookupInScope_Directly (InScope in_scope _) uniq
 123   = lookupVarEnv_Directly in_scope uniq
 124
 125 unionInScope :: InScopeSet -> InScopeSet -> InScopeSet
 126 unionInScope (InScope s1 _) (InScope s2 n2)
 127   = InScope (s1 `plusVarEnv` s2) n2
 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.
 247 rnBndrL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
 248   = (RV2 { envL     = extendVarEnv envL bL new_b
 249          , envR     = envR
 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.
 257 rnBndrR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
 258   = (RV2 { envR     = extendVarEnv envR bR new_b
 259          , envL     = envL
 260          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 261   where
 262     new_b = uniqAway in_scope bR
 263
 264 rnEtaL :: RnEnv2 -> Var -> (RnEnv2, Var)
 265 -- ^ Similar to 'rnBndrL' but used for eta expansion
 266 -- See Note [Eta expansion]
 267 rnEtaL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
 268   = (RV2 { envL     = extendVarEnv envL bL new_b
 269          , envR     = extendVarEnv envR new_b new_b     -- Note [Eta expansion]
 270          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 271   where
 272     new_b = uniqAway in_scope bL
 273
 274 rnEtaR :: RnEnv2 -> Var -> (RnEnv2, Var)
 275 -- ^ Similar to 'rnBndr2' but used for eta expansion
 276 -- See Note [Eta expansion]
 277 rnEtaR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
 278   = (RV2 { envL     = extendVarEnv envL new_b new_b     -- Note [Eta expansion]
 279          , envR     = extendVarEnv envR bR new_b
 280          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 281   where
 282     new_b = uniqAway in_scope bR
 283
 284 rnOccL, rnOccR :: RnEnv2 -> Var -> Var
 285 -- ^ Look up the renaming of an occurrence in the left or right term
 286 rnOccL (RV2 { envL = env }) v = lookupVarEnv env v `orElse` v
 287 rnOccR (RV2 { envR = env }) v = lookupVarEnv env v `orElse` v
 288
 289 inRnEnvL, inRnEnvR :: RnEnv2 -> Var -> Bool
 290 -- ^ Tells whether a variable is locally bound
 291 inRnEnvL (RV2 { envL = env }) v = v `elemVarEnv` env
 292 inRnEnvR (RV2 { envR = env }) v = v `elemVarEnv` env
 293
 294 lookupRnInScope :: RnEnv2 -> Var -> Var
 295 lookupRnInScope env v = lookupInScope (in_scope env) v `orElse` v
 296
 297 nukeRnEnvL, nukeRnEnvR :: RnEnv2 -> RnEnv2
 298 -- ^ Wipe the left or right side renaming
 299 nukeRnEnvL env = env { envL = emptyVarEnv }
 300 nukeRnEnvR env = env { envR = emptyVarEnv }
 301 \end{code}
 302
 303 Note [Eta expansion]
 304 ~~~~~~~~~~~~~~~~~~~~
 305 When matching
 306      (\x.M) ~ N
 307 we rename x to x' with, where x' is not in scope in
 308 either term.  Then we want to behave as if we'd seen
 309      (\x'.M) ~ (\x'.N x')
 310 Since x' isn't in scope in N, the form (\x'. N x') doesn't
 311 capture any variables in N.  But we must nevertheless extend
 312 the envR with a binding [x' -> x'], to support the occurs check.
 313 For example, if we don't do this, we can get silly matches like
 314         forall a.  (\y.a)  ~   v
 315 succeeding with [a -> v y], which is bogus of course.
 316
 317
 318 %************************************************************************
 319 %*                                                                      *
 320                 Tidying
 321 %*                                                                      *
 322 %************************************************************************
 323
 324 \begin{code}
 325 -- | When tidying up print names, we keep a mapping of in-scope occ-names
 326 -- (the 'TidyOccEnv') and a Var-to-Var of the current renamings
 327 type TidyEnv = (TidyOccEnv, VarEnv Var)
 328
 329 emptyTidyEnv :: TidyEnv
 330 emptyTidyEnv = (emptyTidyOccEnv, emptyVarEnv)
 331 \end{code}
 332
 333
 334 %************************************************************************
 335 %*                                                                      *
 336 \subsection{@VarEnv@s}
 337 %*                                                                      *
 338 %************************************************************************
 339
 340 \begin{code}
 341 type VarEnv elt   = UniqFM elt
 342 type IdEnv elt    = VarEnv elt
 343 type TyVarEnv elt = VarEnv elt
 344
 345 emptyVarEnv       :: VarEnv a
 346 mkVarEnv          :: [(Var, a)] -> VarEnv a
 347 zipVarEnv         :: [Var] -> [a] -> VarEnv a
 348 unitVarEnv        :: Var -> a -> VarEnv a
 349 extendVarEnv      :: VarEnv a -> Var -> a -> VarEnv a
 350 extendVarEnv_C    :: (a->a->a) -> VarEnv a -> Var -> a -> VarEnv a
 351 extendVarEnv_Acc  :: (a->b->b) -> (a->b) -> VarEnv b -> Var -> a -> VarEnv b
 352 plusVarEnv        :: VarEnv a -> VarEnv a -> VarEnv a
 353 extendVarEnvList  :: VarEnv a -> [(Var, a)] -> VarEnv a
 354
 355 lookupVarEnv_Directly :: VarEnv a -> Unique -> Maybe a
 356 filterVarEnv_Directly :: (Unique -> a -> Bool) -> VarEnv a -> VarEnv a
 357 restrictVarEnv    :: VarEnv a -> VarSet -> VarEnv a
 358 delVarEnvList     :: VarEnv a -> [Var] -> VarEnv a
 359 delVarEnv         :: VarEnv a -> Var -> VarEnv a
 360 minusVarEnv       :: VarEnv a -> VarEnv a -> VarEnv a
 361 intersectsVarEnv  :: VarEnv a -> VarEnv a -> Bool
 362 plusVarEnv_C      :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
 363 mapVarEnv         :: (a -> b) -> VarEnv a -> VarEnv b
 364 modifyVarEnv      :: (a -> a) -> VarEnv a -> Var -> VarEnv a
 365 varEnvElts        :: VarEnv a -> [a]
 366 varEnvKeys        :: VarEnv a -> [Unique]
 367
 368 isEmptyVarEnv     :: VarEnv a -> Bool
 369 lookupVarEnv      :: VarEnv a -> Var -> Maybe a
 370 lookupVarEnv_NF   :: VarEnv a -> Var -> a
 371 lookupWithDefaultVarEnv :: VarEnv a -> a -> Var -> a
 372 elemVarEnv        :: Var -> VarEnv a -> Bool
 373 elemVarEnvByKey   :: Unique -> VarEnv a -> Bool
 374 foldVarEnv        :: (a -> b -> b) -> b -> VarEnv a -> b
 375 \end{code}
 376
 377 \begin{code}
 378 elemVarEnv       = elemUFM
 379 elemVarEnvByKey  = elemUFM_Directly
 380 extendVarEnv     = addToUFM
 381 extendVarEnv_C   = addToUFM_C
 382 extendVarEnv_Acc = addToUFM_Acc
 383 extendVarEnvList = addListToUFM
 384 plusVarEnv_C     = plusUFM_C
 385 delVarEnvList    = delListFromUFM
 386 delVarEnv        = delFromUFM
 387 minusVarEnv      = minusUFM
 388 intersectsVarEnv e1 e2 = not (isEmptyVarEnv (e1 `intersectUFM` e2))
 389 plusVarEnv       = plusUFM
 390 lookupVarEnv     = lookupUFM
 391 lookupWithDefaultVarEnv = lookupWithDefaultUFM
 392 mapVarEnv        = mapUFM
 393 mkVarEnv         = listToUFM
 394 emptyVarEnv      = emptyUFM
 395 varEnvElts       = eltsUFM
 396 varEnvKeys       = keysUFM
 397 unitVarEnv       = unitUFM
 398 isEmptyVarEnv    = isNullUFM
 399 foldVarEnv       = foldUFM
 400 lookupVarEnv_Directly = lookupUFM_Directly
 401 filterVarEnv_Directly = filterUFM_Directly
 402
 403 restrictVarEnv env vs = filterVarEnv_Directly keep env
 404   where
 405     keep u _ = u `elemVarSetByKey` vs
 406
 407 zipVarEnv tyvars tys   = mkVarEnv (zipEqual "zipVarEnv" tyvars tys)
 408 lookupVarEnv_NF env id = case lookupVarEnv env id of
 409                          Just xx -> xx
 410                          Nothing -> panic "lookupVarEnv_NF: Nothing"
 411 \end{code}
 412
 413 @modifyVarEnv@: Look up a thing in the VarEnv,
 414 then mash it with the modify function, and put it back.
 415
 416 \begin{code}
 417 modifyVarEnv mangle_fn env key
 418   = case (lookupVarEnv env key) of
 419       Nothing -> env
 420       Just xx -> extendVarEnv env key (mangle_fn xx)
 421
 422 modifyVarEnv_Directly :: (a -> a) -> UniqFM a -> Unique -> UniqFM a
 423 modifyVarEnv_Directly mangle_fn env key
 424   = case (lookupUFM_Directly env key) of
 425       Nothing -> env
 426       Just xx -> addToUFM_Directly env key (mangle_fn xx)
 427 \end{code}