Add adjustUFM, adjustUFM_Directly and joinUFM to UniqFM.

Renamed adjustUFM in GraphOps to adjustUFM_C, to account
for alternate argument order.

Signed-off-by: Edward Z. Yang <ezyang@mit.edu>

diff --git a/compiler/utils/GraphOps.hs b/compiler/utils/GraphOps.hs
index 388b968..1fa4199 100644 (file)
--- a/compiler/utils/GraphOps.hs
+++ b/compiler/utils/GraphOps.hs
@@ -61,14 +61,14 @@ addNode k node graph
        -- add back conflict edges from other nodes to this one
        map_conflict    
                = foldUniqSet 
-                       (adjustUFM (\n -> n { nodeConflicts = addOneToUniqSet (nodeConflicts n) k}))
+                       (adjustUFM_C (\n -> n { nodeConflicts = addOneToUniqSet (nodeConflicts n) k}))
                        (graphMap graph)
                        (nodeConflicts node)
                        
        -- add back coalesce edges from other nodes to this one
        map_coalesce
                = foldUniqSet
-                       (adjustUFM (\n -> n { nodeCoalesce = addOneToUniqSet (nodeCoalesce n) k}))
+                       (adjustUFM_C (\n -> n { nodeCoalesce = addOneToUniqSet (nodeCoalesce n) k}))
                        map_conflict
                        (nodeCoalesce node)
        
@@ -434,7 +434,7 @@ freezeNode k
                else node       -- panic "GraphOps.freezeNode: edge to freeze wasn't in the coalesce set"
                                -- If the edge isn't actually in the coelesce set then just ignore it.
 
-       fm2     = foldUniqSet (adjustUFM (freezeEdge k)) fm1
+       fm2     = foldUniqSet (adjustUFM_C (freezeEdge k)) fm1
                        $ nodeCoalesce node
 
     in fm2
@@ -604,7 +604,7 @@ setColor
        
 setColor u color
        = graphMapModify
-       $ adjustUFM
+       $ adjustUFM_C
                (\n -> n { nodeColor = Just color })
                u 
        
@@ -621,13 +621,14 @@ adjustWithDefaultUFM f def k map
                map
                k def
                
-{-# INLINE adjustUFM #-}
-adjustUFM 
+-- Argument order different from UniqFM's adjustUFM
+{-# INLINE adjustUFM_C #-}
+adjustUFM_C 
        :: Uniquable k
        => (a -> a)
        -> k -> UniqFM a -> UniqFM a
 
-adjustUFM f k map
+adjustUFM_C f k map
  = case lookupUFM map k of
        Nothing -> map
        Just a  -> addToUFM map k (f a)

diff --git a/compiler/utils/UniqFM.lhs b/compiler/utils/UniqFM.lhs
index 31d1e87..7302b02 100644 (file)
--- a/compiler/utils/UniqFM.lhs
+++ b/compiler/utils/UniqFM.lhs
@@ -36,6 +36,8 @@ module UniqFM (
        addListToUFM,addListToUFM_C,
        addToUFM_Directly,
        addListToUFM_Directly,
+       adjustUFM,
+       adjustUFM_Directly,
        delFromUFM,
        delFromUFM_Directly,
        delListFromUFM,
@@ -53,12 +55,15 @@ module UniqFM (
        lookupUFM, lookupUFM_Directly,
        lookupWithDefaultUFM, lookupWithDefaultUFM_Directly,
        eltsUFM, keysUFM, splitUFM,
-       ufmToList 
+       ufmToList,
+       joinUFM
     ) where
 
 import Unique           ( Uniquable(..), Unique, getKey )
 import Outputable
 
+import Compiler.Hoopl   hiding (Unique)
+
 import qualified Data.IntMap as M
 \end{code}
 
@@ -103,6 +108,9 @@ addListToUFM_C      :: Uniquable key => (elt -> elt -> elt)
                           -> UniqFM elt -> [(key,elt)]
                           -> UniqFM elt
 
+adjustUFM      :: Uniquable key => (elt -> elt) -> UniqFM elt -> key -> UniqFM elt
+adjustUFM_Directly :: (elt -> elt) -> UniqFM elt -> Unique -> UniqFM elt
+
 delFromUFM     :: Uniquable key => UniqFM elt -> key    -> UniqFM elt
 delListFromUFM :: Uniquable key => UniqFM elt -> [key] -> UniqFM elt
 delFromUFM_Directly :: UniqFM elt -> Unique -> UniqFM elt
@@ -175,6 +183,9 @@ addToUFM_Acc exi new (UFM m) k v =
   UFM (M.insertWith (\_new old -> exi v old) (getKey $ getUnique k) (new v) m)
 addListToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v)
 
+adjustUFM f (UFM m) k = UFM (M.adjust f (getKey $ getUnique k) m)
+adjustUFM_Directly f (UFM m) u = UFM (M.adjust f (getKey u) m)
+
 delFromUFM (UFM m) k = UFM (M.delete (getKey $ getUnique k) m)
 delListFromUFM = foldl delFromUFM
 delFromUFM_Directly (UFM m) u = UFM (M.delete (getKey u) m)
@@ -207,6 +218,16 @@ keysUFM (UFM m) = map getUnique $ M.keys m
 eltsUFM (UFM m) = M.elems m
 ufmToList (UFM m) = map (\(k, v) -> (getUnique k, v)) $ M.toList m
 
+-- Hoopl
+joinUFM :: JoinFun v -> JoinFun (UniqFM v)
+joinUFM eltJoin l (OldFact old) (NewFact new) = foldUFM_Directly add (NoChange, old) new
+    where add k new_v (ch, joinmap) =
+            case lookupUFM_Directly joinmap k of
+                Nothing -> (SomeChange, addToUFM_Directly joinmap k new_v)
+                Just old_v -> case eltJoin l (OldFact old_v) (NewFact new_v) of
+                                (SomeChange, v') -> (SomeChange, addToUFM_Directly joinmap k v')
+                                (NoChange, _) -> (ch, joinmap)
+
 \end{code}
 
 %************************************************************************