compiler: add util functions for UniqFM and UniqMap

This patch adds addToUFM_L (backed by insertLookupWithKey),
addToUniqMap_L and intersectUniqMap_C. These UniqFM/UniqMap util
functions are used by the wasm32 NCG.

2 files changed, 41 insertions, 1 deletions

diff --git a/compiler/GHC/Types/Unique/FM.hs b/compiler/GHC/Types/Unique/FM.hs
index 8f96731599..137e985f92 100644
--- a/compiler/GHC/Types/Unique/FM.hs
+++ b/compiler/GHC/Types/Unique/FM.hs
@@ -40,7 +40,7 @@ module GHC.Types.Unique.FM (
         listToUFM_Directly,
         listToUFM_C,
         listToIdentityUFM,
-        addToUFM,addToUFM_C,addToUFM_Acc,
+        addToUFM,addToUFM_C,addToUFM_Acc,addToUFM_L,
         addListToUFM,addListToUFM_C,
         addToUFM_Directly,
         addListToUFM_Directly,
@@ -183,6 +183,24 @@ addToUFM_Acc
 addToUFM_Acc exi new (UFM m) k v =
   UFM (M.insertWith (\_new old -> exi v old) (getKey $ getUnique k) (new v) m)
 
+-- | Add an element, returns previous lookup result and new map. If
+-- old element doesn't exist, add the passed element directly,
+-- otherwise compute the element to add using the passed function.
+addToUFM_L
+  :: Uniquable key
+  => (key -> elt -> elt -> elt) -- key,old,new
+  -> key
+  -> elt -- new
+  -> UniqFM key elt
+  -> (Maybe elt, UniqFM key elt) -- old, result
+addToUFM_L f k v (UFM m) =
+  coerce $
+    M.insertLookupWithKey
+      (\_ _n _o -> f k _o _n)
+      (getKey $ getUnique k)
+      v
+      m
+
 alterUFM
   :: Uniquable key
   => (Maybe elt -> Maybe elt)  -- How to adjust
diff --git a/compiler/GHC/Types/Unique/Map.hs b/compiler/GHC/Types/Unique/Map.hs
index 18d3b2a73a..aef23be566 100644
--- a/compiler/GHC/Types/Unique/Map.hs
+++ b/compiler/GHC/Types/Unique/Map.hs
@@ -20,6 +20,7 @@ module GHC.Types.Unique.Map (
     addListToUniqMap,
     addToUniqMap_C,
     addToUniqMap_Acc,
+    addToUniqMap_L,
     alterUniqMap,
     addListToUniqMap_C,
     adjustUniqMap,
@@ -31,6 +32,7 @@ module GHC.Types.Unique.Map (
     plusUniqMapList,
     minusUniqMap,
     intersectUniqMap,
+    intersectUniqMap_C,
     disjointUniqMap,
     mapUniqMap,
     filterUniqMap,
@@ -123,6 +125,22 @@ addToUniqMap_Acc exi new (UniqMap m) k0 v0 = UniqMap $
                  (\b -> (k0, new b))
                  m k0 v0
 
+-- | Add an element, returns previous lookup result and new map. If
+-- old element doesn't exist, add the passed element directly,
+-- otherwise compute the element to add using the passed function.
+addToUniqMap_L :: Uniquable k
+               => (k -> a -> a -> a) -- key,old,new
+               -> k
+               -> a -- new
+               -> UniqMap k a
+               -> (Maybe a, UniqMap k a)
+addToUniqMap_L f k v (UniqMap m) = case addToUFM_L
+  (\_k (_, _o) (_, _n) -> (_k, f _k _o _n))
+  k
+  (k, v)
+  m of
+  (_maybe, _ufm) -> (snd <$> _maybe, UniqMap _ufm)
+
 alterUniqMap :: Uniquable k
              => (Maybe a -> Maybe a)
              -> UniqMap k a
@@ -176,6 +194,10 @@ minusUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ minusUFM m1 m2
 intersectUniqMap :: UniqMap k a -> UniqMap k b -> UniqMap k a
 intersectUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ intersectUFM m1 m2
 
+-- | Intersection with a combining function.
+intersectUniqMap_C :: (a -> b -> c) -> UniqMap k a -> UniqMap k b -> UniqMap k c
+intersectUniqMap_C f (UniqMap m1) (UniqMap m2) = UniqMap $ intersectUFM_C (\(k, a) (_, b) -> (k, f a b)) m1 m2
+
 disjointUniqMap :: UniqMap k a -> UniqMap k b -> Bool
 disjointUniqMap (UniqMap m1) (UniqMap m2) = disjointUFM m1 m2