Add "alter" functions to UniqFM, VarEnv, NameEnv

I need these for a trie data structure I'm working on

3 files changed, 15 insertions, 4 deletions

diff --git a/compiler/basicTypes/NameEnv.lhs b/compiler/basicTypes/NameEnv.lhs
index 984f0963cc..0dc5c32c7a 100644
--- a/compiler/basicTypes/NameEnv.lhs
+++ b/compiler/basicTypes/NameEnv.lhs
@@ -15,7 +15,7 @@ module NameEnv (
 	extendNameEnv_C, extendNameEnv_Acc, extendNameEnv,
         extendNameEnvList, extendNameEnvList_C,
 	foldNameEnv, filterNameEnv,
-	plusNameEnv, plusNameEnv_C, 
+	plusNameEnv, plusNameEnv_C, alterNameEnv,
 	lookupNameEnv, lookupNameEnv_NF, delFromNameEnv, delListFromNameEnv,
 	elemNameEnv, mapNameEnv
     ) where
@@ -41,6 +41,7 @@ emptyNameEnv   	   :: NameEnv a
 mkNameEnv	   :: [(Name,a)] -> NameEnv a
 nameEnvElts    	   :: NameEnv a -> [a]
 nameEnvUniqueElts  :: NameEnv a -> [(Unique, a)]
+alterNameEnv       :: (Maybe a-> Maybe a) -> NameEnv a -> Name -> NameEnv a
 extendNameEnv_C    :: (a->a->a) -> NameEnv a -> Name -> a -> NameEnv a
 extendNameEnv_Acc  :: (a->b->b) -> (a->b) -> NameEnv b -> Name -> a -> NameEnv b
 extendNameEnv  	   :: NameEnv a -> Name -> a -> NameEnv a
@@ -64,6 +65,7 @@ unitNameEnv x y       = unitUFM x y
 extendNameEnv x y z   = addToUFM x y z
 extendNameEnvList x l = addListToUFM x l
 lookupNameEnv x y     = lookupUFM x y
+alterNameEnv          = alterUFM
 mkNameEnv     l       = listToUFM l
 elemNameEnv x y 	 = elemUFM x y
 foldNameEnv a b c	 = foldUFM a b c 
diff --git a/compiler/basicTypes/VarEnv.lhs b/compiler/basicTypes/VarEnv.lhs
index 07fabb0345..4ae0ff4307 100644
--- a/compiler/basicTypes/VarEnv.lhs
+++ b/compiler/basicTypes/VarEnv.lhs
@@ -12,7 +12,7 @@ module VarEnv (
 	emptyVarEnv, unitVarEnv, mkVarEnv,
 	elemVarEnv, varEnvElts, varEnvKeys,
 	extendVarEnv, extendVarEnv_C, extendVarEnv_Acc, extendVarEnvList,
-	plusVarEnv, plusVarEnv_C,
+	plusVarEnv, plusVarEnv_C, alterVarEnv,
 	delVarEnvList, delVarEnv,
         minusVarEnv, intersectsVarEnv,
 	lookupVarEnv, lookupVarEnv_NF, lookupWithDefaultVarEnv,
@@ -364,6 +364,7 @@ emptyVarEnv	  :: VarEnv a
 mkVarEnv	  :: [(Var, a)] -> VarEnv a
 zipVarEnv	  :: [Var] -> [a] -> VarEnv a
 unitVarEnv	  :: Var -> a -> VarEnv a
+alterVarEnv	  :: (Maybe a -> Maybe a) -> VarEnv a -> Var -> 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
@@ -395,6 +396,7 @@ foldVarEnv	  :: (a -> b -> b) -> b -> VarEnv a -> b
 \begin{code}
 elemVarEnv       = elemUFM
 elemVarEnvByKey  = elemUFM_Directly
+alterVarEnv      = alterUFM
 extendVarEnv	 = addToUFM
 extendVarEnv_C	 = addToUFM_C
 extendVarEnv_Acc = addToUFM_Acc
diff --git a/compiler/utils/UniqFM.lhs b/compiler/utils/UniqFM.lhs
index 7cbc3dbcfb..0cd9235cad 100644
--- a/compiler/utils/UniqFM.lhs
+++ b/compiler/utils/UniqFM.lhs
@@ -36,7 +36,7 @@ module UniqFM (
 	addListToUFM,addListToUFM_C,
 	addToUFM_Directly,
 	addListToUFM_Directly,
-	adjustUFM,
+	adjustUFM, alterUFM,
 	adjustUFM_Directly,
 	delFromUFM,
 	delFromUFM_Directly,
@@ -108,6 +108,12 @@ addToUFM_Acc	:: Uniquable key =>
 			   -> key -> elt 		-- new
 			   -> UniqFM elts		-- result
 
+alterUFM        :: Uniquable key => 
+                              (Maybe elt -> Maybe elt)  -- How to adjust
+			   -> UniqFM elt 		-- old
+			   -> key        		-- new
+			   -> UniqFM elt		-- result
+
 addListToUFM_C	:: Uniquable key => (elt -> elt -> elt)
 			   -> UniqFM elt -> [(key,elt)]
 			   -> UniqFM elt
@@ -182,7 +188,8 @@ listToUFM = foldl (\m (k, v) -> addToUFM m k v) emptyUFM
 listToUFM_Directly = foldl (\m (u, v) -> addToUFM_Directly m u v) emptyUFM
 listToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v) emptyUFM
 
-addToUFM (UFM m) k v = UFM (M.insert (getKey $ getUnique k) v m)
+alterUFM f (UFM m) k = UFM (M.alter f (getKey $ getUnique k) m)
+addToUFM (UFM m) k v   = UFM (M.insert (getKey $ getUnique k) v m)
 addListToUFM = foldl (\m (k, v) -> addToUFM m k v)
 addListToUFM_Directly = foldl (\m (k, v) -> addToUFM_Directly m k v)
 addToUFM_Directly (UFM m) u v = UFM (M.insert (getKey u) v m)