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