diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs index 0233935e..f6561890 100644 --- a/Data/HashMap/Internal.hs +++ b/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 diff --git a/Data/HashMap/Internal/Strict.hs b/Data/HashMap/Internal/Strict.hs index 8b891690..a7cfcdfe 100644 --- a/Data/HashMap/Internal/Strict.hs +++ b/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) = diff --git a/tests/Regressions.hs b/tests/Regressions.hs index 808a96e3..30997649 100644 --- a/tests/Regressions.hs +++ b/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 + ] ] diff --git a/unordered-containers.cabal b/unordered-containers.cabal index 363ab3dc..47d8dbff 100644 --- a/unordered-containers.cabal +++ b/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