breadthFirst.hs

import Data.Tree

-- | Lists of nodes at each level of the tree.
levels :: Tree a -> [[a]]
levels (Node x xs) = [x] : levelsForest xs

-- | Lists of nodes at each level of the forest.
levelsForest :: Forest a -> [[a]]
levelsForest ts = foldl f b ts [] []
  where
    f k (Node x xs) ls qs = k (x : ls) (xs : qs)

    b _ [] = []
    b k qs = k : foldl (foldl f) b qs [] []

-- | Perform a breadth-first traversal of the tree.
breadthFirst :: Tree a -> [a]
breadthFirst (Node x xs) = x : breadthFirstForest xs

-- | Perform a breadth-first traversal of the forest.
breadthFirstForest :: Forest a -> [a]
breadthFirstForest ts = foldr f b ts []
  where
    f (Node x xs) fw bw = x : fw (xs : bw)

    b [] = []
    b qs = foldl (foldr f) b qs []

-- | Monadic forest builder, in breadth-first order.
unfoldForestM_BF :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a)
unfoldForestM_BF = unfoldForestMWith_BF concat

-- | Monadic tree builder, in breadth-first order.
unfoldTreeM_BF :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a)
unfoldTreeM_BF f b = unfoldForestMWith_BF (head . head) f [b]

unfoldForestMWith_BF :: Monad m => ([Forest a] -> c) -> (b -> m (a, [b])) -> [b] -> m c
unfoldForestMWith_BF r f ts = b [ts] (\ls -> r . ls)
  where
    b [] k = pure (k id [])
    b qs k = foldl g b qs [] (\ls -> k id . ls)

    g a xs qs k = foldr t (\ls ys -> a ys (k . run ls)) xs [] qs

    t a fw xs bw = f a >>= \(x,cs) -> fw (x:xs) (cs:bw)

    run x xs = uncurry (:) . foldl go ((,) [] . xs) x
      where
        go ys y (z:zs) = (Node y z : ys', zs')
          where
            (ys',zs') = ys zs
        go _ _ [] = errorWithoutStackTrace "unfoldTreeM_BF: bug!"