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