module RegAlloc.Data.Trees (
    -- * Trees
    -- ** Type
    Tree,

    -- ** Construction
    tree,

    -- ** Properties
    root,
    subtrees,

    -- ** Operations
    downAccuTree,
    upAccuTree,
    flattenTree,
    treeLevels,

    -- * Forests
    -- ** Type
    Forest,

    -- ** Construction
    forest,
    emptyForest,
    oneTreeForest,

    -- ** Properties
    trees,

    -- ** Operations
    mapTrees,
    downAccuForest,
    upAccuForest,
    flattenForest,
    forestLevels
) where
    -- * Trees
    -- ** Type
    data Tree node = Tree node (Forest node) deriving Eq

    instance Functor Tree where
        fmap mapping (Tree root subtrees) = Tree (mapping root) (fmap mapping subtrees)

    -- ** Construction
    tree :: node -> Forest node -> Tree node
    tree = Tree

    -- ** Properties
    root :: Tree node -> node
    root (Tree root _) = root

    subtrees :: Tree node -> Forest node
    subtrees (Tree _ subtrees) = subtrees

    -- ** Operations
    -- runs in O(n)
    downAccuTree :: (node' -> node -> node') -> node' -> Tree node -> Tree node'
    downAccuTree modification initialValue (Tree root subtrees)
        = let
              root' = modification initialValue root
          in
          Tree root' (downAccuForest modification root' subtrees)

    -- runs in O(n?) or similar
    -- think it runs in O(n) when the children count of the nodes is greater than one
    upAccuTree :: (node -> node' -> node') -> node' -> Tree node -> Tree node'
    upAccuTree modification initialValue (Tree root subtrees)
        = fmap (modification root)
               (Tree initialValue (upAccuForest modification initialValue subtrees))

    flattenTree :: Tree node -> [node]
    flattenTree = flattenForest . oneTreeForest

    treeLevels :: Tree node -> [[node]]
    treeLevels = forestLevels . oneTreeForest

    -- * Forests
    -- ** Type
    newtype Forest node = Forest [Tree node] deriving Eq

    instance Functor Forest where
        fmap mapping (Forest trees) = Forest (map (fmap mapping) trees)

    -- ** Construction
    forest :: [Tree node] -> Forest node
    forest = Forest

    emptyForest :: Forest node
    emptyForest = Forest []

    oneTreeForest :: Tree node -> Forest node
    oneTreeForest = Forest . return

    -- ** Properties
    trees :: Forest node -> [Tree node]
    trees (Forest trees) = trees

    -- ** Operations
    mapTrees :: (Tree node -> Tree node') -> Forest node -> Forest node'
    mapTrees mapping (Forest trees) = Forest (map mapping trees)

    downAccuForest :: (node' -> node -> node') -> node' -> Forest node -> Forest node'
    downAccuForest modification initialValue = mapTrees (downAccuTree modification initialValue)

    upAccuForest :: (node -> node' -> node') -> node' -> Forest node -> Forest node'
    upAccuForest modification initialValue = mapTrees (upAccuTree modification initialValue)

    flattenForest :: Forest node -> [node]
    flattenForest = concat . forestLevels

    forestLevels :: Forest node -> [[node]]
    forestLevels = levels . trees
                 where
                     levels :: [Tree node] -> [[node]]
                     levels treeList = map root treeList :
                                       levels (concatMap (trees . subtrees) treeList)