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