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, 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 StaticFlags
  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 = uniqAway in_scope bL
 235
 236 rnBndrR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
 237   = (RV2 { envL     = envL
 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 rnOccL, rnOccR :: RnEnv2 -> Var -> Var
 244 -- Look up the renaming of an occurrence in the left or right term
 245 rnOccL (RV2 { envL = env }) v = lookupVarEnv env v `orElse` v
 246 rnOccR (RV2 { envR = env }) v = lookupVarEnv env v `orElse` v
 247
 248 inRnEnvL, inRnEnvR :: RnEnv2 -> Var -> Bool
 249 -- Tells whether a variable is locally bound
 250 inRnEnvL (RV2 { envL = env }) v = v `elemVarEnv` env
 251 inRnEnvR (RV2 { envR = env }) v = v `elemVarEnv` env
 252
 253 lookupRnInScope :: RnEnv2 -> Var -> Var
 254 lookupRnInScope env v = lookupInScope (in_scope env) v `orElse` v
 255
 256 nukeRnEnvL, nukeRnEnvR :: RnEnv2 -> RnEnv2
 257 nukeRnEnvL env = env { envL = emptyVarEnv }
 258 nukeRnEnvR env = env { envR = emptyVarEnv }
 259 \end{code}
 260
 261
 262 %************************************************************************
 263 %*                                                                      *
 264                 Tidying
 265 %*                                                                      *
 266 %************************************************************************
 267
 268 When tidying up print names, we keep a mapping of in-scope occ-names
 269 (the TidyOccEnv) and a Var-to-Var of the current renamings.
 270
 271 \begin{code}
 272 type TidyEnv = (TidyOccEnv, VarEnv Var)
 273
 274 emptyTidyEnv :: TidyEnv
 275 emptyTidyEnv = (emptyTidyOccEnv, emptyVarEnv)
 276 \end{code}
 277
 278
 279 %************************************************************************
 280 %*                                                                      *
 281 \subsection{@VarEnv@s}
 282 %*                                                                      *
 283 %************************************************************************
 284
 285 \begin{code}
 286 type VarEnv elt   = UniqFM elt
 287 type IdEnv elt    = VarEnv elt
 288 type TyVarEnv elt = VarEnv elt
 289
 290 emptyVarEnv       :: VarEnv a
 291 mkVarEnv          :: [(Var, a)] -> VarEnv a
 292 zipVarEnv         :: [Var] -> [a] -> VarEnv a
 293 unitVarEnv        :: Var -> a -> VarEnv a
 294 extendVarEnv      :: VarEnv a -> Var -> a -> VarEnv a
 295 extendVarEnv_C    :: (a->a->a) -> VarEnv a -> Var -> a -> VarEnv a
 296 plusVarEnv        :: VarEnv a -> VarEnv a -> VarEnv a
 297 extendVarEnvList  :: VarEnv a -> [(Var, a)] -> VarEnv a
 298
 299 lookupVarEnv_Directly :: VarEnv a -> Unique -> Maybe a
 300 filterVarEnv_Directly :: (Unique -> a -> Bool) -> VarEnv a -> VarEnv a
 301 delVarEnvList     :: VarEnv a -> [Var] -> VarEnv a
 302 delVarEnv         :: VarEnv a -> Var -> VarEnv a
 303 plusVarEnv_C      :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
 304 mapVarEnv         :: (a -> b) -> VarEnv a -> VarEnv b
 305 modifyVarEnv      :: (a -> a) -> VarEnv a -> Var -> VarEnv a
 306 varEnvElts        :: VarEnv a -> [a]
 307 varEnvKeys        :: VarEnv a -> [Unique]
 308
 309 isEmptyVarEnv     :: VarEnv a -> Bool
 310 lookupVarEnv      :: VarEnv a -> Var -> Maybe a
 311 lookupVarEnv_NF   :: VarEnv a -> Var -> a
 312 lookupWithDefaultVarEnv :: VarEnv a -> a -> Var -> a
 313 elemVarEnv        :: Var -> VarEnv a -> Bool
 314 elemVarEnvByKey   :: Unique -> VarEnv a -> Bool
 315 foldVarEnv        :: (a -> b -> b) -> b -> VarEnv a -> b
 316 \end{code}
 317
 318 \begin{code}
 319 elemVarEnv       = elemUFM
 320 elemVarEnvByKey  = elemUFM_Directly
 321 extendVarEnv     = addToUFM
 322 extendVarEnv_C   = addToUFM_C
 323 extendVarEnvList = addListToUFM
 324 plusVarEnv_C     = plusUFM_C
 325 delVarEnvList    = delListFromUFM
 326 delVarEnv        = delFromUFM
 327 plusVarEnv       = plusUFM
 328 lookupVarEnv     = lookupUFM
 329 lookupWithDefaultVarEnv = lookupWithDefaultUFM
 330 mapVarEnv        = mapUFM
 331 mkVarEnv         = listToUFM
 332 emptyVarEnv      = emptyUFM
 333 varEnvElts       = eltsUFM
 334 varEnvKeys       = keysUFM
 335 unitVarEnv       = unitUFM
 336 isEmptyVarEnv    = isNullUFM
 337 foldVarEnv       = foldUFM
 338 lookupVarEnv_Directly = lookupUFM_Directly
 339 filterVarEnv_Directly = filterUFM_Directly
 340
 341 zipVarEnv tyvars tys   = mkVarEnv (zipEqual "zipVarEnv" tyvars tys)
 342 lookupVarEnv_NF env id = case (lookupVarEnv env id) of { Just xx -> xx }
 343 \end{code}
 344
 345 @modifyVarEnv@: Look up a thing in the VarEnv,
 346 then mash it with the modify function, and put it back.
 347
 348 \begin{code}
 349 modifyVarEnv mangle_fn env key
 350   = case (lookupVarEnv env key) of
 351       Nothing -> env
 352       Just xx -> extendVarEnv env key (mangle_fn xx)
 353
 354 modifyVarEnv_Directly mangle_fn env key
 355   = case (lookupUFM_Directly env key) of
 356       Nothing -> env
 357       Just xx -> addToUFM_Directly env key (mangle_fn xx)
 358 \end{code}