-- | Multi-way trees (also known as /rose trees/) and forests,
-- similar to @Data.Tree@ from the /containers/ library.
module Data.Forest
  ( -- * Importing
    -- $imports

    -- * Types
    Forest,
    Tree,

    -- * Constructing
    forest,
    tree,
    leaf,
    leaves,

    -- * Deconstructing
    trees,
    root,
    subforest,
    subtrees,

    -- * Folds
    foldForest,
    foldTree,

    -- * Forest functor
    -- $functor
  )
where

import Data.Eq (Eq)
import Data.Foldable (Foldable)
import Data.Function (($))
import Data.Functor (Functor, fmap, (<$>))
import Data.Monoid (Monoid, mempty)
import Data.Semigroup (Semigroup)
import Data.Traversable (Traversable)
import Prelude (Show)

-- | A forest is defined completely by its 'trees'.
--
-- To construct a forest, use 'forest' or 'leaves'.
newtype Forest a = Forest
  { -- | The trees that constitute the forest.
    trees :: [Tree a]
  }
  deriving (Eq, Show, Functor, Foldable, Traversable, Semigroup, Monoid)

-- | A tree is defined completely by its 'root' and its 'subforest'.
--
-- To construct a tree, use 'tree' or 'leaf'.
data Tree a = Tree
  { -- | The value at the root node of the tree.
    root :: a,
    -- | The forest containing all descendants
    --   of the tree's 'root'.
    subforest :: Forest a
  }
  deriving (Eq, Show, Functor, Foldable, Traversable)

-- | Construct a forest from a list of trees.
--
-- /@'forest' []@ is equivalent to 'mempty'./
forest :: [Tree a] -> Forest a
forest = Forest

-- | Construct a tree with a single root and no subforest.
--
--   /@'leaf' x@ is equivalent to @'tree' x 'mempty'@./
leaf :: a -> Tree a
leaf a = tree a mempty

-- | Construct a forest of depth 1, where each tree contains only a root.
--
-- /'leaves' is equivalent to @'forest' . fmap 'leaf'@/
leaves :: [a] -> Forest a
leaves xs =
  forest (fmap leaf xs)

-- | Construct a tree with a root and subforest.
tree :: a -> Forest a -> Tree a
tree = Tree

-- | The tree's immediate subtrees.
--
-- /'subtrees' is equivalent to @'trees' . 'subforest'@./
subtrees :: Tree a -> [Tree a]
subtrees t = trees (subforest t)

-- | Catamorphism on forests.
--
-- >>>
-- :{
-- example :: Forest Char
-- example = forest
--    [ tree 'a' $ leaves "bc"
--    , tree 'd' $ forest
--        [ leaf 'e'
--        , tree 'f' $ leaves "g"
--        ]
--   ]
-- :}
--
-- >>> foldForest (intercalate ", " . fmap (\(a, b) -> [a] <> " [" <> b <> "]")) example
-- "a [b [], c []], d [e [], f [g []]]"
foldForest :: ([(a, b)] -> b) -> Forest a -> b
foldForest f =
  go
  where
    go (Forest ts) = f $ (\t -> (root t, go (subforest t))) <$> ts

-- | Catamorphism on trees.
--
-- >>>
-- :{
-- example :: Tree Char
-- example = tree 'a' $ forest
--    [ tree 'b' $ leaves "cd"
--    , tree 'e' $ forest
--        [ leaf 'f'
--        , tree 'g' $ leaves "h"
--        ]
--   ]
-- :}
--
-- >>> foldTree (\a bs -> [a] <> " [" <> intercalate ", " bs <> "]") example
-- "a [b [c [], d []], e [f [], g [h []]]]"
foldTree :: (a -> [b] -> b) -> Tree a -> b
foldTree f =
  go
  where
    go t = f (root t) (go <$> subtrees t)

-- $setup
--
-- >>> import Prelude
-- >>> import Data.Char
-- >>> import Data.Foldable
-- >>> import Data.Function
-- >>> import Data.List
-- >>> import Data.Semigroup

-- $imports
--
-- Recommended imports:
--
-- > import Data.Forest (Forest, Tree)
-- > import qualified Data.Forest as Forest

-- $functor
--
-- One notable difference of this 'Forest' from that of the /containers/ library is
-- that this 'Forest' is a newtype rather than a type alias, and so it provides a
-- more appropriate 'Functor' instance:
--
-- >>>
-- :{
-- example :: Forest Char
-- example = forest
--     [ tree 'a' $ leaves "bc"
--     , tree 'd' $ forest
--         [ leaf 'e'
--         , tree 'f' $ leaves "g"
--         ]
--    ]
-- :}
--
-- >>>
-- :{
-- showCharForest f =
--     intercalate ", " (showCharTree <$> trees f)
--   where
--     showCharTree t = case trees (subforest t) of
--       []   -> [root t]
--       [t'] -> [root t] <> ": " <> showCharTree t'
--       _    -> [root t] <> ": (" <> showCharForest (subforest t) <> ")"
-- :}
--
-- >>> showCharForest example
-- "a: (b, c), d: (e, f: g)"
--
-- >>> showCharForest (fmap toUpper example)
-- "A: (B, C), D: (E, F: G)"
--
-- Likewise, 'Forest''s 'Foldable' instance folds over the elements of the forest.
--
-- >>> toList example
-- "abcdefg"