Merge pull request #880 from meooow25/tree

Faster Tree folds

Author: treeowl

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.