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