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