Strict.mapMaybe[WithKey]: Fix strictness for collisions (#385)

Fixes #381.

Author: sjakobi

Data/HashMap/Internal/Strict.hs

@@ -548,7 +548,7 @@ mapMaybeWithKey f = filterMapAux onLeaf onColl
   where onLeaf (Leaf h (L k v)) | Just v' <- f k v = Just (leaf h k v')
         onLeaf _ = Nothing
 
-        onColl (L k v) | Just v' <- f k v = Just (L k v')
+        onColl (L k v) | Just !v' <- f k v = Just (L k v')
                        | otherwise = Nothing
 {-# INLINE mapMaybeWithKey #-}
 

tests/Regressions.hs

@@ -187,6 +187,27 @@ issue379LazyUnionWith = do
   touch v -- makes sure that we didn't GC away the combined value
   assert $ isNothing res
 
+------------------------------------------------------------------------
+-- Issue #381
+
+#ifdef HAVE_NOTHUNKS
+
+issue381mapMaybe :: Assertion
+issue381mapMaybe = do
+  let m0 = HMS.fromList [(KC 1, 10), (KC 2, 20 :: Int)]
+  let m1 = HMS.mapMaybe (Just . (+ 1)) m0
+  mThunkInfo <- noThunksInValues mempty (Foldable.toList m1)
+  assert $ isNothing mThunkInfo
+
+issue381mapMaybeWithKey :: Assertion
+issue381mapMaybeWithKey = do
+  let m0 = HMS.fromList [(KC 1, 10), (KC 2, 20 :: Int)]
+  let m1 = HMS.mapMaybeWithKey (\(KC k) v -> Just (k + v)) m0
+  mThunkInfo <- noThunksInValues mempty (Foldable.toList m1)
+  assert $ isNothing mThunkInfo
+
+#endif
+
 ------------------------------------------------------------------------
 -- * Test list
 
@@ -206,4 +227,10 @@ tests = testGroup "Regression tests"
           , testCase "Strict.unionWithKey" issue379StrictUnionWithKey
 #endif
           ]
+#ifdef HAVE_NOTHUNKS
+    , testGroup "issue381"
+          [ testCase "mapMaybe" issue381mapMaybe
+          , testCase "mapMaybeWithKey" issue381mapMaybeWithKey
+          ]
+#endif
     ]