module Data.Tree.Extended (
      traverseTree
    , depthFirst
    , breadthFirst
    , module Data.Tree
) where



import           Data.List
import qualified Data.Sequence as Seq
import           Data.Tree



traverseTree
    :: (Tree a -> mem -> mem)       -- ^ Put a tree onto the to-visit list
    -> (mem -> Maybe (Tree a, mem)) -- ^ Get a tree from the to-visit list
    -> mem                          -- ^ To-visit list; initially the first node of the tree
    -> [a]
traverseTree :: forall a mem.
(Tree a -> mem -> mem)
-> (mem -> Maybe (Tree a, mem)) -> mem -> [a]
traverseTree Tree a -> mem -> mem
put mem -> Maybe (Tree a, mem)
get mem
memory = case mem -> Maybe (Tree a, mem)
get mem
memory of
    Maybe (Tree a, mem)
Nothing -> []
    Just (Node a
x [Tree a]
children, mem
rest) ->
        a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (Tree a -> mem -> mem)
-> (mem -> Maybe (Tree a, mem)) -> mem -> [a]
forall a mem.
(Tree a -> mem -> mem)
-> (mem -> Maybe (Tree a, mem)) -> mem -> [a]
traverseTree Tree a -> mem -> mem
put mem -> Maybe (Tree a, mem)
get ((mem -> Tree a -> mem) -> mem -> [Tree a] -> mem
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\mem
acc Tree a
y -> Tree a -> mem -> mem
put Tree a
y mem
acc) mem
rest [Tree a]
children)

depthFirst :: Tree a -> [a]
depthFirst :: forall a. Tree a -> [a]
depthFirst (Node a
x [Tree a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: ([Tree a]
xs [Tree a] -> (Tree a -> [a]) -> [a]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Tree a -> [a]
forall a. Tree a -> [a]
depthFirst)

-- | Breadth-first 'toList'.
breadthFirst :: Tree a -> [a]
breadthFirst :: forall a. Tree a -> [a]
breadthFirst Tree a
tree = (Tree a -> Seq (Tree a) -> Seq (Tree a))
-> (Seq (Tree a) -> Maybe (Tree a, Seq (Tree a)))
-> Seq (Tree a)
-> [a]
forall a mem.
(Tree a -> mem -> mem)
-> (mem -> Maybe (Tree a, mem)) -> mem -> [a]
traverseTree Tree a -> Seq (Tree a) -> Seq (Tree a)
forall {a}. a -> Seq a -> Seq a
putLeft Seq (Tree a) -> Maybe (Tree a, Seq (Tree a))
forall {a}. Seq a -> Maybe (a, Seq a)
getRight (Tree a -> Seq (Tree a) -> Seq (Tree a)
forall {a}. a -> Seq a -> Seq a
putLeft Tree a
tree Seq (Tree a)
forall a. Monoid a => a
mempty)
  where
    putLeft :: a -> Seq a -> Seq a
putLeft = a -> Seq a -> Seq a
forall {a}. a -> Seq a -> Seq a
(Seq.<|)
    getRight :: Seq a -> Maybe (a, Seq a)
getRight Seq a
xs = case Seq a -> ViewR a
forall a. Seq a -> ViewR a
Seq.viewr Seq a
xs of
        ViewR a
Seq.EmptyR -> Maybe (a, Seq a)
forall a. Maybe a
Nothing
        Seq a
ys Seq.:> a
y -> (a, Seq a) -> Maybe (a, Seq a)
forall a. a -> Maybe a
Just (a
y, Seq a
ys)