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