-- | Abstracts away tree traversals.
-- Mostly used by callers including (soon) XML Conduit Stylist,
-- but also used internally for generating counter text.
module Stylist.Tree(StyleTree(..), treeOrder, treeOrder',
    Path, treeMap, treeFind, treeFlatten, treeFlattenAll, preorder, preorder', postorder) where

data StyleTree p = StyleTree {
    style :: p,
    children :: [StyleTree p]
}

type Path = [Integer]
treeOrder :: (c -> c -> Path -> p -> (c, p')) ->
    c -> StyleTree p -> StyleTree p'
treeOrder cb ctxt tree = StyleTree
    (snd $ cb ctxt ctxt [] $ style tree)
    (snd $ treeOrder' cb ctxt ctxt [0] $ children tree)
treeOrder' :: (c -> c -> Path -> p -> (c, p')) ->
    c -> c -> Path -> [StyleTree p] -> (c, [StyleTree p'])
treeOrder' cb prevContext context (num:path) (node:nodes) = (tailContext, StyleTree node' children' : nodes')
    where
        (selfContext, node') = cb prevContext context (num:path) $ style node
        (childContext, children') = treeOrder' cb selfContext selfContext (0:num:path) $ children node
        (tailContext, nodes') = treeOrder' cb selfContext childContext (num + 1:path) nodes
treeOrder' _ _ context _ [] = (context, [])
treeOrder' _ _ _ [] _ = error "Invalid path during tree traversal!"

treeMap :: (p -> p') -> StyleTree p -> StyleTree p'
treeMap cb = treeOrder (\_ _ _ p -> ((), cb p)) ()

treeFlatten :: StyleTree p -> [p]
treeFlatten = treeFlatten' . children
treeFlatten' :: [StyleTree p] -> [p]
treeFlatten' (StyleTree p []:ps) = p : treeFlatten' ps
treeFlatten' (StyleTree _ childs:sibs) = treeFlatten' childs ++ treeFlatten' sibs
treeFlatten' [] = []

treeFlattenAll :: StyleTree p -> [p]
treeFlattenAll = treeFlattenAll' . children
treeFlattenAll' :: [StyleTree p] -> [p]
treeFlattenAll' (StyleTree p []:ps) = p : treeFlattenAll' ps
treeFlattenAll' (StyleTree p childs:sibs) = p : treeFlattenAll' childs ++ treeFlattenAll' sibs
treeFlattenAll' [] = []

treeFind :: StyleTree p -> (p -> Bool) -> [p]
treeFind p test = filter test $ treeFlattenAll p

preorder :: (Maybe b -> Maybe b -> a -> b) -> StyleTree a -> StyleTree b
preorder cb self = head $ preorder' cb Nothing Nothing [self]
preorder' :: (Maybe b -> Maybe b -> a -> b) ->
        Maybe b -> Maybe b -> [StyleTree a] -> [StyleTree b]
preorder' cb parent previous (self:sibs) = let self' = cb parent previous $ style self
        in StyleTree self' (preorder' cb (Just self') Nothing $ children self) :
            preorder' cb parent (Just self') sibs
preorder' _ _ _ [] = []

postorder :: (a -> [b] -> [b]) -> StyleTree a -> [StyleTree b]
postorder cb (StyleTree self childs) =
    [StyleTree self' children' | self' <- cb self $ Prelude.map style children']
  where children' = concat $ Prelude.map (postorder cb) childs