Make Map's and Set's: foldr' & firends ' more strict

* Map.foldr'
* Map.foldl'
* Map.foldrWithKey'
* Map.foldlWithKey'
* Set.foldr'
* Set.foldl'

They now evaluate intermediate results of the accumulator before
appling the folding operator again.

This patch contains tests based on `nothunks` library.

Benchmark summary:

Map old:
```
benchmarking foldlWithKey' ...
benchmarked foldlWithKey'
time                 28.57 μs   (28.00 μs .. 29.35 μs)
                     0.991 R²   (0.979 R² .. 0.999 R²)
mean                 28.43 μs   (28.12 μs .. 28.99 μs)
std dev              1.358 μs   (729.0 ns .. 2.101 μs)
variance introduced by outliers: 27% (moderately inflated)

benchmarking foldrWithKey' ...
benchmarked foldrWithKey'
time                 80.25 ns   (79.45 ns .. 81.13 ns)
                     0.998 R²   (0.996 R² .. 0.999 R²)
mean                 79.84 ns   (79.27 ns .. 80.68 ns)
std dev              2.184 ns   (1.572 ns .. 2.921 ns)
variance introduced by outliers: 11% (moderately inflated)
```

Map new:
```
benchmarking foldlWithKey' ...
benchmarked foldlWithKey'
time                 27.76 μs   (27.15 μs .. 28.34 μs)
                     0.996 R²   (0.992 R² .. 0.998 R²)
mean                 27.08 μs   (26.84 μs .. 27.43 μs)
std dev              964.2 ns   (720.7 ns .. 1.350 μs)
variance introduced by outliers: 18% (moderately inflated)

benchmarking foldrWithKey' ...
benchmarked foldrWithKey'
time                 74.02 ns   (73.01 ns .. 75.82 ns)
                     0.998 R²   (0.994 R² .. 1.000 R²)
mean                 73.14 ns   (72.91 ns .. 73.80 ns)
std dev              1.245 ns   (434.3 ps .. 2.625 ns)
```

Set old:

```benchmarking member ...
benchmarked member
time                 237.1 μs   (231.5 μs .. 246.0 μs)
                     0.993 R²   (0.986 R² .. 0.999 R²)
mean                 234.3 μs   (232.1 μs .. 238.0 μs)
std dev              9.031 μs   (5.737 μs .. 15.46 μs)
variance introduced by outliers: 20% (moderately inflated)
```

Set new:
```
benchmarking member ...
benchmarked member
time                 219.4 μs   (216.3 μs .. 222.6 μs)
                     0.999 R²   (0.997 R² .. 1.000 R²)
mean                 221.3 μs   (219.3 μs .. 225.3 μs)
std dev              10.25 μs   (5.408 μs .. 18.71 μs)
variance introduced by outliers: 28% (moderately inflated)
```

Author: coot

containers-tests/benchmarks/Map.hs

@@ -72,8 +72,9 @@ main = do
         , bench "insertLookupWithKey' present" $ whnf (insLookupWithKey' elems_even) m_even
         , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m
         , bench "foldlWithKey" $ whnf (ins elems) m
---         , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sum 0) m
+        , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sum 0) m
         , bench "foldrWithKey" $ whnf (M.foldrWithKey consPair []) m
+        , bench "foldrWithKey'" $ whnf (M.foldrWithKey' consPair []) m
         , bench "update absent" $ whnf (upd Just evens) m_odd
         , bench "update present" $ whnf (upd Just evens) m_even
         , bench "update delete" $ whnf (upd (const Nothing) evens) m

containers-tests/containers-tests.cabal

@@ -41,9 +41,13 @@ library
       array    >=0.4.0.0
     , base     >=4.6     && <5
     , deepseq  >=1.2     && <1.5
+  if impl(ghc >= 8.6.0)
+    build-depends:
+        nothunks
+      , QuickCheck
 
   include-dirs:     include
-  hs-source-dirs:   src
+  hs-source-dirs:   src, tests
   ghc-options:      -O2 -Wall
   other-extensions:
     BangPatterns
@@ -79,6 +83,9 @@ library
     Utils.Containers.Internal.BitQueue
     Utils.Containers.Internal.BitUtil
     Utils.Containers.Internal.StrictPair
+  if impl(ghc >= 8.6.0)
+    exposed-modules:
+      Utils.NoThunks
 
   other-modules:
     Utils.Containers.Internal.Coercions
@@ -332,6 +339,12 @@ test-suite set-properties
     , test-framework-quickcheck2
     , transformers
 
+  if impl(ghc >= 8.6)
+    build-depends:
+      nothunks
+    other-modules:
+      Utils.NoThunks
+
 test-suite intmap-lazy-properties
   default-language: Haskell2010
   hs-source-dirs:   tests
@@ -475,6 +488,12 @@ test-suite map-strictness-properties
   other-modules:
     Utils.IsUnit
 
+  if impl(ghc >= 8.6)
+    build-depends:
+      nothunks
+    other-modules:
+      Utils.NoThunks
+
 test-suite intmap-strictness-properties
   default-language: Haskell2010
   hs-source-dirs:   tests
@@ -501,6 +520,12 @@ test-suite intmap-strictness-properties
   other-modules:
     Utils.IsUnit
 
+  if impl(ghc >= 8.6)
+    build-depends:
+      nothunks
+    other-modules:
+      Utils.NoThunks
+
 test-suite intset-strictness-properties
   default-language: Haskell2010
   hs-source-dirs:   tests
@@ -522,6 +547,12 @@ test-suite intset-strictness-properties
 
   ghc-options:      -Wall
 
+  if impl(ghc >= 8.6)
+    build-depends:
+      nothunks
+    other-modules:
+      Utils.NoThunks
+
 test-suite listutils-properties
   default-language: Haskell2010
   hs-source-dirs: tests

containers-tests/tests/Utils/NoThunks.hs

@@ -0,0 +1,15 @@
+module Utils.NoThunks (whnfHasNoThunks) where
+
+import Data.Maybe (isNothing)
+
+import NoThunks.Class (NoThunks, noThunks)
+import Test.QuickCheck (Property, ioProperty)
+
+-- | 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

containers-tests/tests/intmap-strictness.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE CPP #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
 module Main (main) where
@@ -6,6 +7,9 @@ import Test.ChasingBottoms.IsBottom
 import Test.Framework (Test, TestName, defaultMain, testGroup)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
 import Test.QuickCheck (Arbitrary(arbitrary))
+#if __GLASGOW_HASKELL__ >= 806
+import Test.QuickCheck (Property)
+#endif
 import Test.QuickCheck.Function (Fun(..), apply)
 import Test.Framework.Providers.HUnit
 import Test.HUnit hiding (Test)
@@ -16,6 +20,9 @@ import qualified Data.IntMap as L
 import Data.Containers.ListUtils
 
 import Utils.IsUnit
+#if __GLASGOW_HASKELL__ >= 806
+import Utils.NoThunks
+#endif
 
 instance Arbitrary v => Arbitrary (IntMap v) where
     arbitrary = M.fromList `fmap` arbitrary
@@ -101,6 +108,16 @@ pFromAscListStrict ks
   where
     elems = [(k, v) | k <- nubInt ks, v <- [undefined, undefined, ()]]
 
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldr' :: IntMap Int -> Property
+pStrictFoldr' m = whnfHasNoThunks (M.foldr' (:) [] m)
+#endif
+
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldl' :: IntMap Int -> Property
+pStrictFoldl' m = whnfHasNoThunks (M.foldl' (flip (:)) [] m)
+#endif
+
 ------------------------------------------------------------------------
 -- check for extra thunks
 --
@@ -184,6 +201,10 @@ tests =
         pInsertLookupWithKeyValueStrict
       , testProperty "fromAscList is somewhat value-lazy" pFromAscListLazy
       , testProperty "fromAscList is somewhat value-strict" pFromAscListStrict
+#if __GLASGOW_HASKELL__ >= 806
+      , testProperty "strict foldr'" pStrictFoldr'
+      , testProperty "strict foldl'" pStrictFoldl'
+#endif
       ]
       , tExtraThunksM
       , tExtraThunksL

containers-tests/tests/intset-strictness.hs

@@ -1,13 +1,34 @@
+{-# LANGUAGE CPP #-}
+#if __GLASGOW_HASKELL__ >= 800
+{-# OPTIONS_GHC -Wno-orphans #-}
+#endif
 module Main (main) where
 
 import Prelude hiding (foldl)
 
 import Test.ChasingBottoms.IsBottom
 import Test.Framework (Test, defaultMain, testGroup)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
+import Test.QuickCheck (Arbitrary (..))
+#if __GLASGOW_HASKELL__ >= 806
+import Test.QuickCheck (Property)
+#endif
 
 import Data.IntSet
 
+#if __GLASGOW_HASKELL__ >= 806
+import Utils.NoThunks
+#endif
+
+
+{--------------------------------------------------------------------
+  Arbitrary, reasonably balanced trees
+--------------------------------------------------------------------}
+instance Arbitrary IntSet where
+  arbitrary = do{ xs <- arbitrary
+                ; return (fromList xs)
+                }
+
 ------------------------------------------------------------------------
 -- * Properties
 
@@ -18,6 +39,16 @@ pFoldlAccLazy :: Int -> Bool
 pFoldlAccLazy k =
   isn'tBottom $ foldl (\_ x -> x) (bottom :: Int) (singleton k)
 
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldr' :: IntSet -> Property
+pStrictFoldr' m = whnfHasNoThunks (foldr' (:) [] m)
+#endif
+
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldl' :: IntSet -> Property
+pStrictFoldl' m = whnfHasNoThunks (foldl' (flip (:)) [] m)
+#endif
+
 ------------------------------------------------------------------------
 -- * Test list
 
@@ -27,6 +58,10 @@ tests =
     -- Basic interface
       testGroup "IntSet"
       [ testProperty "foldl is lazy in accumulator" pFoldlAccLazy
+#if __GLASGOW_HASKELL__ >= 806
+      , testProperty "strict foldr'" pStrictFoldr'
+      , testProperty "strict foldl'" pStrictFoldl'
+#endif
       ]
     ]
 

containers-tests/tests/map-strictness.hs

@@ -1,3 +1,5 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP          #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
 module Main (main) where
@@ -6,6 +8,9 @@ import Test.ChasingBottoms.IsBottom
 import Test.Framework (Test, TestName, defaultMain, testGroup)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
 import Test.QuickCheck (Arbitrary(arbitrary))
+#if __GLASGOW_HASKELL__ >= 806
+import Test.QuickCheck (Property)
+#endif
 import Test.QuickCheck.Function (Fun(..), apply)
 import Test.Framework.Providers.HUnit
 import Test.HUnit hiding (Test)
@@ -15,6 +20,9 @@ import qualified Data.Map.Strict as M
 import qualified Data.Map as L
 
 import Utils.IsUnit
+#if __GLASGOW_HASKELL__ >= 806
+import Utils.NoThunks
+#endif
 
 instance (Arbitrary k, Arbitrary v, Ord k) =>
          Arbitrary (Map k v) where
@@ -82,6 +90,26 @@ pInsertLookupWithKeyValueStrict f k v m
                      not (isBottom $ M.insertLookupWithKey (const3 1) k bottom m)
     | otherwise    = isBottom $ M.insertLookupWithKey (apply3 f) k bottom m
 
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldr' :: Map Int Int -> Property
+pStrictFoldr' m = whnfHasNoThunks (M.foldr' (:) [] m)
+#endif
+
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldl' :: Map Int Int -> Property
+pStrictFoldl' m = whnfHasNoThunks (M.foldl' (flip (:)) [] m)
+#endif
+
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldrWithKey' :: Map Int Int -> Property
+pStrictFoldrWithKey' m = whnfHasNoThunks (M.foldrWithKey' (\_ a as -> a : as) [] m)
+#endif
+
+#if __GLASGOW_HASKELL__ >= 806
+pStrictFoldlWithKey' :: Map Int Int -> Property
+pStrictFoldlWithKey' m = whnfHasNoThunks (M.foldlWithKey' (\as _ a -> a : as) [] m)
+#endif
+
 ------------------------------------------------------------------------
 -- check for extra thunks
 --
@@ -162,6 +190,12 @@ tests =
         pInsertLookupWithKeyKeyStrict
       , testProperty "insertLookupWithKey is value-strict"
         pInsertLookupWithKeyValueStrict
+#if __GLASGOW_HASKELL__ >= 806
+      , testProperty "strict foldr'" pStrictFoldr'
+      , testProperty "strict foldl'" pStrictFoldl'
+      , testProperty "strict foldrWithKey'" pStrictFoldrWithKey'
+      , testProperty "strict foldlWithKey'" pStrictFoldlWithKey'
+#endif
       ]
       , tExtraThunksM
       , tExtraThunksL
@@ -184,3 +218,4 @@ const2 x _ _ = x
 
 const3 :: a -> b -> c -> d -> a
 const3 x _ _ _ = x
+

containers-tests/tests/set-properties.hs

@@ -23,6 +23,10 @@ import Control.Applicative (Applicative (..), (<$>))
 #endif
 import Control.Applicative (liftA2)
 
+#if __GLASGOW_HASKELL__ >= 806
+import Utils.NoThunks (whnfHasNoThunks)
+#endif
+
 main :: IO ()
 main = defaultMain [ testCase "lookupLT" test_lookupLT
                    , testCase "lookupGT" test_lookupGT
@@ -104,6 +108,10 @@ main = defaultMain [ testCase "lookupLT" test_lookupLT
                    , testProperty "powerSet"             prop_powerSet
                    , testProperty "cartesianProduct"     prop_cartesianProduct
                    , testProperty "disjointUnion"        prop_disjointUnion
+#if __GLASGOW_HASKELL__ >= 806
+                   , testProperty "strict foldr"         prop_strictFoldr'
+                   , testProperty "strict foldr"         prop_strictFoldl'
+#endif
                    ]
 
 -- A type with a peculiar Eq instance designed to make sure keys
@@ -690,3 +698,13 @@ prop_disjointUnion xs ys =
 isLeft :: Either a b -> Bool
 isLeft (Left _) = True
 isLeft _ = False
+
+#if __GLASGOW_HASKELL__ >= 806
+prop_strictFoldr' :: Set Int -> Property
+prop_strictFoldr' m = whnfHasNoThunks (foldr' (:) [] m)
+#endif
+
+#if __GLASGOW_HASKELL__ >= 806
+prop_strictFoldl' :: Set Int -> Property
+prop_strictFoldl' m = whnfHasNoThunks (foldl' (flip (:)) [] m)
+#endif

containers/src/Data/Map/Internal.hs

@@ -3274,8 +3274,8 @@ foldr f z = go z
 foldr' :: (a -> b -> b) -> b -> Map k a -> b
 foldr' f z = go z
   where
-    go !z' Tip             = z'
-    go z' (Bin _ _ x l r) = go (f x (go z' r)) l
+    go !z' Tip            = z'
+    go z' (Bin _ _ x l r) = go (f x $! go z' r) l
 {-# INLINE foldr' #-}
 
 -- | /O(n)/. Fold the values in the map using the given left-associative
@@ -3300,8 +3300,10 @@ foldl f z = go z
 foldl' :: (a -> b -> a) -> a -> Map k b -> a
 foldl' f z = go z
   where
-    go !z' Tip             = z'
-    go z' (Bin _ _ x l r) = go (f (go z' l) x) r
+    go !z' Tip            = z'
+    go z' (Bin _ _ x l r) =
+      let !z'' = go z' l
+      in go (f z'' x) r
 {-# INLINE foldl' #-}
 
 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
@@ -3328,7 +3330,7 @@ foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b
 foldrWithKey' f z = go z
   where
     go !z' Tip              = z'
-    go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l
+    go z' (Bin _ kx x l r) = go (f kx x $! go z' r) l
 {-# INLINE foldrWithKey' #-}
 
 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
@@ -3354,8 +3356,10 @@ foldlWithKey f z = go z
 foldlWithKey' :: (a -> k -> b -> a) -> a -> Map k b -> a
 foldlWithKey' f z = go z
   where
-    go !z' Tip              = z'
-    go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r
+    go !z' Tip             = z'
+    go z' (Bin _ kx x l r) =
+      let !z'' = go z' l
+      in go (f z'' kx x) r
 {-# INLINE foldlWithKey' #-}
 
 -- | /O(n)/. Fold the keys and values in the map using the given monoid, such that

containers/src/Data/Set/Internal.hs

@@ -991,7 +991,7 @@ foldr' :: (a -> b -> b) -> b -> Set a -> b
 foldr' f z = go z
   where
     go !z' Tip           = z'
-    go z' (Bin _ x l r) = go (f x (go z' r)) l
+    go z' (Bin _ x l r) = go (f x $! go z' r) l
 {-# INLINE foldr' #-}
 
 -- | /O(n)/. Fold the elements in the set using the given left-associative
@@ -1014,7 +1014,9 @@ foldl' :: (a -> b -> a) -> a -> Set b -> a
 foldl' f z = go z
   where
     go !z' Tip           = z'
-    go z' (Bin _ x l r) = go (f (go z' l) x) r
+    go z' (Bin _ x l r) =
+      let !z'' = go z' l
+      in go (f z'' x) r
 {-# INLINE foldl' #-}
 
 {--------------------------------------------------------------------