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