{- Inspired by Seres, Spivey, Hoare, "Algebra of Logic Programming,"
   Proc. Int. Conf. on Logic Programming, Nov. 1999, pp. 184-199.
-}

module Forest (Forest,
	       Tree,
	       lift,
	       dfs,
	       bfs)
where
	              
import Prelude 
import Monad

newtype Forest a = Forest [Tree a]
     deriving Show
data Tree a = Leaf a | Fork (Forest a)
     deriving Show

instance Monad Forest where
  m >>= k = forestjoin (forestmap k m)
  return a = Forest [Leaf a]
instance MonadPlus Forest where
  mzero = Forest []
  (Forest m1) `mplus` (Forest m2) = Forest (m1 ++ m2)

lift :: Forest a -> Forest a
lift f = Forest [Fork f]

forestjoin :: Forest (Forest a) -> Forest a
forestjoin (Forest ts) = Forest (concat (map join' ts))
  where join' :: Tree (Forest a) -> [Tree a]
	join' (Leaf (Forest ts)) = ts
	join' (Fork xff) = [Fork (forestjoin xff)]

treemap :: (a -> b) -> Tree a -> Tree b
treemap f (Leaf x) = Leaf (f x)
treemap f (Fork xf) = Fork (forestmap f xf)

forestmap :: (a -> b) -> Forest a -> Forest b
forestmap f (Forest ts) = Forest (map (treemap f) ts)

dfs :: Forest a -> [a]
dfs (Forest ts) = concat (map dfs' ts)
 where dfs' :: Tree a -> [a]
       dfs' (Leaf x) = [x]
       dfs' (Fork xf) = dfs xf

bfs :: Forest a -> [a]
bfs (Forest ts) = concat (bfs' ts)
 where bfs' :: [Tree a]  -> [[a]]
       bfs' ts = combine (map levels ts)

       levels :: Tree a -> [[a]]
       levels (Leaf x) = [[x]]
       levels (Fork (Forest xf)) = []:bfs' xf

       combine :: [[[a]]] -> [[a]]
       combine = foldr merge []
 
       merge :: [[a]] -> [[a]] -> [[a]]
       merge (x:xs) (y:ys) = (x ++ y):(merge xs ys)
       merge xs [] = xs
       merge [] ys = ys