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