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