compiler/basicTypes/VarEnv.lhs

   1
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section{@VarEnvs@: Variable environments}
   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, lookupRnInScope,
  32
  33         -- TidyEnvs
  34         TidyEnv, emptyTidyEnv
  35     ) where
  36
  37 #include "HsVersions.h"
  38
  39 import OccName    ( TidyOccEnv, emptyTidyOccEnv )
  40 import Var        ( Var, setVarUnique )
  41 import VarSet
  42 import UniqFM
  43 import Unique     ( Unique, deriveUnique, getUnique )
  44 import Util       ( zipEqual, foldl2 )
  45 import Maybes     ( orElse )
  46 import StaticFlags( opt_PprStyle_Debug )
  47 import Outputable
  48 import FastTypes
  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 i) = ptext SLIT("InScope") <+> ppr s
  66
  67 emptyInScopeSet :: InScopeSet
  68 emptyInScopeSet = InScope emptyVarSet 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 1#
  75
  76 extendInScopeSet :: InScopeSet -> Var -> InScopeSet
  77 extendInScopeSet (InScope in_scope n) v = InScope (extendVarEnv in_scope v v) (n +# 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 +# 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 n) = 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 n) 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 1#
 131   where
 132     orig_unique = getUnique var
 133     try k
 134 #ifdef DEBUG
 135           | k ># 1000#
 136           = pprPanic "uniqAway loop:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n))
 137 #endif
 138           | uniq `elemVarSetByKey` set = try (k +# 1#)
 139 #ifdef DEBUG
 140           | opt_PprStyle_Debug && k ># 3#
 141           = pprTrace "uniqAway:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n))
 142             setVarUnique var uniq
 143 #endif
 144           | otherwise = setVarUnique var uniq
 145           where
 146             uniq = deriveUnique orig_unique (iBox (n *# k))
 147 \end{code}
 148
 149
 150 %************************************************************************
 151 %*                                                                      *
 152                 Dual renaming
 153 %*                                                                      *
 154 %************************************************************************
 155
 156 When we are comparing (or matching) types or terms, we are faced with
 157 "going under" corresponding binders.  E.g. when comparing
 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         a) a renaming for the left-hand expression
 165         b) a renaming for the right-hand expressions
 166         c) an in-scope set
 167
 168 Furthermore, when matching, we want to be able to have an 'occurs check',
 169 to prevent
 170         \x. f   ~   \y. y
 171 matching with f->y.  So for each expression we want to know that set of
 172 locally-bound variables. That is precisely the domain of the mappings (a)
 173 and (b), but we must ensure that we always extend the mappings as we go in.
 174
 175
 176 \begin{code}
 177 data RnEnv2
 178   = RV2 { envL     :: VarEnv Var        -- Renaming for Left term
 179         , envR     :: VarEnv Var        -- Renaming for Right term
 180         , in_scope :: InScopeSet }      -- In scope in left or right terms
 181
 182 -- The renamings envL and envR are *guaranteed* to contain a binding
 183 -- for every variable bound as we go into the term, even if it is not
 184 -- renamed.  That way we can ask what variables are locally bound
 185 -- (inRnEnvL, inRnEnvR)
 186
 187 mkRnEnv2 :: InScopeSet -> RnEnv2
 188 mkRnEnv2 vars = RV2     { envL     = emptyVarEnv
 189                         , envR     = emptyVarEnv
 190                         , in_scope = vars }
 191
 192 extendRnInScopeList :: RnEnv2 -> [Var] -> RnEnv2
 193 extendRnInScopeList env vs
 194   = env { in_scope = extendInScopeSetList (in_scope env) vs }
 195
 196 rnInScope :: Var -> RnEnv2 -> Bool
 197 rnInScope x env = x `elemInScopeSet` in_scope env
 198
 199 rnBndrs2 :: RnEnv2 -> [Var] -> [Var] -> RnEnv2
 200 -- Arg lists must be of equal length
 201 rnBndrs2 env bsL bsR = foldl2 rnBndr2 env bsL bsR
 202
 203 rnBndr2 :: RnEnv2 -> Var -> Var -> RnEnv2
 204 -- (rnBndr2 env bL bR) go under a binder bL in the Left term 1,
 205 --                     and binder bR in the Right term
 206 -- It finds a new binder, new_b,
 207 -- and returns an environment mapping bL->new_b and bR->new_b resp.
 208 rnBndr2 (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL bR
 209   = RV2 { envL     = extendVarEnv envL bL new_b   -- See Note
 210         , envR     = extendVarEnv envR bR new_b   -- [Rebinding]
 211         , in_scope = extendInScopeSet in_scope new_b }
 212   where
 213         -- Find a new binder not in scope in either term
 214     new_b | not (bL `elemInScopeSet` in_scope) = bL
 215           | not (bR `elemInScopeSet` in_scope) = bR
 216           | otherwise                          = uniqAway' in_scope bL
 217
 218         -- Note [Rebinding]
 219         -- If the new var is the same as the old one, note that
 220         -- the extendVarEnv *deletes* any current renaming
 221         -- E.g.   (\x. \x. ...)  ~  (\y. \z. ...)
 222         --
 223         --   Inside \x  \y      { [x->y], [y->y],       {y} }
 224         --       \x  \z         { [x->x], [y->y, z->x], {y,x} }
 225
 226 rnBndrL, rnBndrR :: RnEnv2 -> Var -> (RnEnv2, Var)
 227 -- Used when there's a binder on one side or the other only
 228 -- Useful when eta-expanding
 229 rnBndrL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
 230   = (RV2 { envL     = extendVarEnv envL bL new_b
 231          , envR     = envR
 232          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 233   where
 234     new_b | not (bL `elemInScopeSet` in_scope) = bL
 235           | otherwise                          = uniqAway' in_scope bL
 236
 237 rnBndrR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
 238   = (RV2 { envL     = envL
 239          , envR     = extendVarEnv envR bR new_b
 240          , in_scope = extendInScopeSet in_scope new_b }, new_b)
 241   where
 242     new_b | not (bR `elemInScopeSet` in_scope) = bR
 243           | otherwise                          = uniqAway' in_scope bR
 244
 245 rnOccL, rnOccR :: RnEnv2 -> Var -> Var
 246 -- Look up the renaming of an occurrence in the left or right term
 247 rnOccL (RV2 { envL = env }) v = lookupVarEnv env v `orElse` v
 248 rnOccR (RV2 { envR = env }) v = lookupVarEnv env v `orElse` v
 249
 250 inRnEnvL, inRnEnvR :: RnEnv2 -> Var -> Bool
 251 -- Tells whether a variable is locally bound
 252 inRnEnvL (RV2 { envL = env }) v = v `elemVarEnv` env
 253 inRnEnvR (RV2 { envR = env }) v = v `elemVarEnv` env
 254
 255 lookupRnInScope :: RnEnv2 -> Var -> Var
 256 lookupRnInScope env v = lookupInScope (in_scope env) v `orElse` v
 257
 258 nukeRnEnvL, nukeRnEnvR :: RnEnv2 -> RnEnv2
 259 nukeRnEnvL env = env { envL = emptyVarEnv }
 260 nukeRnEnvR env = env { envR = emptyVarEnv }
 261 \end{code}
 262
 263
 264 %************************************************************************
 265 %*                                                                      *
 266                 Tidying
 267 %*                                                                      *
 268 %************************************************************************
 269
 270 When tidying up print names, we keep a mapping of in-scope occ-names
 271 (the TidyOccEnv) and a Var-to-Var of the current renamings.
 272
 273 \begin{code}
 274 type TidyEnv = (TidyOccEnv, VarEnv Var)
 275
 276 emptyTidyEnv :: TidyEnv
 277 emptyTidyEnv = (emptyTidyOccEnv, emptyVarEnv)
 278 \end{code}
 279
 280
 281 %************************************************************************
 282 %*                                                                      *
 283 \subsection{@VarEnv@s}
 284 %*                                                                      *
 285 %************************************************************************
 286
 287 \begin{code}
 288 type VarEnv elt   = UniqFM elt
 289 type IdEnv elt    = VarEnv elt
 290 type TyVarEnv elt = VarEnv elt
 291
 292 emptyVarEnv       :: VarEnv a
 293 mkVarEnv          :: [(Var, a)] -> VarEnv a
 294 zipVarEnv         :: [Var] -> [a] -> VarEnv a
 295 unitVarEnv        :: Var -> a -> VarEnv a
 296 extendVarEnv      :: VarEnv a -> Var -> a -> VarEnv a
 297 extendVarEnv_C    :: (a->a->a) -> VarEnv a -> Var -> a -> VarEnv a
 298 plusVarEnv        :: VarEnv a -> VarEnv a -> VarEnv a
 299 extendVarEnvList  :: VarEnv a -> [(Var, a)] -> VarEnv a
 300
 301 lookupVarEnv_Directly :: VarEnv a -> Unique -> Maybe a
 302 filterVarEnv_Directly :: (Unique -> a -> Bool) -> VarEnv a -> VarEnv a
 303 delVarEnvList     :: VarEnv a -> [Var] -> VarEnv a
 304 delVarEnv         :: VarEnv a -> Var -> VarEnv a
 305 plusVarEnv_C      :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
 306 mapVarEnv         :: (a -> b) -> VarEnv a -> VarEnv b
 307 modifyVarEnv      :: (a -> a) -> VarEnv a -> Var -> VarEnv a
 308 varEnvElts        :: VarEnv a -> [a]
 309 varEnvKeys        :: VarEnv a -> [Unique]
 310
 311 isEmptyVarEnv     :: VarEnv a -> Bool
 312 lookupVarEnv      :: VarEnv a -> Var -> Maybe a
 313 lookupVarEnv_NF   :: VarEnv a -> Var -> a
 314 lookupWithDefaultVarEnv :: VarEnv a -> a -> Var -> a
 315 elemVarEnv        :: Var -> VarEnv a -> Bool
 316 elemVarEnvByKey   :: Unique -> VarEnv a -> Bool
 317 foldVarEnv        :: (a -> b -> b) -> b -> VarEnv a -> b
 318 \end{code}
 319
 320 \begin{code}
 321 elemVarEnv       = elemUFM
 322 elemVarEnvByKey  = elemUFM_Directly
 323 extendVarEnv     = addToUFM
 324 extendVarEnv_C   = addToUFM_C
 325 extendVarEnvList = addListToUFM
 326 plusVarEnv_C     = plusUFM_C
 327 delVarEnvList    = delListFromUFM
 328 delVarEnv        = delFromUFM
 329 plusVarEnv       = plusUFM
 330 lookupVarEnv     = lookupUFM
 331 lookupWithDefaultVarEnv = lookupWithDefaultUFM
 332 mapVarEnv        = mapUFM
 333 mkVarEnv         = listToUFM
 334 emptyVarEnv      = emptyUFM
 335 varEnvElts       = eltsUFM
 336 varEnvKeys       = keysUFM
 337 unitVarEnv       = unitUFM
 338 isEmptyVarEnv    = isNullUFM
 339 foldVarEnv       = foldUFM
 340 lookupVarEnv_Directly = lookupUFM_Directly
 341 filterVarEnv_Directly = filterUFM_Directly
 342
 343 zipVarEnv tyvars tys   = mkVarEnv (zipEqual "zipVarEnv" tyvars tys)
 344 lookupVarEnv_NF env id = case (lookupVarEnv env id) of { Just xx -> xx }
 345 \end{code}
 346
 347 @modifyVarEnv@: Look up a thing in the VarEnv,
 348 then mash it with the modify function, and put it back.
 349
 350 \begin{code}
 351 modifyVarEnv mangle_fn env key
 352   = case (lookupVarEnv env key) of
 353       Nothing -> env
 354       Just xx -> extendVarEnv env key (mangle_fn xx)
 355
 356 modifyVarEnv_Directly mangle_fn env key
 357   = case (lookupUFM_Directly env key) of
 358       Nothing -> env
 359       Just xx -> addToUFM_Directly env key (mangle_fn xx)
 360 \end{code}