Fix space leaks in union[With[Key]] (#379)

Data/HashMap/Internal.hs

@@ -122,7 +122,6 @@ module Data.HashMap.Internal
     , update32
     , update32M
     , update32With'
-    , updateOrConcatWith
     , updateOrConcatWithKey
     , filterMapAux
     , equalKeys
@@ -1551,7 +1550,7 @@ unionWithKey f = go 0
         | h1 == h2  = Collision h1 (updateOrSnocWithKey (\k a b -> (# f k b a #)) k2 v2 ls1)
         | otherwise = goDifferentHash s h1 h2 t1 t2
     go s t1@(Collision h1 ls1) t2@(Collision h2 ls2)
-        | h1 == h2  = Collision h1 (updateOrConcatWithKey f ls1 ls2)
+        | h1 == h2  = Collision h1 (updateOrConcatWithKey (\k a b -> (# f k a b #)) ls1 ls2)
         | otherwise = goDifferentHash s h1 h2 t1 t2
     -- branch vs. branch
     go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) =
@@ -2177,11 +2176,7 @@ updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
             = go k v ary (i+1) n
 {-# INLINABLE updateOrSnocWithKey #-}
 
-updateOrConcatWith :: Eq k => (v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
-updateOrConcatWith f = updateOrConcatWithKey (const f)
-{-# INLINABLE updateOrConcatWith #-}
-
-updateOrConcatWithKey :: Eq k => (k -> v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
+updateOrConcatWithKey :: Eq k => (k -> v -> v -> (# v #)) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
 updateOrConcatWithKey f ary1 ary2 = A.run $ do
     -- TODO: instead of mapping and then folding, should we traverse?
     -- We'll have to be careful to avoid allocating pairs or similar.
@@ -2203,7 +2198,7 @@ updateOrConcatWithKey f ary1 ary2 = A.run $ do
                Just i1 -> do -- key occurs in both arrays, store combination in position i1
                              L k v1 <- A.indexM ary1 i1
                              L _ v2 <- A.indexM ary2 i2
-                             A.write mary i1 (L k (f k v1 v2))
+                             case f k v1 v2 of (# v3 #) -> A.write mary i1 (L k v3)
                              go iEnd (i2+1)
                Nothing -> do -- key is only in ary2, append to end
                              A.write mary iEnd =<< A.indexM ary2 i2

Data/HashMap/Internal/Strict.hs

@@ -453,7 +453,7 @@ unionWithKey f = go 0
         | h1 == h2  = Collision h1 (updateOrSnocWithKey (flip . f) k2 v2 ls1)
         | otherwise = goDifferentHash s h1 h2 t1 t2
     go s t1@(Collision h1 ls1) t2@(Collision h2 ls2)
-        | h1 == h2  = Collision h1 (updateOrConcatWithKey f ls1 ls2)
+        | h1 == h2  = Collision h1 (updateOrConcatWithKey (\k a b -> let !v = f k a b in (# v #)) ls1 ls2)
         | otherwise = goDifferentHash s h1 h2 t1 t2
     -- branch vs. branch
     go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) =

tests/Regressions.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE CPP                 #-}
 {-# LANGUAGE MagicHash           #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE UnboxedTuples       #-}
@@ -22,6 +23,13 @@ import Test.Tasty.QuickCheck (testProperty)
 import qualified Data.HashMap.Lazy   as HML
 import qualified Data.HashMap.Strict as HMS
 
+#if MIN_VERSION_base(4,12,0)
+-- nothunks requires base >= 4.12
+#define HAVE_NOTHUNKS
+import qualified Data.Foldable  as Foldable
+import           NoThunks.Class (noThunksInValues)
+#endif
+
 issue32 :: Assertion
 issue32 = assert $ isJust $ HMS.lookup 7 m'
   where
@@ -124,6 +132,61 @@ issue254Strict = do
   touch mp
   assert $ isNothing res
 
+------------------------------------------------------------------------
+-- Issue #379
+
+#ifdef HAVE_NOTHUNKS
+
+issue379Union :: Assertion
+issue379Union = do
+  let m0 = HMS.fromList [(KC 1, ()), (KC 2, ())]
+  let m1 = HMS.fromList [(KC 2, ()), (KC 3, ())]
+  let u = m0 `HMS.union` m1
+  mThunkInfo <- noThunksInValues mempty (Foldable.toList u)
+  assert $ isNothing mThunkInfo
+
+issue379StrictUnionWith :: Assertion
+issue379StrictUnionWith = do
+  let m0 = HMS.fromList [(KC 1, 10), (KC 2, 20 :: Int)]
+  let m1 = HMS.fromList [(KC 2, 20), (KC 3, 30)]
+  let u = HMS.unionWith (+) m0 m1
+  mThunkInfo <- noThunksInValues mempty (Foldable.toList u)
+  assert $ isNothing mThunkInfo
+
+issue379StrictUnionWithKey :: Assertion
+issue379StrictUnionWithKey = do
+  let m0 = HMS.fromList [(KC 1, 10), (KC 2, 20 :: Int)]
+  let m1 = HMS.fromList [(KC 2, 20), (KC 3, 30)]
+  let u = HMS.unionWithKey (\(KC i) v0 v1 -> i + v0 + v1) m0 m1
+  mThunkInfo <- noThunksInValues mempty (Foldable.toList u)
+  assert $ isNothing mThunkInfo
+
+#endif
+
+-- Another key type that always collides.
+--
+-- Note (sjakobi): The KC newtype of Int somehow can't be used to demonstrate
+-- the space leak in issue379LazyUnionWith. This type does the trick.
+newtype SC = SC String
+  deriving (Eq, Ord, Show)
+instance Hashable SC where
+  hashWithSalt salt _ = salt
+
+issue379LazyUnionWith :: Assertion
+issue379LazyUnionWith = do
+  i :: Int <- randomIO
+  let k = SC (show i)
+  weakK <- mkWeakPtr k Nothing -- add the ability to test whether k is alive
+  let f :: Int -> Int
+      f x = error ("Should not be evaluated " ++ show x)
+  let m = HML.fromList [(SC "1", f 1), (SC "2", f 2), (k, f 3)]
+  let u = HML.unionWith (+) m m
+  Just v <- evaluate $ HML.lookup k u
+  performGC
+  res <- deRefWeak weakK -- gives Just if k is still alive
+  touch v -- makes sure that we didn't GC away the combined value
+  assert $ isNothing res
+
 ------------------------------------------------------------------------
 -- * Test list
 
@@ -135,4 +198,12 @@ tests = testGroup "Regression tests"
     , testProperty "issue39b" propEqAfterDelete
     , testCase "issue254 lazy" issue254Lazy
     , testCase "issue254 strict" issue254Strict
+    , testGroup "issue379"
+          [ testCase "Lazy.unionWith" issue379LazyUnionWith
+#ifdef HAVE_NOTHUNKS
+          , testCase "union" issue379Union
+          , testCase "Strict.unionWith" issue379StrictUnionWith
+          , testCase "Strict.unionWithKey" issue379StrictUnionWithKey
+#endif
+          ]
     ]

unordered-containers.cabal

@@ -98,6 +98,10 @@ test-suite unordered-containers-tests
     tasty-quickcheck >= 0.10.1.2,
     unordered-containers
 
+  if impl(ghc >= 8.6)
+    build-depends:
+      nothunks >= 0.1.3
+
   default-language: Haskell2010
   ghc-options: -Wall
   cpp-options: -DASSERTS