Improve HsSyn pretty printing

[ghc-hetmet.git] / compiler / basicTypes / VarEnv.lhs

diff --git a/compiler/basicTypes/VarEnv.lhs b/compiler/basicTypes/VarEnv.lhs
index 7e28d1a..bf3f96d 100644 (file)
--- a/compiler/basicTypes/VarEnv.lhs
+++ b/compiler/basicTypes/VarEnv.lhs
@@ -11,7 +11,7 @@ module VarEnv (
        -- ** Manipulating these environments
        emptyVarEnv, unitVarEnv, mkVarEnv,
        elemVarEnv, varEnvElts, varEnvKeys,
-       extendVarEnv, extendVarEnv_C, extendVarEnvList,
+       extendVarEnv, extendVarEnv_C, extendVarEnv_Acc, extendVarEnvList,
        plusVarEnv, plusVarEnv_C,
        delVarEnvList, delVarEnv,
        lookupVarEnv, lookupVarEnv_NF, lookupWithDefaultVarEnv,
@@ -19,7 +19,7 @@ module VarEnv (
        modifyVarEnv, modifyVarEnv_Directly,
        isEmptyVarEnv, foldVarEnv, 
        elemVarEnvByKey, lookupVarEnv_Directly,
-       filterVarEnv_Directly,
+       filterVarEnv_Directly, restrictVarEnv,
 
        -- * The InScopeSet type
        InScopeSet, 
@@ -27,7 +27,8 @@ module VarEnv (
        -- ** Operations on InScopeSets
        emptyInScopeSet, mkInScopeSet, delInScopeSet,
        extendInScopeSet, extendInScopeSetList, extendInScopeSetSet, 
-       getInScopeVars, lookupInScope, elemInScopeSet, uniqAway, 
+       getInScopeVars, lookupInScope, lookupInScope_Directly, 
+        elemInScopeSet, uniqAway, 
 
        -- * The RnEnv2 type
        RnEnv2, 
@@ -114,6 +115,10 @@ elemInScopeSet v (InScope in_scope _) = v `elemVarEnv` in_scope
 -- the variable's identity (unique) to its full value.
 lookupInScope :: InScopeSet -> Var -> Maybe Var
 lookupInScope (InScope in_scope _) v  = lookupVarEnv in_scope v
+
+lookupInScope_Directly :: InScopeSet -> Unique -> Maybe Var
+lookupInScope_Directly (InScope in_scope _) uniq
+  = lookupVarEnv_Directly in_scope uniq
 \end{code}
 
 \begin{code}
@@ -316,11 +321,13 @@ zipVarEnv   :: [Var] -> [a] -> VarEnv a
 unitVarEnv       :: Var -> a -> VarEnv a
 extendVarEnv     :: VarEnv a -> Var -> a -> VarEnv a
 extendVarEnv_C   :: (a->a->a) -> VarEnv a -> Var -> a -> VarEnv a
+extendVarEnv_Acc  :: (a->b->b) -> (a->b) -> VarEnv b -> Var -> a -> VarEnv b
 plusVarEnv       :: VarEnv a -> VarEnv a -> VarEnv a
 extendVarEnvList  :: VarEnv a -> [(Var, a)] -> VarEnv a
                  
 lookupVarEnv_Directly :: VarEnv a -> Unique -> Maybe a
 filterVarEnv_Directly :: (Unique -> a -> Bool) -> VarEnv a -> VarEnv a
+restrictVarEnv    :: VarEnv a -> VarSet -> VarEnv a
 delVarEnvList     :: VarEnv a -> [Var] -> VarEnv a
 delVarEnv        :: VarEnv a -> Var -> VarEnv a
 plusVarEnv_C     :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
@@ -343,6 +350,7 @@ elemVarEnv       = elemUFM
 elemVarEnvByKey  = elemUFM_Directly
 extendVarEnv    = addToUFM
 extendVarEnv_C  = addToUFM_C
+extendVarEnv_Acc = addToUFM_Acc
 extendVarEnvList = addListToUFM
 plusVarEnv_C    = plusUFM_C
 delVarEnvList   = delListFromUFM
@@ -361,6 +369,10 @@ foldVarEnv  = foldUFM
 lookupVarEnv_Directly = lookupUFM_Directly
 filterVarEnv_Directly = filterUFM_Directly
 
+restrictVarEnv env vs = filterVarEnv_Directly keep env
+  where
+    keep u _ = u `elemVarSetByKey` vs
+    
 zipVarEnv tyvars tys   = mkVarEnv (zipEqual "zipVarEnv" tyvars tys)
 lookupVarEnv_NF env id = case lookupVarEnv env id of
                          Just xx -> xx