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