Fix strictness bugs in fromDistinctAscList and fromDistinctDescList (#996)

* Add a counterexample annotation to whnfHasNoThunks

* Add NoThunks tests for `fromDistinctAscList` and `fromDistinctDescList`

* Fix strictness bugs in `fromDistinctAscList` and `fromDistinctDescList`

Author: neilmayhew

containers-tests/tests/Utils/NoThunks.hs

@@ -1,15 +1,12 @@
 module Utils.NoThunks (whnfHasNoThunks) where
 
-import Data.Maybe (isNothing)
-
 import NoThunks.Class (NoThunks, noThunks)
-import Test.QuickCheck (Property, ioProperty)
+import Test.QuickCheck (Property, counterexample, ioProperty, property)
 
 -- | Check that after evaluating the argument to weak head normal form there
 -- are no thunks.
 --
 whnfHasNoThunks :: NoThunks a => a -> Property
-whnfHasNoThunks a = ioProperty
-                  . fmap isNothing
-                  . noThunks []
-                 $! a
+whnfHasNoThunks a = ioProperty $
+  maybe (property True) ((`counterexample` False) . show)
+    <$> (noThunks [] $! a)

containers-tests/tests/intmap-strictness.hs

@@ -116,6 +116,13 @@ pStrictFoldl' :: IntMap Int -> Property
 pStrictFoldl' m = whnfHasNoThunks (M.foldl' (flip (:)) [] m)
 #endif
 
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFromDistinctAscList :: [Int] -> Property
+pStrictFromDistinctAscList = whnfHasNoThunks . evalSpine . M.elems . M.fromDistinctAscList . zip [0::Int ..] . map (Just $!)
+  where
+    evalSpine xs = length xs `seq` xs
+#endif
+
 ------------------------------------------------------------------------
 -- check for extra thunks
 --
@@ -202,6 +209,7 @@ tests =
 #if __GLASGOW_HASKELL__ >= 806
       , testProperty "strict foldr'" pStrictFoldr'
       , testProperty "strict foldl'" pStrictFoldl'
+      , testProperty "strict fromDistinctAscList" pStrictFromDistinctAscList
 #endif
       ]
       , tExtraThunksM

containers-tests/tests/map-strictness.hs

@@ -108,6 +108,20 @@ pStrictFoldlWithKey' :: Map Int Int -> Property
 pStrictFoldlWithKey' m = whnfHasNoThunks (M.foldlWithKey' (\as _ a -> a : as) [] m)
 #endif
 
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFromDistinctAscList :: [Int] -> Property
+pStrictFromDistinctAscList = whnfHasNoThunks . evalSpine . M.elems . M.fromDistinctAscList . zip [0::Int ..] . map (Just $!)
+  where
+    evalSpine xs = length xs `seq` xs
+#endif
+
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFromDistinctDescList :: [Int] -> Property
+pStrictFromDistinctDescList = whnfHasNoThunks . evalSpine . M.elems . M.fromDistinctDescList . zip [0::Int, -1 ..] . map (Just $!)
+  where
+    evalSpine xs = length xs `seq` xs
+#endif
+
 ------------------------------------------------------------------------
 -- check for extra thunks
 --
@@ -193,6 +207,8 @@ tests =
       , testProperty "strict foldl'" pStrictFoldl'
       , testProperty "strict foldrWithKey'" pStrictFoldrWithKey'
       , testProperty "strict foldlWithKey'" pStrictFoldlWithKey'
+      , testProperty "strict fromDistinctAscList" pStrictFromDistinctAscList
+      , testProperty "strict fromDistinctDescList" pStrictFromDistinctDescList
 #endif
       ]
       , tExtraThunksM

containers/src/Data/Map/Strict/Internal.hs

@@ -1762,7 +1762,7 @@ fromDistinctAscList = fromDistinctAscList_linkAll . Foldable.foldl' next (State0
   where
     next :: FromDistinctMonoState k a -> (k,a) -> FromDistinctMonoState k a
     next (State0 stk) (!kx, !x) = fromDistinctAscList_linkTop (Bin 1 kx x Tip Tip) stk
-    next (State1 l stk) (kx, x) = State0 (Push kx x l stk)
+    next (State1 l stk) (!kx, !x) = State0 (Push kx x l stk)
 {-# INLINE fromDistinctAscList #-}  -- INLINE for fusion
 
 -- | \(O(n)\). Build a map from a descending list of distinct elements in linear time.
@@ -1781,5 +1781,5 @@ fromDistinctDescList = fromDistinctDescList_linkAll . Foldable.foldl' next (Stat
   where
     next :: FromDistinctMonoState k a -> (k,a) -> FromDistinctMonoState k a
     next (State0 stk) (!kx, !x) = fromDistinctDescList_linkTop (Bin 1 kx x Tip Tip) stk
-    next (State1 r stk) (kx, x) = State0 (Push kx x r stk)
+    next (State1 r stk) (!kx, !x) = State0 (Push kx x r stk)
 {-# INLINE fromDistinctDescList #-}  -- INLINE for fusion