BFS tree fold

bf-fold-tree-strict.hs

#!/usr/bin/env stack
-- stack exec ghc --resolver lts-16.2 --package criterion -- -threaded -O2 -rtsopts -with-rtsopts=-N -eventlog

import           Control.Monad  (replicateM_, void)
import           Criterion.Main
import           Data.Foldable  (foldl')
import           Data.Tree

treeFoldStrict :: (b -> [a] -> b) -> b -> Tree a -> b
treeFoldStrict f init (Node val children) =
    seq acc $ treeFoldChildren f acc children
    where
        acc = f init [val]

treeFoldChildren :: (b -> [a] -> b) -> b -> [Tree a] -> b
treeFoldChildren f b children =
    seq acc $ treeFoldChildren f acc (concat branches)
    where
        acc = f b childValues
        childValues = fmap getNodeValue children
        branches = fmap getBranches children
        getNodeValue (Node v _) = v
        getBranches (Node _ b) = b


main :: IO ()
main = do
    print $ treeFoldStrict addOneToMany 0 tree
    where
        tree :: Tree Int
        tree = Node 1 [Node 2 [], Node 3 []]
        addOneToMany = foldl' (+)

TODO benchmark