Add performant Foldable functions for Data.Tree

Implement Foldable functions for Data.Tree, which only had the bare
minimum foldMap so far. Add tests and benchmarks for the added
functions. Benchmarks show very good improvements.

Author: meooow25

containers-tests/benchmarks/Tree.hs

@@ -0,0 +1,53 @@
+module Main where
+
+import Control.DeepSeq (NFData, rnf)
+import Control.Exception (evaluate)
+import Data.Coerce (coerce)
+import Data.Foldable (fold, foldl', toList)
+import Data.Monoid (All(..))
+import Test.Tasty.Bench (Benchmark, Benchmarkable, bench, bgroup, defaultMain, whnf, nf)
+import qualified Data.Tree as T
+
+main :: IO ()
+main = do
+  evaluate $ rnf ts `seq` rnf tsBool
+  defaultMain
+    [ bgroup "fold" $ forTs tsBool $ whnf fold . (coerce :: T.Tree Bool -> T.Tree All)
+    , bgroup "foldMap" $ forTs tsBool $ whnf (foldMap All)
+    , bgroup "foldr_1" $ forTs tsBool $ whnf (foldr (&&) True)
+    , bgroup "foldr_2" $ forTs ts $ whnf (length . foldr (:) [])
+    , bgroup "foldr_3" $ forTs ts $ whnf (\t -> foldr (\x k acc -> if acc < 0 then acc else k $! acc + x) id t 0)
+    , bgroup "foldl'" $ forTs ts $ whnf (foldl' (+) 0)
+    , bgroup "foldr1" $ forTs tsBool $ whnf (foldr1 (&&))
+    , bgroup "foldl1" $ forTs ts $ whnf (foldl1 (+))
+    , bgroup "toList" $ forTs ts $ nf toList
+    , bgroup "elem" $ forTs ts $ whnf (elem 0)
+    , bgroup "maximum" $ forTs ts $ whnf maximum
+    , bgroup "sum" $ forTs ts $ whnf sum
+    ]
+  where
+    ts = [binaryTree, lineTree] <*> [1000, 1000000]
+    tsBool = [t { getT = True <$ getT t } | t <- ts]
+
+forTs :: [Tree a] -> (T.Tree a -> Benchmarkable) -> [Benchmark]
+forTs ts f = [bench label (f t) | Tree label t <- ts]
+
+data Tree a = Tree
+  { getLabel :: String
+  , getT :: T.Tree a
+  }
+
+instance NFData a => NFData (Tree a) where
+  rnf (Tree label t) = rnf label `seq` rnf t
+
+binaryTree :: Int -> Tree Int
+binaryTree n = Tree label t
+  where
+    label = "bin,n=" ++ show n
+    t = T.unfoldTree (\x -> (x, takeWhile (<=n) [2*x, 2*x+1])) 1
+
+lineTree :: Int -> Tree Int
+lineTree n = Tree label t
+  where
+    label = "line,n=" ++ show n
+    t = T.unfoldTree (\x -> (x, [x+1 | x+1 <= n])) 1

containers-tests/containers-tests.cabal

@@ -152,6 +152,14 @@ benchmark map-benchmarks
   main-is:        Map.hs
   ghc-options:    -O2
 
+benchmark tree-benchmarks
+  import: benchmark-deps
+  default-language: Haskell2010
+  type:           exitcode-stdio-1.0
+  hs-source-dirs: benchmarks
+  main-is:        Tree.hs
+  ghc-options:    -O2
+
 benchmark sequence-benchmarks
   import: benchmark-deps
   default-language: Haskell2010

containers-tests/tests/tree-properties.hs

@@ -5,23 +5,37 @@ import Data.Tree as T
 import Control.Applicative (Const(Const, getConst), pure, (<$>), (<*>), liftA2)
 
 import Test.Tasty
+import Test.Tasty.HUnit
 import Test.Tasty.QuickCheck
 import Test.QuickCheck.Function (apply)
-import Test.QuickCheck.Poly (A, B, C)
+import Test.QuickCheck.Poly (A, B, C, OrdA)
 import Control.Monad.Fix (MonadFix (..))
 import Control.Monad (ap)
+import Data.Foldable (foldl', toList)
+import Data.Traversable (foldMapDefault)
 
 default (Int)
 
 main :: IO ()
 main = defaultMain $ testGroup "tree-properties"
          [
-           testProperty "monad_id1"                prop_monad_id1
+           testCase     "foldr"                    test_foldr
+         , testProperty "monad_id1"                prop_monad_id1
          , testProperty "monad_id2"                prop_monad_id2
          , testProperty "monad_assoc"              prop_monad_assoc
          , testProperty "ap_ap"                    prop_ap_ap
          , testProperty "ap_liftA2"                prop_ap_liftA2
          , testProperty "monadFix_ls"              prop_monadFix_ls
+         , testProperty "toList"                   prop_toList
+         , testProperty "foldMap"                  prop_foldMap
+         , testProperty "foldl'"                   prop_foldl'
+         , testProperty "foldr1"                   prop_foldr1
+         , testProperty "foldl1"                   prop_foldl1
+         , testProperty "foldr_infinite"           prop_foldr_infinite
+         , testProperty "maximum"                  prop_maximum
+         , testProperty "minimum"                  prop_minimum
+         , testProperty "sum"                      prop_sum
+         , testProperty "product"                  prop_product
          ]
 
 {--------------------------------------------------------------------
@@ -52,10 +66,25 @@ instance Arbitrary a => Arbitrary (Tree a) where
   shrink = genericShrink
 #endif
 
+----------------------------------------------------------------
+-- Utilities
+----------------------------------------------------------------
+
+data Magma a
+  = Inj a
+  | Magma a :* Magma a
+  deriving (Eq, Show)
+
 ----------------------------------------------------------------
 -- Unit tests
 ----------------------------------------------------------------
 
+test_foldr :: Assertion
+test_foldr = do
+  foldr (:) [] (Node 1 []) @?= [1]
+  foldr (:) [] (Node 1 [Node 2 [Node 3 []]]) @?= [1..3]
+  foldr (:) [] (Node 1 [Node 2 [Node 3 [], Node 4 []], Node 5 [Node 6 [], Node 7 []]]) @?= [1..7]
+
 ----------------------------------------------------------------
 -- QuickCheck
 ----------------------------------------------------------------
@@ -101,3 +130,39 @@ prop_monadFix_ls val ta ti =
     f :: (Int -> Int) -> Int -> Tree (Int -> Int)
     f q y = let t = apply ti y
             in fmap (\w -> fact w q) t
+
+prop_toList :: Tree A -> Property
+prop_toList t = toList t === foldr (:) [] t
+
+prop_foldMap :: Tree A -> Property
+prop_foldMap t =
+  foldMap (:[]) t === toList t .&&.
+  foldMap (:[]) t === foldMapDefault (:[]) t
+
+prop_foldl' :: Tree A -> Property
+prop_foldl' t = foldl' (flip (:)) [] t === reverse (toList t)
+
+prop_foldr1 :: Tree A -> Property
+prop_foldr1 t = foldr1 (:*) (fmap Inj t) === foldr1 (:*) (map Inj (toList t))
+
+prop_foldl1 :: Tree A -> Property
+prop_foldl1 t = foldl1 (:*) (fmap Inj t) === foldl1 (:*) (map Inj (toList t))
+
+prop_foldr_infinite :: NonNegative Int -> Property
+prop_foldr_infinite (NonNegative n) =
+    forAllShow genInf (const "<possibly infinite tree>") $
+        \t -> length (take n (foldr (:) [] t)) <= n
+  where
+    genInf = Node () <$> oneof [listOf genInf, infiniteListOf genInf]
+
+prop_maximum :: Tree OrdA -> Property
+prop_maximum t = maximum t === maximum (toList t)
+
+prop_minimum :: Tree OrdA -> Property
+prop_minimum t = minimum t === minimum (toList t)
+
+prop_sum :: Tree OrdA -> Property
+prop_sum t = sum t === sum (toList t)
+
+prop_product :: Tree OrdA -> Property
+prop_product t = product t === product (toList t)

containers/src/Data/Tree.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE PatternGuards #-}
 {-# LANGUAGE CPP #-}
 #if __GLASGOW_HASKELL__
@@ -52,7 +53,8 @@ module Data.Tree(
 
     ) where
 
-import Data.Foldable (toList)
+import Data.Foldable (fold, foldl', toList)
+import Data.Traversable (foldMapDefault)
 import Control.Monad (liftM)
 import Control.Monad.Fix (MonadFix (..), fix)
 import Data.Sequence (Seq, empty, singleton, (<|), (|>), fromList,
@@ -179,16 +181,59 @@ mfixTree f
                     [0..] children)
 
 instance Traversable Tree where
-    traverse f (Node x ts) = liftA2 Node (f x) (traverse (traverse f) ts)
+  traverse f = go
+    where go (Node x ts) = liftA2 Node (f x) (traverse go ts)
+  {-# INLINE traverse #-}
 
+-- | Folds in preorder
+
+-- See Note [Implemented Foldable Tree functions]
 instance Foldable Tree where
-    foldMap f (Node x ts) = f x `mappend` foldMap (foldMap f) ts
+    fold = foldMap id
+    {-# INLINABLE fold #-}
+
+    foldMap = foldMapDefault
+    {-# INLINE foldMap #-}
+
+    foldr f z = \t -> go t z  -- Use a lambda to allow inlining with two arguments
+      where
+        go (Node x ts) = f x . foldr (\t k -> go t . k) id ts
+        -- This is equivalent to the following simpler definition, but has been found to optimize
+        -- better in benchmarks:
+        -- go (Node x ts) z' = f x (foldr go z' ts)
+    {-# INLINE foldr #-}
+
+    foldl' f = go
+      where go !z (Node x ts) = foldl' go (f z x) ts
+    {-# INLINE foldl' #-}
+
+    foldr1 f = go id
+      where go k (Node x ts) = foldr (\n k' prev -> f prev (go k' n)) k ts x
+    {-# INLINE foldr1 #-}
+
+    foldl1 f (Node x ts) = foldl (foldl f) x ts
 
     null _ = False
     {-# INLINE null #-}
 
-    toList = flatten
-    {-# INLINE toList #-}
+    elem = any . (==)
+    {-# INLINABLE elem #-}
+
+    maximum = foldl1' max
+    {-# INLINABLE maximum #-}
+
+    minimum = foldl1' min
+    {-# INLINABLE minimum #-}
+
+    sum = foldl1' (+)
+    {-# INLINABLE sum #-}
+
+    product = foldl1' (*)
+    {-# INLINABLE product #-}
+
+foldl1' :: (a -> a -> a) -> Tree a -> a
+foldl1' f = \(Node x ts) -> foldl' (foldl' f) x ts
+{-# INLINE foldl1' #-}
 
 instance NFData a => NFData (Tree a) where
     rnf (Node x ts) = rnf x `seq` rnf ts
@@ -262,8 +307,7 @@ draw (Node x ts0) = lines x ++ drawSubTrees ts0
 --
 -- > flatten (Node 1 [Node 2 [], Node 3 []]) == [1,2,3]
 flatten :: Tree a -> [a]
-flatten t = squish t []
-  where squish (Node x ts) xs = x:Prelude.foldr squish xs ts
+flatten = toList
 
 -- | Returns the list of nodes at each level of the tree.
 --
@@ -415,3 +459,23 @@ unfoldForestQ f aQ = case viewl aQ of
     splitOnto as (_:bs) q = case viewr q of
         q' :> a -> splitOnto (a:as) bs q'
         EmptyR -> error "unfoldForestQ"
+
+--------------------------------------------------------------------------------
+
+-- Note [Implemented Foldable Tree functions]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- Implemented:
+--
+-- foldMap, foldr, foldl': Basic functions.
+-- fold, elem: Implemented same as the default definition, but INLINABLE to
+-- allow specialization.
+-- foldr1, foldl1, null, maximum, minimum: Implemented more efficiently than
+-- defaults since trees are non-empty.
+-- sum, product: Implemented as strict left folds. Defaults use the lazy foldMap
+-- before base 4.15.1.
+--
+-- Not implemented:
+--
+-- foldMap', toList, length: Defaults perform well.
+-- foldr', foldl: Unlikely to be used.