Implement foldl-related folds for PostOrder

containers-tests/benchmarks/Tree.hs

@@ -29,7 +29,7 @@ main = do
         | Tree label t <- ts
         ]
       , bgroup "foldr1" $ forTs tsBool $ whnf (foldr1 (&&))
-      , bgroup "foldl1" $ forTs ts $ whnf (foldl1 (+))
+      , bgroup "foldl1" $ forTs tsBool $ whnf (foldl1 (flip (&&)))
       , bgroup "toList" $ forTs ts $ nf F.toList
       , bgroup "elem" $ forTs ts $ whnf (elem 0)
       , bgroup "maximum" $ forTs ts $ whnf maximum
@@ -46,7 +46,7 @@ main = do
       , bgroup "foldrMap1_1" $ forTs tsBool $ whnf (Foldable1.foldrMap1 id (&&))
       , bgroup "foldrMap1_2" $ forTs ts $ whnf (length . Foldable1.foldrMap1 (:[]) (:))
       , bgroup "foldlMap1'" $ forTs ts $ whnf (Foldable1.foldlMap1' id (+))
-      , bgroup "foldlMap1" $ forTs ts $ whnf (Foldable1.foldlMap1 id (+))
+      , bgroup "foldlMap1" $ forTs tsBool $ whnf (Foldable1.foldlMap1 id (flip (&&)))
       ]
 #endif
     , bgroup "leaves" $ forTs ts $ nf T.leaves
@@ -61,7 +61,7 @@ main = do
           | Tree label t <- ts
           ]
         , bgroup "foldr1" $ forPostOrders tsBool $ whnf (foldr1 (&&))
-        , bgroup "foldl1" $ forPostOrders ts $ whnf (foldl1 (+))
+        , bgroup "foldl1" $ forPostOrders tsBool $ whnf (foldl1 (flip (&&)))
         , bgroup "toList" $ forPostOrders ts $ nf F.toList
         , bgroup "elem" $ forPostOrders ts $ whnf (elem 0)
         , bgroup "maximum" $ forPostOrders ts $ whnf maximum
@@ -77,7 +77,7 @@ main = do
         , bgroup "foldrMap1_1" $ forPostOrders tsBool $ whnf (Foldable1.foldrMap1 id (&&))
         , bgroup "foldrMap1_2" $ forPostOrders ts $ whnf (length . Foldable1.foldrMap1 (:[]) (:))
         , bgroup "foldlMap1'" $ forPostOrders ts $ whnf (Foldable1.foldlMap1' id (+))
-        , bgroup "foldlMap1" $ forPostOrders ts $ whnf (Foldable1.foldlMap1 id (+))
+        , bgroup "foldlMap1" $ forTs tsBool $ whnf (Foldable1.foldlMap1 id (flip (&&)))
         ]
 #endif
       ]

containers-tests/containers-tests.cabal

@@ -466,6 +466,7 @@ test-suite tree-properties
 
   build-depends:
       test-utils
+    , ChasingBottoms
 
   other-extensions:
     BangPatterns

containers-tests/tests/tree-properties.hs

@@ -10,9 +10,12 @@ import Test.Tasty.QuickCheck
 import Test.QuickCheck.Function (apply)
 import Test.QuickCheck.Poly (A, B, C, OrdA)
 import qualified Test.QuickCheck.Classes.Base as Laws
+import Test.ChasingBottoms.IsBottom (isBottom)
 import Control.Monad.Fix (MonadFix (..))
 import Control.Monad (ap)
+import Data.Coerce (coerce)
 import Data.Foldable (fold, foldl', toList)
+import qualified Data.Foldable as Foldable
 import qualified Data.List as List
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as NE
@@ -26,6 +29,7 @@ import qualified Data.Foldable1 as Foldable1
 #endif
 
 import Utils.QuickCheckClasses (testLaws)
+import Utils.Strictness (Bot(..), Func2, applyFunc2)
 
 default (Int)
 
@@ -108,6 +112,16 @@ main = defaultMain $ testGroup "tree-properties"
            , testLaws $ Laws.traversableLaws (Proxy :: Proxy PostOrder)
 #endif
            ]
+         , testGroup "strictness"
+           [ testProperty "foldr" prop_strictness_foldr
+           , testProperty "foldl'" prop_strictness_foldl'
+           , testGroup "PostOrder"
+             [ testProperty "foldr" prop_PostOrder_strictness_foldr
+             , testProperty "foldl" prop_PostOrder_strictness_foldl
+             , testProperty "foldr'" prop_PostOrder_strictness_foldr'
+             , testProperty "foldl'" prop_PostOrder_strictness_foldl'
+             ]
+           ]
          ]
 
 {--------------------------------------------------------------------
@@ -463,3 +477,58 @@ prop_PostOrder_foldlMap1 t =
   where
     f z x = z :* Inj x
 #endif
+
+-- * Strictness tests
+
+-- See Note [Testing strictness of folds] in map-strictness.hs
+
+prop_strictness_foldr :: Tree A -> Func2 A B (Bot B) -> Bot B -> Property
+prop_strictness_foldr t fun (Bot z) =
+  isBottom (Foldable.foldr f z t) ===
+  isBottom (Foldable.foldr f z (T.flatten t))
+  where
+    f = coerce (applyFunc2 fun) :: A -> B -> B
+
+prop_strictness_foldl' :: Tree A -> Func2 B A (Bot B) -> Bot B -> Property
+prop_strictness_foldl' t fun (Bot z) =
+  isBottom (Foldable.foldl' f z t) ===
+  isBottom (Foldable.foldl' f z (T.flatten t))
+  where
+    f = coerce (applyFunc2 fun) :: B -> A -> B
+
+prop_PostOrder_strictness_foldr
+  :: PostOrder A -> Func2 A B (Bot B) -> Bot B -> Property
+prop_PostOrder_strictness_foldr t fun (Bot z) =
+  isBottom (Foldable.foldr f z t) ===
+  isBottom (Foldable.foldr f z (toPostOrderList t))
+  where
+    f = coerce (applyFunc2 fun) :: A -> B -> B
+
+prop_PostOrder_strictness_foldl
+  :: PostOrder A -> Func2 B A (Bot B) -> Bot B -> Property
+prop_PostOrder_strictness_foldl t fun (Bot z) =
+  isBottom (Foldable.foldl f z t) ===
+  isBottom (Foldable.foldl f z (toPostOrderList t))
+  where
+    f = coerce (applyFunc2 fun) :: B -> A -> B
+
+prop_PostOrder_strictness_foldr'
+  :: PostOrder A -> Func2 A B (Bot B) -> Bot B -> Property
+prop_PostOrder_strictness_foldr' t fun (Bot z) =
+  isBottom (Foldable.foldr' f z t) ===
+  isBottom (z `seq` Foldable.foldr' f z (toPostOrderList t))
+  where
+    f = coerce (applyFunc2 fun) :: A -> B -> B
+
+prop_PostOrder_strictness_foldl'
+  :: PostOrder A -> Func2 B A (Bot B) -> Bot B -> Property
+prop_PostOrder_strictness_foldl' t fun (Bot z) =
+  isBottom (Foldable.foldl' f z t) ===
+  isBottom (Foldable.foldl' f z (toPostOrderList t))
+  where
+    f = coerce (applyFunc2 fun) :: B -> A -> B
+
+toPostOrderList :: PostOrder a -> [a]
+toPostOrderList (PostOrder t) = go t []
+  where
+    go (Node x ts) xs = foldr go (x:xs) ts

containers/src/Data/Tree.hs

@@ -706,6 +706,25 @@ toNonEmptyBQ (BQ x0 _ f r) = case r of
 
 -- | A newtype over 'Tree' that folds and traverses in post-order.
 --
+-- ==== __@Foldable@ examples__
+--
+-- >>> import Data.Foldable (toList)
+-- >>> toList $ PostOrder $ Node 1 [Node 2 [Node 3 [], Node 4 []], Node 5 []]
+-- [3,4,2,5,1]
+--
+-- @foldr@ produces elements incrementally, inspecting the structure of the
+-- @Tree@ just as much as necessary.
+--
+-- >>> take 3 $ foldr (:) [] $ PostOrder $ Node 1 ([Node 2 [Node 3 [], Node 4 []]] ++ undefined)
+-- [3,4,2]
+--
+-- @foldl@ also produces elements incrementally.
+--
+-- >>> foldl (flip (:)) [] $ PostOrder $ Node 1 [Node 2 [Node 3 [], Node 4 []], Node 5 []]
+-- [1,5,2,4,3]
+-- >>> take 4 $ foldl (flip (:)) [] $ PostOrder $ Node 1 [Node 2 [undefined, Node 4 []], Node 5 []]
+-- [1,5,2,4]
+--
 -- @since 0.8
 newtype PostOrder a = PostOrder { unPostOrder :: Tree a }
 #ifdef __GLASGOW_HASKELL__
@@ -747,9 +766,25 @@ instance Foldable PostOrder where
           in f z' x
     {-# INLINE foldl' #-}
 
+    foldl f z0 = -- Inline with two arguments
+      \(PostOrder t) -> go z0 t
+      where
+        go z (Node x ts) = f (Foldable.foldl go z ts) x
+    {-# INLINE foldl #-}
+
+    foldr' f z0 = -- Inline with two arguments
+      \(PostOrder t) -> go t z0
+      where
+        go (Node x ts) !z =
+          let !z' = f x z
+          in foldrTreeList go z' ts
+    {-# INLINE foldr' #-}
+
     foldr1 = foldrMap1PostOrder id
+    {-# INLINE foldr1 #-}
 
     foldl1 = foldlMap1PostOrder id
+    {-# INLINE foldl1 #-}
 
     null _ = False
     {-# INLINE null #-}
@@ -801,10 +836,18 @@ instance Foldable1.Foldable1 PostOrder where
   {-# INLINABLE minimum #-}
 
   foldrMap1 = foldrMap1PostOrder
+  {-# INLINE foldrMap1 #-}
 
   foldlMap1' = foldlMap1'PostOrder
+  {-# INLINE foldlMap1' #-}
 
   foldlMap1 = foldlMap1PostOrder
+  {-# INLINE foldlMap1 #-}
+
+  foldrMap1' f g = -- Inline with two arguments
+    \(PostOrder (Node x ts)) ->
+      foldr (\t !z -> Foldable.foldr' g z (PostOrder t)) (f x) ts
+  {-# INLINE foldrMap1' #-}
 #endif
 
 foldrMap1PostOrder :: (a -> b) -> (a -> b -> b) -> PostOrder a -> b