#254, fix space leak on collisions

Data/HashMap/Base.hs

@@ -757,7 +757,7 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
             else Full (update16 ary i st')
       where i = index h s
     go h k x s t@(Collision hy v)
-        | h == hy   = Collision h (updateOrSnocWith const k x v)
+        | h == hy   = Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v)
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
 {-# INLINABLE insert' #-}
 
@@ -880,7 +880,7 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
         return t
       where i = index h s
     go h k x s t@(Collision hy v)
-        | h == hy   = return $! Collision h (updateOrSnocWith const k x v)
+        | h == hy   = return $! Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v)
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
 {-# INLINABLE unsafeInsert #-}
 
@@ -1026,7 +1026,7 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
         return t
       where i = index h s
     go h k x s t@(Collision hy v)
-        | h == hy   = return $! Collision h (updateOrSnocWith f k x v)
+        | h == hy   = return $! Collision h (updateOrSnocWith (\a b -> (# f a b #)) k x v)
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
 {-# INLINABLE unsafeInsertWith #-}
 
@@ -1394,10 +1394,10 @@ unionWithKey f = go 0
                       else collision h1 l1 l2
         | otherwise = goDifferentHash s h1 h2 t1 t2
     go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2)
-        | h1 == h2  = Collision h1 (updateOrSnocWithKey f k1 v1 ls2)
+        | h1 == h2  = Collision h1 (updateOrSnocWithKey (\k a b -> (# f k a b #)) k1 v1 ls2)
         | otherwise = goDifferentHash s h1 h2 t1 t2
     go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2))
-        | h1 == h2  = Collision h1 (updateOrSnocWithKey (flip . f) k2 v2 ls1)
+        | 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)
@@ -1932,12 +1932,12 @@ updateWith# f k0 ary0 = go k0 ary0 0 (A.length ary0)
                      | otherwise -> go k ary (i+1) n
 {-# INLINABLE updateWith# #-}
 
-updateOrSnocWith :: Eq k => (v -> v -> v) -> k -> v -> A.Array (Leaf k v)
+updateOrSnocWith :: Eq k => (v -> v -> (# v #)) -> k -> v -> A.Array (Leaf k v)
                  -> A.Array (Leaf k v)
 updateOrSnocWith f = updateOrSnocWithKey (const f)
 {-# INLINABLE updateOrSnocWith #-}
 
-updateOrSnocWithKey :: Eq k => (k -> v -> v -> v) -> k -> v -> A.Array (Leaf k v)
+updateOrSnocWithKey :: Eq k => (k -> v -> v -> (# v #)) -> k -> v -> A.Array (Leaf k v)
                  -> A.Array (Leaf k v)
 updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
   where
@@ -1948,9 +1948,12 @@ updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
             A.copy ary 0 mary 0 n
             A.write mary n (L k v)
             return mary
-        | otherwise = case A.index ary i of
-            (L kx y) | k == kx   -> A.update ary i (L k (f k v y))
-                     | otherwise -> go k v ary (i+1) n
+        | L kx y <- A.index ary i
+        , k == kx
+        , (# v2 #) <- f k v y
+            = A.update ary i (L k v2)
+        | otherwise
+            = 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)

tests/Regressions.hs

@@ -1,10 +1,21 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
 module Main where
 
 import Control.Applicative ((<$>))
+import Control.Exception (evaluate)
 import Control.Monad (replicateM)
+import Data.Hashable (Hashable(..))
 import qualified Data.HashMap.Strict as HM
+import qualified Data.HashMap.Lazy as HML
 import Data.List (delete)
 import Data.Maybe
+import GHC.Exts (touch#)
+import GHC.IO (IO (..))
+import System.Mem (performGC)
+import System.Mem.Weak (mkWeakPtr, deRefWeak)
+import System.Random (randomIO)
 import Test.HUnit (Assertion, assert)
 import Test.Framework (Test, defaultMain)
 import Test.Framework.Providers.HUnit (testCase)
@@ -71,6 +82,48 @@ propEqAfterDelete (Keys keys) =
 mapFromKeys :: [Int] -> HM.HashMap Int ()
 mapFromKeys keys = HM.fromList (zip keys (repeat ()))
 
+------------------------------------------------------------------------
+-- Issue #254
+
+-- Key type that always collides.
+newtype KC = KC Int
+  deriving (Eq, Ord, Show)
+instance Hashable KC where
+  hashWithSalt salt _ = salt
+
+touch :: a -> IO ()
+touch a = IO (\s -> (# touch# a s, () #))
+
+-- We want to make sure that old values in the HashMap are evicted when new values are inserted,
+-- even if they aren't evaluated. To do that, we use the WeakPtr trick described at
+-- http://simonmar.github.io/posts/2018-06-20-Finding-fixing-space-leaks.html.
+-- We insert a value named oldV into the HashMap, then insert over it, checking oldV is no longer reachable.
+--
+-- To make the test robust, it's important that oldV isn't hoisted up to the top or shared.
+-- To do that, we generate it randomly.
+issue254Lazy :: Assertion
+issue254Lazy = do
+  i :: Int <- randomIO
+  let oldV = error $ "Should not be evaluated: " ++ show i
+  weakV <- mkWeakPtr oldV Nothing -- add the ability to test whether oldV is alive
+  mp <- evaluate $ HML.insert (KC 1) (error "Should not be evaluated") $ HML.fromList [(KC 0, "1"), (KC 1, oldV)]
+  performGC
+  res <- deRefWeak weakV -- gives Just if oldV is still alive
+  touch mp -- makes sure that we didn't GC away the whole HashMap, just oldV
+  assert $ isNothing res
+
+-- Like issue254Lazy, but using strict HashMap
+issue254Strict :: Assertion
+issue254Strict = do
+  i :: Int <- randomIO
+  let oldV = show i
+  weakV <- mkWeakPtr oldV Nothing
+  mp <- evaluate $ HM.insert (KC 1) "3" $ HM.fromList [(KC 0, "1"), (KC 1, oldV)]
+  performGC
+  res <- deRefWeak weakV
+  touch mp
+  assert $ isNothing res
+
 ------------------------------------------------------------------------
 -- * Test list
 
@@ -80,6 +133,8 @@ tests =
       testCase "issue32" issue32
     , testCase "issue39a" issue39
     , testProperty "issue39b" propEqAfterDelete
+    , testCase "issue254 lazy" issue254Lazy
+    , testCase "issue254 strict" issue254Strict
     ]
 
 ------------------------------------------------------------------------

unordered-containers.cabal

@@ -155,6 +155,7 @@ test-suite regressions
     hashable >= 1.0.1.1,
     HUnit,
     QuickCheck >= 2.4.0.1,
+    random,
     test-framework >= 0.3.3,
     test-framework-hunit,
     test-framework-quickcheck2,