Rooted.hs

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}

-- |
-- Module      :  ELynx.Tree.Rooted
-- Description :  Rooted trees with labeled branches
-- Copyright   :  2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Jan 17 09:57:29 2019.
--
-- Rooted 'Tree's differs from a classical rose 'Data.Tree.Tree's in that they
-- have labeled branches.
--
-- For rooted topologies, please see 'ELynx.Topology.Rooted'.
--
-- A 'Tree' is defined as:
--
-- @
-- data Tree e a = Node
--   { branch :: e,
--     label :: a,
--     forest :: Forest e a
--   }
-- @
--
-- where
--
-- @
-- type Forest e a = [Tree e a]
-- @
--
-- This means, that the word 'Node' is reserved for the constructor of a tree,
-- and that a 'Node' has an attached 'branch', a 'label', and a sub-'forest'.
-- The terms /Node/ and /label/ referring to the value constructor 'Node' and
-- the record function 'label', respectively, are not to be confused. The
-- elements of the sub-forest are often called /children/.
--
-- In mathematical terms: A 'Tree' is a directed acyclic graph without loops,
-- with vertex labels, and with edge labels.
--
-- A short recap of recursive tree traversals:
--
-- - Pre-order: Root first, then sub trees from left to right. Also called depth
--   first.
--
-- - In-order: Only valid for bifurcating trees. Left sub tree first, then root,
--   then right sub tree.
--
-- - Post-order: Sub trees from left to right, then the root. Also called
--   breadth first.
--
-- Here, pre-order traversals are used exclusively, for example, by accessor
-- functions such as 'branches', or 'labels' which is the same as 'toList'.
-- Please let me know, if post-order algorithms are required.
module ELynx.Tree.Rooted
  ( -- * Tree with branch labels
    Tree (..),
    Forest,
    fromRoseTree,
    toTreeBranchLabels,
    toTreeNodeLabels,

    -- * Access leaves, branches and labels
    leaves,
    duplicateLeaves,
    setStem,
    modifyStem,
    branches,
    setBranches,
    setLabel,
    modifyLabel,
    labels,
    duplicateLabels,
    setLabels,
    identify,

    -- * Structure
    degree,
    depth,
    prune,
    dropNodesWith,
    dropLeavesWith,
    zipTreesWith,
    zipTrees,
    flipLabels,

    -- * Newtypes with specific instances
    ZipTree (..),
    BranchTree (..),
    ZipBranchTree (..),
  )
where

import Control.Comonad
import Control.DeepSeq
import Control.Monad
import Data.Aeson
import Data.Bifoldable
import Data.Bifunctor
import Data.Bitraversable
import Data.Data
import Data.Foldable
import Data.Functor
import Data.List
import Data.Maybe
import qualified Data.Set as S
import qualified Data.Tree as T
import GHC.Generics

-- | Rooted rose trees with branch labels.
--
-- Unary instances such as 'Functor' act on node labels, and not on branch
-- labels. Binary instances such as 'Bifunctor' act on both labels (`first` acts
-- on branches, `second` on node labels).
data Tree e a = Node
  { branch :: e,
    label :: a,
    forest :: Forest e a
  }
  deriving (Eq, Read, Show, Data, Generic)

-- | Shorthand.
type Forest e a = [Tree e a]

-- | Map over node labels.
instance Functor (Tree e) where
  fmap f ~(Node br lb ts) = Node br (f lb) $ map (fmap f) ts
  lb <$ ~(Node br _ ts) = Node br lb (map (lb <$) ts)

-- | The function 'first' acts on branch labels, 'second' on node labels.
instance Bifunctor Tree where
  bimap f g ~(Node br lb ts) = Node (f br) (g lb) $ map (bimap f g) ts
  first f ~(Node br lb ts) = Node (f br) lb $ map (first f) ts
  second g ~(Node br lb ts) = Node br (g lb) $ map (second g) ts

-- | Combine node labels in pre-order.
instance Foldable (Tree e) where
  foldMap f ~(Node _ lb ts) = f lb <> foldMap (foldMap f) ts
  null _ = False
  {-# INLINE null #-}
  toList = labels
  {-# INLINE toList #-}

instance Bifoldable Tree where
  bifoldMap f g ~(Node br lb ts) = f br <> g lb <> foldMap (bifoldMap f g) ts

instance Traversable (Tree e) where
  traverse g ~(Node br lb ts) = Node br <$> g lb <*> traverse (traverse g) ts

instance Bitraversable Tree where
  bitraverse f g ~(Node br lb ts) = Node <$> f br <*> g lb <*> traverse (bitraverse f g) ts

-- | The 'Semigroup' instance of the branch labels determines how the
-- branches are combined. For example, distances can be summed using
-- 'Data.Semigroup.Sum'.
--
-- The 'Monoid' instance of the branch labels determines the default branch
-- label when using 'pure'.
--
-- This instance is similar to the one provided by 'Data.Tree.Tree'. For an
-- alternative, see 'ZipTree'.
instance (Semigroup e, Monoid e) => Applicative (Tree e) where
  pure lb = Node mempty lb []
  ~(Node brF lbF tsF) <*> ~tx@(Node brX lbX tsX) =
    Node (brF <> brX) (lbF lbX) (map (bimap (brF <>) lbF) tsX ++ map (<*> tx) tsF)
  liftA2 f ~(Node brX lbX tsX) ~ty@(Node brY lbY tsY) =
    Node (brX <> brY) (f lbX lbY) (map (bimap (brX <>) (f lbX)) tsY ++ map (\tx -> liftA2 f tx ty) tsX)
  ~(Node brX _ tsX) *> ~ty@(Node brY lbY tsY) =
    Node (brX <> brY) lbY (map (first (brX <>)) tsY ++ (tsX <&> (*> ty)))
  ~(Node brX lbX tsX) <* ~ty@(Node brY _ tsY) =
    Node (brX <> brY) lbX (map (bimap (brX <>) (const lbX)) tsY ++ (tsX <&> (<* ty)))

-- | The 'Semigroup' instance of the branch labels determines how the branches
-- are combined. For example, distances can be summed using
-- 'Data.Semigroup.Sum'.
--
-- The 'Monoid' instance of the branch labels determines the default branch
-- label when using 'return'.
instance (Semigroup e, Monoid e) => Monad (Tree e) where
  ~(Node br lb ts) >>= f = case f lb of
    Node br' lb' ts' -> Node (br <> br') lb' (map (first (br <>)) ts' ++ map (>>= f) ts)

-- -- NOTE: We cannot provide a MonadZip instance because branch labels cannot
-- -- be recovered from the combined label.
--
-- instance Monoid e => MonadZip (Tree e) where
--   mzipWith f (Node brL lbL tsL) (Node brR lbR tsR) =
--     Node (brL <> brR) (f lbL lbR) (mzipWith (mzipWith f) tsL tsR)
--
--   munzip (Node br (lbL, lbR) ts) = (Node ? lbL tsL, Node ? lbR tsR)
--     where
--       (tsL, tsR) = munzip (map munzip ts)

-- -- NOTE: I don't really know much about 'MonadFix', and so do not provide the
-- -- instance.
--
-- instance Monoid e => MonadFix (Tree e) where
--   mfix = mfixTree

-- mfixTree :: (a -> Tree e a) -> Tree e a
-- mfixTree f
--   | Node br lb ts <- fix (f . label) =
--     Node
--       br
--       lb
--       ( zipWith
--           (\i _ -> mfixTree ((!! i) . forest . f))
--           [0 ..]
--           ts
--       )

instance Comonad (Tree e) where
  duplicate t@(Node br _ ts) = Node br t (map duplicate ts)
  extract = label
  {-# INLINE extract #-}

instance (NFData e, NFData a) => NFData (Tree e a) where
  rnf (Node br lb ts) = rnf br `seq` rnf lb `seq` rnf ts

instance (ToJSON e, ToJSON a) => ToJSON (Tree e a)

instance (FromJSON e, FromJSON a) => FromJSON (Tree e a)

-- | Conversion from 'T.Tree'.
fromRoseTree :: T.Tree a -> Tree () a
fromRoseTree (T.Node l ts) = Node () l $ map fromRoseTree ts

-- | Conversion to 'T.Tree' using branch labels.
toTreeBranchLabels :: Tree e a -> T.Tree e
toTreeBranchLabels (Node br _ ts) = T.Node br (map toTreeBranchLabels ts)

-- | Conversion to 'T.Tree' using node labels.
toTreeNodeLabels :: Tree e a -> T.Tree a
toTreeNodeLabels (Node _ lb ts) = T.Node lb (map toTreeNodeLabels ts)

-- | List of leaves.
leaves :: Tree e a -> [a]
leaves t = squish t []
  where
    squish (Node _ lb []) xs = lb : xs
    squish (Node _ _ ts) xs = foldr squish xs ts

duplicates :: (Ord a) => [a] -> Bool
duplicates = go S.empty
  where
    go _ [] = False
    go seen (x : xs) = x `S.member` seen || go (S.insert x seen) xs

-- | Check if a tree has duplicate leaves.
duplicateLeaves :: (Ord a) => Tree e a -> Bool
duplicateLeaves = duplicates . leaves

-- | Set the stem to a given value.
setStem :: e -> Tree e a -> Tree e a
setStem br (Node _ lb ts) = Node br lb ts

-- | Modify the stem of a tree.
modifyStem :: (e -> e) -> Tree e a -> Tree e a
modifyStem f t = t {branch = f $ branch t}

-- | Get branch labels in pre-order.
branches :: Tree e a -> [e]
branches t = squish t []
  where
    squish (Node br _ ts) xs = br : foldr squish xs ts

-- | Set branch labels in pre-order.
--
-- Return 'Nothing' if the provided list of branch labels is too short.
setBranches :: (Bitraversable t) => [f] -> t e a -> Maybe (t f a)
setBranches xs = bisequenceA . snd . bimapAccumL setBranch noChange xs
  where
    setBranch [] _ = ([], Nothing)
    setBranch (y : ys) _ = (ys, Just y)
    noChange ys z = (ys, Just z)

-- | Set label.
setLabel :: a -> Tree e a -> Tree e a
setLabel lb (Node br _ ts) = Node br lb ts

-- | Modify the root label of a tree.
modifyLabel :: (a -> a) -> Tree e a -> Tree e a
modifyLabel f t = t {label = f $ label t}

-- | Return node labels in pre-order.
labels :: Tree e a -> [a]
labels t = squish t []
  where
    squish (Node _ lb ts) xs = lb : foldr squish xs ts

-- | Check if a tree has duplicate labels.
duplicateLabels :: (Ord a) => Tree e a -> Bool
duplicateLabels = duplicates . labels

-- | Set node labels in pre-order.
--
-- Return 'Nothing' if the provided list of node labels is too short.
setLabels :: (Traversable t) => [b] -> t a -> Maybe (t b)
setLabels xs = sequenceA . snd . mapAccumL setLabelM xs
  where
    setLabelM [] _ = ([], Nothing)
    setLabelM (y : ys) _ = (ys, Just y)

-- | Label the nodes in pre-order with unique indices starting at 0.
identify :: (Traversable t) => t a -> t Int
identify = snd . mapAccumL (\i _ -> (i + 1, i)) (0 :: Int)

-- | Degree of the root node.
--
-- The degree of a node is the number of branches attached to the node.
degree :: Tree e a -> Int
degree = (+ 1) . length . forest

-- | Depth of a tree.
--
-- The [depth of a tree](https://en.wikipedia.org/wiki/Tree-depth) is the
-- largest number of nodes traversed on a path from the root to a leaf.
--
-- By convention, the depth is larger equal 1. That is, the depth of a leaf tree
-- is 1.
depth :: Tree e a -> Int
depth = maximum . go 1
  where
    go n (Node _ _ []) = [n]
    go n (Node _ _ xs) = concatMap (go (n + 1)) xs

-- | Prune degree two nodes.
--
-- The label of a pruned node is lost. The branches are combined according to
-- their 'Semigroup' instance of the form
--
-- @\daughterBranch parentBranch -> combinedBranch@.
prune :: (Semigroup e) => Tree e a -> Tree e a
prune t@(Node _ _ []) = t
prune (Node paBr _ [Node daBr daLb daTs]) = Node (daBr <> paBr) daLb daTs
prune (Node paBr paLb paTs) = Node paBr paLb $ map prune paTs

-- | Drop nodes satisfying predicate.
--
-- Degree two nodes may arise.
--
-- Also drop nodes of which all daughter nodes are dropped.
--
-- Return 'Nothing' if
--
-- - The root node satisfies the predicate.
--
-- - All daughter nodes of the root are dropped.
dropNodesWith :: (a -> Bool) -> Tree e a -> Maybe (Tree e a)
dropNodesWith p (Node br lb ts)
  | p lb = Nothing
  | otherwise =
      if null ts'
        then Nothing
        else Just $ Node br lb ts'
  where
    ts' = mapMaybe (dropNodesWith p) ts

-- | Drop leaves satisfying predicate.
--
-- Degree two nodes may arise.
--
-- Also drop nodes of which all daughter nodes are dropped.
--
-- Return 'Nothing' if all leaves satisfy the predicate.
dropLeavesWith :: (a -> Bool) -> Tree e a -> Maybe (Tree e a)
dropLeavesWith p (Node br lb [])
  | p lb = Nothing
  | otherwise = Just $ Node br lb []
dropLeavesWith p (Node br lb ts) =
  if null ts'
    then Nothing
    else Just $ Node br lb ts'
  where
    ts' = mapMaybe (dropLeavesWith p) ts

-- | Zip two trees with the same topology.
--
-- This function differs from the 'Applicative' instance of 'ZipTree' in that it
-- fails when the topologies don't match. Further, it allows specification of a
-- zipping function for the branches.
--
-- Return 'Nothing' if the topologies are different.
zipTreesWith ::
  (e1 -> e2 -> e) ->
  (a1 -> a2 -> a) ->
  Tree e1 a1 ->
  Tree e2 a2 ->
  Maybe (Tree e a)
zipTreesWith f g (Node brL lbL tsL) (Node brR lbR tsR) =
  if length tsL == length tsR
    then -- I am proud of that :)).
      zipWithM (zipTreesWith f g) tsL tsR >>= Just . Node (f brL brR) (g lbL lbR)
    else Nothing

-- | See 'zipTreesWith'.
zipTrees :: Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree (e1, e2) (a1, a2))
zipTrees = zipTreesWith (,) (,)

-- | Flip the branch and node lables.
flipLabels :: Tree e a -> Tree a e
flipLabels (Node x y zs) = Node y x $ map flipLabels zs

-- | This newtype provides instances acting on the branch labels, and not on the
-- node labels as it is the case in 'Tree'.
newtype BranchTree a e = BranchTree {getBranchTree :: Tree e a}
  deriving (Eq, Read, Show, Data, Generic)

-- | Map over branch labels.
instance Functor (BranchTree a) where
  fmap f ~(BranchTree (Node br lb ts)) =
    BranchTree $ Node (f br) lb $ map (getBranchTree . fmap f . BranchTree) ts
  br <$ ~(BranchTree (Node _ lb ts)) =
    BranchTree $ Node br lb (map (getBranchTree . (br <$) . BranchTree) ts)

-- | Combine branch labels in pre-order.
instance Foldable (BranchTree a) where
  foldMap f ~(BranchTree (Node br _ ts)) =
    f br <> foldMap (foldMap f . BranchTree) ts
  null _ = False
  {-# INLINE null #-}
  toList = branches . getBranchTree
  {-# INLINE toList #-}

instance Traversable (BranchTree a) where
  traverse g ~(BranchTree (Node br lb ts)) =
    assemble lb <$> fbr' <*> fts'
    where
      assemble lb' br' ts' = BranchTree $ Node br' lb' ts'
      fbr' = g br
      fts' = map getBranchTree <$> traverse (traverse g . BranchTree) ts

instance Comonad (BranchTree a) where
  duplicate (BranchTree t@(Node _ lb ts)) =
    BranchTree $
      Node (BranchTree t) lb $
        map (getBranchTree . duplicate . BranchTree) ts
  extract = branch . getBranchTree

instance (Monoid a) => Applicative (BranchTree a) where
  -- Infinite layers with infinite subtrees.
  pure br = BranchTree $ Node br mempty []
  (BranchTree ~(Node brF lbF tsF)) <*> tx@(BranchTree ~(Node brX lbX tsX)) =
    BranchTree $
      Node
        (brF brX)
        (lbF <> lbX)
        ( map (bimap brF (lbF <>)) tsX
            ++ map (getBranchTree . (<*> tx) . BranchTree) tsF
        )
  liftA2 f (BranchTree ~(Node brX lbX tsX)) ty@(BranchTree ~(Node brY lbY tsY)) =
    BranchTree $
      Node
        (f brX brY)
        (lbX <> lbY)
        ( map (bimap (f brX) (lbX <>)) tsY
            ++ map (\tx -> getBranchTree $ liftA2 f (BranchTree tx) ty) tsX
        )
  (BranchTree ~(Node _ lbX tsX)) *> ty@(BranchTree ~(Node brY lbY tsY)) =
    BranchTree $
      Node
        brY
        (lbX <> lbY)
        ( getBranchTree
            <$> ( map (BranchTree . second (lbX <>)) tsY
                    ++ map ((*> ty) . BranchTree) tsX
                )
        )
  (BranchTree ~(Node brX lbX tsX)) <* ty@(BranchTree ~(Node _ lbY tsY)) =
    BranchTree $
      Node
        brX
        (lbX <> lbY)
        ( map (bimap (const brX) (lbX <>)) tsY
            ++ map (getBranchTree . (<* ty) . BranchTree) tsX
        )

-- | This newtype provides a zip-like applicative instance, similar to
-- 'Control.Applicative.ZipList'.
--
-- The default applicative instance of 'Tree' is not zip-like, because the
-- zip-like instance makes the Monad instance meaningless (similar to the
-- behavior observed with lists).
newtype ZipTree e a = ZipTree {getZipTree :: Tree e a}
  deriving (Eq, Read, Show, Data, Generic)

deriving instance Functor (ZipTree e)

deriving instance Foldable (ZipTree e)

instance Traversable (ZipTree e) where
  traverse f (ZipTree t) = ZipTree <$> traverse f t

instance Comonad (ZipTree e) where
  duplicate (ZipTree t) = ZipTree $ second ZipTree $ duplicate t
  extract = label . getZipTree

-- | The 'Monoid' instance of the branch labels determines the default branch
-- label, and how the branches are combined. For example, distances can be
-- summed using the 'Data.Monoid.Sum' monoid.
--
-- >>> let t = ZipTree $ Node "" 0 [Node "" 1 [], Node "" 2 []] :: ZipTree String Int
-- >>> let f = ZipTree $ Node "+3" (+3) [Node "*5" (*5) [], Node "+10" (+10) []] :: ZipTree String (Int -> Int)
-- >>> f <*> t
--
-- ZipTree {getZipTree = Node {branch = "+3", label = 3, forest = [Node {branch = "*5", label = 5, forest = []},Node {branch = "+10", label = 12, forest = []}]}}
instance (Monoid e) => Applicative (ZipTree e) where
  -- Infinite layers with infinite subtrees.
  pure lb = ZipTree $ Node mempty lb $ repeat (getZipTree $ pure lb)
  (ZipTree ~(Node brF lbF tsF)) <*> (ZipTree ~(Node brX lbX tsX)) =
    ZipTree $ Node (brF <> brX) (lbF lbX) (zipWith f tsF tsX)
    where
      f x y = getZipTree $ ZipTree x <*> ZipTree y
  liftA2 f (ZipTree ~(Node brX lbX tsX)) (ZipTree ~(Node brY lbY tsY)) =
    ZipTree $ Node (brX <> brY) (f lbX lbY) (zipWith g tsX tsY)
    where
      g x y = getZipTree $ liftA2 f (ZipTree x) (ZipTree y)
  (ZipTree ~(Node brX _ tsX)) *> (ZipTree ~(Node brY lbY tsY)) =
    ZipTree $ Node (brX <> brY) lbY (zipWith f tsX tsY)
    where
      f x y = getZipTree $ ZipTree x *> ZipTree y
  (ZipTree ~(Node brX lbX tsX)) <* (ZipTree ~(Node brY _ tsY)) =
    ZipTree $ Node (brX <> brY) lbX (zipWith f tsX tsY)
    where
      f x y = getZipTree $ ZipTree x <* ZipTree y

-- | Like 'ZipTree' but act on branch labels; see 'BranchTree'.
newtype ZipBranchTree a e = ZipBranchTree {getZipBranchTree :: Tree e a}
  deriving (Eq, Read, Show, Data, Generic)

-- | Map over branch labels.
instance Functor (ZipBranchTree a) where
  fmap f ~(ZipBranchTree (Node br lb ts)) =
    ZipBranchTree $ Node (f br) lb $ map g ts
    where
      g = getZipBranchTree . fmap f . ZipBranchTree
  br <$ ~(ZipBranchTree (Node _ lb ts)) =
    ZipBranchTree $ Node br lb (map f ts)
    where
      f = getZipBranchTree . (br <$) . ZipBranchTree

-- | Combine branch labels in pre-order.
instance Foldable (ZipBranchTree a) where
  foldMap f ~(ZipBranchTree (Node br _ ts)) =
    f br <> foldMap g ts
    where
      g = foldMap f . ZipBranchTree
  null _ = False
  {-# INLINE null #-}
  toList = branches . getZipBranchTree
  {-# INLINE toList #-}

instance Traversable (ZipBranchTree a) where
  traverse g ~(ZipBranchTree (Node br lb ts)) =
    assemble lb <$> fbr' <*> fts'
    where
      assemble lb' br' ts' = ZipBranchTree $ Node br' lb' ts'
      fbr' = g br
      fts' = map getZipBranchTree <$> traverse (traverse g . ZipBranchTree) ts

instance Comonad (ZipBranchTree a) where
  duplicate (ZipBranchTree t@(Node _ lb ts)) =
    ZipBranchTree $
      Node (ZipBranchTree t) lb $
        map (getZipBranchTree . duplicate . ZipBranchTree) ts
  extract = branch . getZipBranchTree

-- | See the 'Applicative' instance of 'ZipTree'.
instance (Monoid a) => Applicative (ZipBranchTree a) where
  -- Infinite layers with infinite subtrees.
  pure br = ZipBranchTree $ Node br mempty $ repeat (getZipBranchTree $ pure br)
  (ZipBranchTree ~(Node brF lbF tsF)) <*> (ZipBranchTree ~(Node brX lbX tsX)) =
    ZipBranchTree $ Node (brF brX) (lbF <> lbX) (zipWith f tsF tsX)
    where
      f x y = getZipBranchTree $ ZipBranchTree x <*> ZipBranchTree y
  liftA2 f (ZipBranchTree ~(Node brX lbX tsX)) (ZipBranchTree ~(Node brY lbY tsY)) =
    ZipBranchTree $ Node (f brX brY) (lbX <> lbY) (zipWith g tsX tsY)
    where
      g x y = getZipBranchTree $ liftA2 f (ZipBranchTree x) (ZipBranchTree y)
  (ZipBranchTree ~(Node _ lbX tsX)) *> (ZipBranchTree ~(Node brY lbY tsY)) =
    ZipBranchTree $ Node brY (lbX <> lbY) (zipWith f tsX tsY)
    where
      f x y = getZipBranchTree $ ZipBranchTree x *> ZipBranchTree y
  (ZipBranchTree ~(Node brX lbX tsX)) <* (ZipBranchTree ~(Node _ lbY tsY)) =
    ZipBranchTree $ Node brX (lbX <> lbY) (zipWith f tsX tsY)
    where
      f x y = getZipBranchTree $ ZipBranchTree x <* ZipBranchTree y