Make sure mergeWithKey is inlined and applied strictly

In the particular case of `DmdEnv`, not applying this function strictly
meant 500MB of thunks were accumulated before the values were forced at
the end of demand analysis.

2 files changed, 15 insertions, 6 deletions

diff --git a/compiler/GHC/Types/Name/Env.hs b/compiler/GHC/Types/Name/Env.hs
index 1a94dc4fa0..509ea5f7bb 100644
--- a/compiler/GHC/Types/Name/Env.hs
+++ b/compiler/GHC/Types/Name/Env.hs
@@ -134,6 +134,7 @@ mkNameEnvWith f       = mkNameEnv . map (\a -> (f a, a))
 elemNameEnv x y          = elemUFM x y
 plusNameEnv x y          = plusUFM x y
 plusNameEnv_C f x y      = plusUFM_C f x y
+{-# INLINE plusNameEnv_CD #-}
 plusNameEnv_CD f x d y b = plusUFM_CD f x d y b
 plusNameEnv_CD2 f x y    = plusUFM_CD2 f x y
 extendNameEnv_C f x y z  = addToUFM_C f x y z
diff --git a/compiler/GHC/Types/Unique/FM.hs b/compiler/GHC/Types/Unique/FM.hs
index 6d13436169..6c2eec6a6d 100644
--- a/compiler/GHC/Types/Unique/FM.hs
+++ b/compiler/GHC/Types/Unique/FM.hs
@@ -86,6 +86,7 @@ import GHC.Utils.Outputable
 import GHC.Utils.Panic (assertPanic)
 import GHC.Utils.Misc (debugIsOn)
 import qualified Data.IntMap as M
+import qualified Data.IntMap.Strict as MS
 import qualified Data.IntSet as S
 import Data.Data
 import qualified Data.Semigroup as Semi
@@ -229,12 +230,16 @@ plusUFM_C f (UFM x) (UFM y) = UFM (M.unionWith f x y)
 -- there is no entry in `m1` reps. `m2`. The domain is the union of
 -- the domains of `m1` and `m2`.
 --
+-- IMPORTANT NOTE: This function strictly applies the modification function
+-- and forces the result unlike most the other functions in this module.
+--
 -- Representative example:
 --
 -- @
 -- plusUFM_CD f {A: 1, B: 2} 23 {B: 3, C: 4} 42
 --    == {A: f 1 42, B: f 2 3, C: f 23 4 }
 -- @
+{-# INLINE plusUFM_CD #-}
 plusUFM_CD
   :: (elta -> eltb -> eltc)
   -> UniqFM key elta  -- map X
@@ -243,10 +248,10 @@ plusUFM_CD
   -> eltb         -- default for Y
   -> UniqFM key eltc
 plusUFM_CD f (UFM xm) dx (UFM ym) dy
-  = UFM $ M.mergeWithKey
+  = UFM $ MS.mergeWithKey
       (\_ x y -> Just (x `f` y))
-      (M.map (\x -> x `f` dy))
-      (M.map (\y -> dx `f` y))
+      (MS.map (\x -> x `f` dy))
+      (MS.map (\y -> dx `f` y))
       xm ym
 
 -- | `plusUFM_CD2 f m1 m2` merges the maps using `f` as the combining
@@ -254,6 +259,9 @@ plusUFM_CD f (UFM xm) dx (UFM ym) dy
 -- instead passed as `Nothing` to `f`. `f` can never have both its arguments
 -- be `Nothing`.
 --
+-- IMPORTANT NOTE: This function strictly applies the modification function
+-- and forces the result.
+--
 -- `plusUFM_CD2 f m1 m2` is the same as `plusUFM_CD f (mapUFM Just m1) Nothing
 -- (mapUFM Just m2) Nothing`.
 plusUFM_CD2
@@ -262,10 +270,10 @@ plusUFM_CD2
   -> UniqFM key eltb  -- map Y
   -> UniqFM key eltc
 plusUFM_CD2 f (UFM xm) (UFM ym)
-  = UFM $ M.mergeWithKey
+  = UFM $ MS.mergeWithKey
       (\_ x y -> Just (Just x `f` Just y))
-      (M.map (\x -> Just x `f` Nothing))
-      (M.map (\y -> Nothing `f` Just y))
+      (MS.map (\x -> Just x `f` Nothing))
+      (MS.map (\y -> Nothing `f` Just y))
       xm ym
 
 plusMaybeUFM_C :: (elt -> elt -> Maybe elt)