Skip to content

A Tree fold does not optimize very well

Summary

The following fold function on a rose tree does not optimize very well.

data Tree a = Node a [Tree a]

foldrTree :: (a -> b -> b) -> b -> Tree a -> b
foldrTree f z t = go t z
  where go (Node x ts) z' = f x (foldr go z' ts)

But some more complex variations of it do. An example with benchmarks on -O2:

import Control.DeepSeq
import Criterion.Main

main :: IO ()
main = defaultMain
  [ env (pure t_) $ \t -> bgroup ""
    [ bench "foldrTree" $ whnf (foldrTree (&&) True) t
    , bench "foldrTreeAlt" $ whnf (foldrTreeAlt (&&) True) t
    ]
  ]
  where
    t_ = Node True [Node True [] | _ <- [1..1000000]]

data Tree a = Node a [Tree a]

instance NFData a => NFData (Tree a) where
  rnf (Node x ts) = rnf x `seq` rnf ts

foldrTree :: (a -> b -> b) -> b -> Tree a -> b
foldrTree f z t = go t z
  where go (Node x ts) z' = f x (foldr go z' ts)

foldrTreeAlt :: (a -> b -> b) -> b -> Tree a -> b
foldrTreeAlt f z t = go t z
  where go (Node x ts) = f x . foldr ((.) . go) id ts
benchmarking foldrTree
time                 11.60 ms   (11.50 ms .. 11.70 ms)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 11.61 ms   (11.58 ms .. 11.65 ms)
std dev              93.16 μs   (76.81 μs .. 111.6 μs)

benchmarking foldrTreeAlt
time                 5.890 ms   (5.869 ms .. 5.910 ms)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 5.943 ms   (5.928 ms .. 5.969 ms)
std dev              57.94 μs   (43.70 μs .. 88.66 μs)

For background, this was encountered in https://github.com/haskell/containers/pull/880

Expected behavior

The simpler function should ideally optimize as well as the other one.

Environment

  • GHC version used: 9.2.5
To upload designs, you'll need to enable LFS and have an admin enable hashed storage. More information