{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}

-- |
-- Module      :  ELynx.Tree.Rooted
-- Description :  Rooted trees with labeled branches
-- Copyright   :  (c) Dominik Schrempf 2020
-- 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 differes from a classical rose 'Data.Tree.Tree' in that it has
-- 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 value constructor /Node/ and the record function /label/ 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
  ( -- * Data type
    Tree (..),
    Forest,
    toTreeBranchLabels,
    toTreeNodeLabels,

    -- * Access leaves, branches and labels
    leaves,
    duplicateLeaves,
    setStem,
    applyStem,
    branches,
    setBranches,
    setLabel,
    applyLabel,
    labels,
    setLabels,
    applyRoot,
    identify,

    -- * Structure
    degree,
    prune,
    dropNodesWith,
    dropLeavesWith,
    zipTreesWith,
    zipTrees,
  )
where

import Control.Applicative
import Control.Comonad
import Control.DeepSeq
import Control.Monad
import Control.Monad.Fix
import Data.Aeson
import Data.Bifoldable
import Data.Bifunctor
import Data.Bitraversable
import Data.Data
import Data.Foldable
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).
--
-- Lifted instances are not provided.
data Tree e a = Node
  { Tree e a -> e
branch :: e,
    Tree e a -> a
label :: a,
    Tree e a -> Forest e a
forest :: Forest e a
  }
  deriving (Tree e a -> Tree e a -> Bool
(Tree e a -> Tree e a -> Bool)
-> (Tree e a -> Tree e a -> Bool) -> Eq (Tree e a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall e a. (Eq e, Eq a) => Tree e a -> Tree e a -> Bool
/= :: Tree e a -> Tree e a -> Bool
$c/= :: forall e a. (Eq e, Eq a) => Tree e a -> Tree e a -> Bool
== :: Tree e a -> Tree e a -> Bool
$c== :: forall e a. (Eq e, Eq a) => Tree e a -> Tree e a -> Bool
Eq, ReadPrec [Tree e a]
ReadPrec (Tree e a)
Int -> ReadS (Tree e a)
ReadS [Tree e a]
(Int -> ReadS (Tree e a))
-> ReadS [Tree e a]
-> ReadPrec (Tree e a)
-> ReadPrec [Tree e a]
-> Read (Tree e a)
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
forall e a. (Read e, Read a) => ReadPrec [Tree e a]
forall e a. (Read e, Read a) => ReadPrec (Tree e a)
forall e a. (Read e, Read a) => Int -> ReadS (Tree e a)
forall e a. (Read e, Read a) => ReadS [Tree e a]
readListPrec :: ReadPrec [Tree e a]
$creadListPrec :: forall e a. (Read e, Read a) => ReadPrec [Tree e a]
readPrec :: ReadPrec (Tree e a)
$creadPrec :: forall e a. (Read e, Read a) => ReadPrec (Tree e a)
readList :: ReadS [Tree e a]
$creadList :: forall e a. (Read e, Read a) => ReadS [Tree e a]
readsPrec :: Int -> ReadS (Tree e a)
$creadsPrec :: forall e a. (Read e, Read a) => Int -> ReadS (Tree e a)
Read, Int -> Tree e a -> ShowS
[Tree e a] -> ShowS
Tree e a -> String
(Int -> Tree e a -> ShowS)
-> (Tree e a -> String) -> ([Tree e a] -> ShowS) -> Show (Tree e a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall e a. (Show e, Show a) => Int -> Tree e a -> ShowS
forall e a. (Show e, Show a) => [Tree e a] -> ShowS
forall e a. (Show e, Show a) => Tree e a -> String
showList :: [Tree e a] -> ShowS
$cshowList :: forall e a. (Show e, Show a) => [Tree e a] -> ShowS
show :: Tree e a -> String
$cshow :: forall e a. (Show e, Show a) => Tree e a -> String
showsPrec :: Int -> Tree e a -> ShowS
$cshowsPrec :: forall e a. (Show e, Show a) => Int -> Tree e a -> ShowS
Show, Typeable (Tree e a)
DataType
Constr
Typeable (Tree e a)
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> Tree e a -> c (Tree e a))
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c (Tree e a))
-> (Tree e a -> Constr)
-> (Tree e a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c (Tree e a)))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e))
    -> Maybe (c (Tree e a)))
-> ((forall b. Data b => b -> b) -> Tree e a -> Tree e a)
-> (forall r r'.
    (r -> r' -> r)
    -> r -> (forall d. Data d => d -> r') -> Tree e a -> r)
-> (forall r r'.
    (r' -> r -> r)
    -> r -> (forall d. Data d => d -> r') -> Tree e a -> r)
-> (forall u. (forall d. Data d => d -> u) -> Tree e a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> Tree e a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a))
-> Data (Tree e a)
Tree e a -> DataType
Tree e a -> Constr
(forall b. Data b => b -> b) -> Tree e a -> Tree e a
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Tree e a -> c (Tree e a)
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Tree e a)
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Tree e a))
forall a.
Typeable a
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> a -> c a)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c a)
-> (a -> Constr)
-> (a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c a))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c a))
-> ((forall b. Data b => b -> b) -> a -> a)
-> (forall r r'.
    (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall r r'.
    (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall u. (forall d. Data d => d -> u) -> a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> Data a
forall u. Int -> (forall d. Data d => d -> u) -> Tree e a -> u
forall u. (forall d. Data d => d -> u) -> Tree e a -> [u]
forall e a. (Data e, Data a) => Typeable (Tree e a)
forall e a. (Data e, Data a) => Tree e a -> DataType
forall e a. (Data e, Data a) => Tree e a -> Constr
forall e a.
(Data e, Data a) =>
(forall b. Data b => b -> b) -> Tree e a -> Tree e a
forall e a u.
(Data e, Data a) =>
Int -> (forall d. Data d => d -> u) -> Tree e a -> u
forall e a u.
(Data e, Data a) =>
(forall d. Data d => d -> u) -> Tree e a -> [u]
forall e a r r'.
(Data e, Data a) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
forall e a r r'.
(Data e, Data a) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
forall e a (m :: * -> *).
(Data e, Data a, Monad m) =>
(forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
forall e a (m :: * -> *).
(Data e, Data a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
forall e a (c :: * -> *).
(Data e, Data a) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Tree e a)
forall e a (c :: * -> *).
(Data e, Data a) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Tree e a -> c (Tree e a)
forall e a (t :: * -> *) (c :: * -> *).
(Data e, Data a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Tree e a))
forall e a (t :: * -> * -> *) (c :: * -> *).
(Data e, Data a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Tree e a))
forall r r'.
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
forall r r'.
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Tree e a)
forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Tree e a -> c (Tree e a)
forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c (Tree e a))
forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Tree e a))
$cNode :: Constr
$tTree :: DataType
gmapMo :: (forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
$cgmapMo :: forall e a (m :: * -> *).
(Data e, Data a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
gmapMp :: (forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
$cgmapMp :: forall e a (m :: * -> *).
(Data e, Data a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
gmapM :: (forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
$cgmapM :: forall e a (m :: * -> *).
(Data e, Data a, Monad m) =>
(forall d. Data d => d -> m d) -> Tree e a -> m (Tree e a)
gmapQi :: Int -> (forall d. Data d => d -> u) -> Tree e a -> u
$cgmapQi :: forall e a u.
(Data e, Data a) =>
Int -> (forall d. Data d => d -> u) -> Tree e a -> u
gmapQ :: (forall d. Data d => d -> u) -> Tree e a -> [u]
$cgmapQ :: forall e a u.
(Data e, Data a) =>
(forall d. Data d => d -> u) -> Tree e a -> [u]
gmapQr :: (r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
$cgmapQr :: forall e a r r'.
(Data e, Data a) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
gmapQl :: (r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
$cgmapQl :: forall e a r r'.
(Data e, Data a) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Tree e a -> r
gmapT :: (forall b. Data b => b -> b) -> Tree e a -> Tree e a
$cgmapT :: forall e a.
(Data e, Data a) =>
(forall b. Data b => b -> b) -> Tree e a -> Tree e a
dataCast2 :: (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Tree e a))
$cdataCast2 :: forall e a (t :: * -> * -> *) (c :: * -> *).
(Data e, Data a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Tree e a))
dataCast1 :: (forall d. Data d => c (t d)) -> Maybe (c (Tree e a))
$cdataCast1 :: forall e a (t :: * -> *) (c :: * -> *).
(Data e, Data a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Tree e a))
dataTypeOf :: Tree e a -> DataType
$cdataTypeOf :: forall e a. (Data e, Data a) => Tree e a -> DataType
toConstr :: Tree e a -> Constr
$ctoConstr :: forall e a. (Data e, Data a) => Tree e a -> Constr
gunfold :: (forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Tree e a)
$cgunfold :: forall e a (c :: * -> *).
(Data e, Data a) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Tree e a)
gfoldl :: (forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Tree e a -> c (Tree e a)
$cgfoldl :: forall e a (c :: * -> *).
(Data e, Data a) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Tree e a -> c (Tree e a)
$cp1Data :: forall e a. (Data e, Data a) => Typeable (Tree e a)
Data, (forall x. Tree e a -> Rep (Tree e a) x)
-> (forall x. Rep (Tree e a) x -> Tree e a) -> Generic (Tree e a)
forall x. Rep (Tree e a) x -> Tree e a
forall x. Tree e a -> Rep (Tree e a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall e a x. Rep (Tree e a) x -> Tree e a
forall e a x. Tree e a -> Rep (Tree e a) x
$cto :: forall e a x. Rep (Tree e a) x -> Tree e a
$cfrom :: forall e a x. Tree e a -> Rep (Tree e a) x
Generic)

-- | A shorthand.
type Forest e a = [Tree e a]

-- | Map over node labels.
instance Functor (Tree e) where
  fmap :: (a -> b) -> Tree e a -> Tree e b
fmap a -> b
f ~(Node e
br a
lb Forest e a
ts) = e -> b -> Forest e b -> Tree e b
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br (a -> b
f a
lb) (Forest e b -> Tree e b) -> Forest e b -> Tree e b
forall a b. (a -> b) -> a -> b
$ (Tree e a -> Tree e b) -> Forest e a -> Forest e b
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> Tree e a -> Tree e b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) Forest e a
ts
  a
x <$ :: a -> Tree e b -> Tree e a
<$ ~(Node e
br b
_ Forest e b
ts) = e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
x ((Tree e b -> Tree e a) -> Forest e b -> Forest e a
forall a b. (a -> b) -> [a] -> [b]
map (a
x a -> Tree e b -> Tree e a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$) Forest e b
ts)

-- | The function 'first' acts on branch labels, 'second' on node labels.
instance Bifunctor Tree where
  bimap :: (a -> b) -> (c -> d) -> Tree a c -> Tree b d
bimap a -> b
f c -> d
g ~(Node a
br c
lb Forest a c
ts) = b -> d -> Forest b d -> Tree b d
forall e a. e -> a -> Forest e a -> Tree e a
Node (a -> b
f a
br) (c -> d
g c
lb) (Forest b d -> Tree b d) -> Forest b d -> Tree b d
forall a b. (a -> b) -> a -> b
$ (Tree a c -> Tree b d) -> Forest a c -> Forest b d
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> (c -> d) -> Tree a c -> Tree b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> b
f c -> d
g) Forest a c
ts
  first :: (a -> b) -> Tree a c -> Tree b c
first a -> b
f ~(Node a
br c
lb Forest a c
ts) = b -> c -> Forest b c -> Tree b c
forall e a. e -> a -> Forest e a -> Tree e a
Node (a -> b
f a
br) c
lb (Forest b c -> Tree b c) -> Forest b c -> Tree b c
forall a b. (a -> b) -> a -> b
$ (Tree a c -> Tree b c) -> Forest a c -> Forest b c
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> Tree a c -> Tree b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f) Forest a c
ts
  second :: (b -> c) -> Tree a b -> Tree a c
second b -> c
g ~(Node a
br b
lb Forest a b
ts) = a -> c -> Forest a c -> Tree a c
forall e a. e -> a -> Forest e a -> Tree e a
Node a
br (b -> c
g b
lb) (Forest a c -> Tree a c) -> Forest a c -> Tree a c
forall a b. (a -> b) -> a -> b
$ (Tree a b -> Tree a c) -> Forest a b -> Forest a c
forall a b. (a -> b) -> [a] -> [b]
map ((b -> c) -> Tree a b -> Tree a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second b -> c
g) Forest a b
ts

-- | Combine node labels in pre-order.
instance Foldable (Tree e) where
  foldMap :: (a -> m) -> Tree e a -> m
foldMap a -> m
f ~(Node e
_ a
lb Forest e a
ts) = a -> m
f a
lb m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (Tree e a -> m) -> Forest e a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((a -> m) -> Tree e a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f) Forest e a
ts
  null :: Tree e a -> Bool
null Tree e a
_ = Bool
False
  {-# INLINE null #-}
  toList :: Tree e a -> [a]
toList = Tree e a -> [a]
forall e a. Tree e a -> [a]
labels
  {-# INLINE toList #-}

instance Bifoldable Tree where
  bifoldMap :: (a -> m) -> (b -> m) -> Tree a b -> m
bifoldMap a -> m
f b -> m
g ~(Node a
br b
lb Forest a b
ts) = a -> m
f a
br m -> m -> m
forall a. Semigroup a => a -> a -> a
<> b -> m
g b
lb m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (Tree a b -> m) -> Forest a b -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((a -> m) -> (b -> m) -> Tree a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap a -> m
f b -> m
g) Forest a b
ts

instance Traversable (Tree e) where
  traverse :: (a -> f b) -> Tree e a -> f (Tree e b)
traverse a -> f b
g ~(Node e
br a
lb Forest e a
ts) = e -> b -> Forest e b -> Tree e b
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br (b -> Forest e b -> Tree e b) -> f b -> f (Forest e b -> Tree e b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
g a
lb f (Forest e b -> Tree e b) -> f (Forest e b) -> f (Tree e b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Tree e a -> f (Tree e b)) -> Forest e a -> f (Forest e b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((a -> f b) -> Tree e a -> f (Tree e b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
g) Forest e a
ts

instance Bitraversable Tree where
  bitraverse :: (a -> f c) -> (b -> f d) -> Tree a b -> f (Tree c d)
bitraverse a -> f c
f b -> f d
g ~(Node a
br b
lb Forest a b
ts) = c -> d -> Forest c d -> Tree c d
forall e a. e -> a -> Forest e a -> Tree e a
Node (c -> d -> Forest c d -> Tree c d)
-> f c -> f (d -> Forest c d -> Tree c d)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f c
f a
br f (d -> Forest c d -> Tree c d)
-> f d -> f (Forest c d -> Tree c d)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> b -> f d
g b
lb f (Forest c d -> Tree c d) -> f (Forest c d) -> f (Tree c d)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Tree a b -> f (Tree c d)) -> Forest a b -> f (Forest c d)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((a -> f c) -> (b -> f d) -> Tree a b -> f (Tree c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse a -> f c
f b -> f d
g) Forest a b
ts

-- The following code provides a zip-like applicative instance. However,
-- the zip-like instance makes the Monad instance meaningless. So, either we
-- provide only 'Applicative' in zip-like form, or we use the classic instance
-- for 'Applicative' and 'Monad'.

-- -- | Note: The 'Applicative' instance of 'Tree' is similar to the one of
-- -- 'Control.Applicative.ZipList', and differs from the instance of
-- -- 'Data.Tree.Tree'!
-- --
-- -- >>> let t = Node "" 0 [Node "" 1 [], Node "" 2 []] :: Tree String Int
-- -- >>> let f = Node "+3" (+3) [Node "*5" (*5) [], Node "+10" (+10) []] :: Tree String (Int -> Int)
-- -- >>> f <*> t
-- -- Node {branch = "+3", label = 3, forest = [Node {branch = "*5", label = 5, forest = []},Node {branch = "+10", label = 12, forest = []}]}
-- --
-- -- Note: The 'Monoid' instance of the branch labels determines how the branches
-- -- are combined. For example, distances can be summed using the
-- -- 'Data.Monoid.Sum' monoid.
-- instance Monoid e => Applicative (Tree e) where
--   pure lb = Node mempty lb []
--   ~(Node brF lbF tsF) <*> ~(Node brX lbX tsX) =
--     Node (brF <> brX) (lbF lbX) (zipWith (<*>) tsF tsX)
--   liftA2 f ~(Node brX lbX tsX) ~(Node brY lbY tsY) =
--     Node (brX <> brY) (f lbX lbY) (zipWith (liftA2 f) tsX tsY)
--   ~(Node brX _ tsX) *> ~(Node brY lbY tsY) =
--     Node (brX <> brY) lbY (zipWith (*>) tsX tsY)
--   ~(Node brX lbX tsX) <* ~(Node brY _ tsY) =
--     Node (brX <> brY) lbX (zipWith (<*) tsX tsY)

-- | 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'.
instance (Semigroup e, Monoid e) => Applicative (Tree e) where
  pure :: a -> Tree e a
pure a
lb = e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
forall a. Monoid a => a
mempty a
lb []
  ~(Node e
brF a -> b
lbF Forest e (a -> b)
tsF) <*> :: Tree e (a -> b) -> Tree e a -> Tree e b
<*> ~tx :: Tree e a
tx@(Node e
brX a
lbX Forest e a
tsX) =
    e -> b -> Forest e b -> Tree e b
forall e a. e -> a -> Forest e a -> Tree e a
Node (e
brF e -> e -> e
forall a. Semigroup a => a -> a -> a
<> e
brX) (a -> b
lbF a
lbX) ((Tree e a -> Tree e b) -> Forest e a -> Forest e b
forall a b. (a -> b) -> [a] -> [b]
map (a -> b
lbF (a -> b) -> Tree e a -> Tree e b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) Forest e a
tsX Forest e b -> Forest e b -> Forest e b
forall a. [a] -> [a] -> [a]
++ (Tree e (a -> b) -> Tree e b) -> Forest e (a -> b) -> Forest e b
forall a b. (a -> b) -> [a] -> [b]
map (Tree e (a -> b) -> Tree e a -> Tree e b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree e a
tx) Forest e (a -> b)
tsF)
  liftA2 :: (a -> b -> c) -> Tree e a -> Tree e b -> Tree e c
liftA2 a -> b -> c
f ~(Node e
brX a
lbX Forest e a
tsX) ~ty :: Tree e b
ty@(Node e
brY b
lbY Forest e b
tsY) =
    e -> c -> Forest e c -> Tree e c
forall e a. e -> a -> Forest e a -> Tree e a
Node (e
brX e -> e -> e
forall a. Semigroup a => a -> a -> a
<> e
brY) (a -> b -> c
f a
lbX b
lbY) ((Tree e b -> Tree e c) -> Forest e b -> Forest e c
forall a b. (a -> b) -> [a] -> [b]
map (a -> b -> c
f a
lbX (b -> c) -> Tree e b -> Tree e c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) Forest e b
tsY Forest e c -> Forest e c -> Forest e c
forall a. [a] -> [a] -> [a]
++ (Tree e a -> Tree e c) -> Forest e a -> Forest e c
forall a b. (a -> b) -> [a] -> [b]
map (\Tree e a
tx -> (a -> b -> c) -> Tree e a -> Tree e b -> Tree e c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
f Tree e a
tx Tree e b
ty) Forest e a
tsX)
  ~(Node e
brX a
_ Forest e a
tsX) *> :: Tree e a -> Tree e b -> Tree e b
*> ~ty :: Tree e b
ty@(Node e
brY b
lbY Forest e b
tsY) =
    e -> b -> Forest e b -> Tree e b
forall e a. e -> a -> Forest e a -> Tree e a
Node (e
brX e -> e -> e
forall a. Semigroup a => a -> a -> a
<> e
brY) b
lbY (Forest e b
tsY Forest e b -> Forest e b -> Forest e b
forall a. [a] -> [a] -> [a]
++ (Tree e a -> Tree e b) -> Forest e a -> Forest e b
forall a b. (a -> b) -> [a] -> [b]
map (Tree e a -> Tree e b -> Tree e b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Tree e b
ty) Forest e a
tsX)
  ~(Node e
brX a
lbX Forest e a
tsX) <* :: Tree e a -> Tree e b -> Tree e a
<* ~ty :: Tree e b
ty@(Node e
brY b
_ Forest e b
tsY) =
    e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node (e
brX e -> e -> e
forall a. Semigroup a => a -> a -> a
<> e
brY) a
lbX ((Tree e b -> Tree e a) -> Forest e b -> Forest e a
forall a b. (a -> b) -> [a] -> [b]
map (a
lbX a -> Tree e b -> Tree e a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$) Forest e b
tsY Forest e a -> Forest e a -> Forest e a
forall a. [a] -> [a] -> [a]
++ (Tree e a -> Tree e a) -> Forest e a -> Forest e a
forall a b. (a -> b) -> [a] -> [b]
map (Tree e a -> Tree e b -> Tree e a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Tree e b
ty) Forest e a
tsX)

-- | 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 e
br a
lb Forest e a
ts) >>= :: Tree e a -> (a -> Tree e b) -> Tree e b
>>= a -> Tree e b
f = case a -> Tree e b
f a
lb of
    Node e
br' b
lb' Forest e b
ts' -> e -> b -> Forest e b -> Tree e b
forall e a. e -> a -> Forest e a -> Tree e a
Node (e
br e -> e -> e
forall a. Semigroup a => a -> a -> a
<> e
br') b
lb' (Forest e b
ts' Forest e b -> Forest e b -> Forest e b
forall a. [a] -> [a] -> [a]
++ (Tree e a -> Tree e b) -> Forest e a -> Forest e b
forall a b. (a -> b) -> [a] -> [b]
map (Tree e a -> (a -> Tree e b) -> Tree e b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Tree e b
f) Forest e a
ts)

-- -- Cannot provide MonadZip instance because branch labels cannot be
-- -- recovered from 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)

instance Monoid e => MonadFix (Tree e) where
  mfix :: (a -> Tree e a) -> Tree e a
mfix = (a -> Tree e a) -> Tree e a
forall a e. (a -> Tree e a) -> Tree e a
mfixTree

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

instance Comonad (Tree e) where
  duplicate :: Tree e a -> Tree e (Tree e a)
duplicate t :: Tree e a
t@(Node e
br a
_ Forest e a
ts) = e -> Tree e a -> Forest e (Tree e a) -> Tree e (Tree e a)
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br Tree e a
t ((Tree e a -> Tree e (Tree e a))
-> Forest e a -> Forest e (Tree e a)
forall a b. (a -> b) -> [a] -> [b]
map Tree e a -> Tree e (Tree e a)
forall (w :: * -> *) a. Comonad w => w a -> w (w a)
duplicate Forest e a
ts)
  extract :: Tree e a -> a
extract (Node e
_ a
lb Forest e a
_) = a
lb
  {-# INLINE extract #-}

instance (NFData e, NFData a) => NFData (Tree e a) where
  rnf :: Tree e a -> ()
rnf (Node e
br a
lb Forest e a
ts) = e -> ()
forall a. NFData a => a -> ()
rnf e
br () -> () -> ()
`seq` a -> ()
forall a. NFData a => a -> ()
rnf a
lb () -> () -> ()
`seq` Forest e a -> ()
forall a. NFData a => a -> ()
rnf Forest e a
ts

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

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

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

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

-- TODO: Maybe use foldr similar to 'labels'.
-- | Get leaves.
leaves :: Tree e a -> [a]
leaves :: Tree e a -> [a]
leaves (Node e
_ a
lb []) = [a
lb]
leaves (Node e
_ a
_ [Tree e a]
ts) = (Tree e a -> [a]) -> [Tree e a] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree e a -> [a]
forall e a. Tree e a -> [a]
leaves [Tree e a]
ts

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

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

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

-- | Change the root branch of a tree.
applyStem :: (e -> e) -> Tree e a -> Tree e a
applyStem :: (e -> e) -> Tree e a -> Tree e a
applyStem e -> e
f Tree e a
t = Tree e a
t {branch :: e
branch = e -> e
f (e -> e) -> e -> e
forall a b. (a -> b) -> a -> b
$ Tree e a -> e
forall e a. Tree e a -> e
branch Tree e a
t}

-- | Get branch labels in pre-order.
branches :: Tree e a -> [e]
branches :: Tree e a -> [e]
branches Tree e a
t = Tree e a -> [e] -> [e]
forall a a. Tree a a -> [a] -> [a]
squish Tree e a
t []
  where
    squish :: Tree a a -> [a] -> [a]
squish (Node a
br a
_ Forest a a
ts) [a]
xs = a
br a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (Tree a a -> [a] -> [a]) -> [a] -> Forest a a -> [a]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tree a a -> [a] -> [a]
squish [a]
xs Forest a a
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 :: [f] -> t e a -> Maybe (t f a)
setBranches [f]
xs = t (Maybe f) (Maybe a) -> Maybe (t f a)
forall (t :: * -> * -> *) (f :: * -> *) a b.
(Bitraversable t, Applicative f) =>
t (f a) (f b) -> f (t a b)
bisequenceA (t (Maybe f) (Maybe a) -> Maybe (t f a))
-> (t e a -> t (Maybe f) (Maybe a)) -> t e a -> Maybe (t f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([f], t (Maybe f) (Maybe a)) -> t (Maybe f) (Maybe a)
forall a b. (a, b) -> b
snd (([f], t (Maybe f) (Maybe a)) -> t (Maybe f) (Maybe a))
-> (t e a -> ([f], t (Maybe f) (Maybe a)))
-> t e a
-> t (Maybe f) (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([f] -> e -> ([f], Maybe f))
-> ([f] -> a -> ([f], Maybe a))
-> [f]
-> t e a
-> ([f], t (Maybe f) (Maybe a))
forall (t :: * -> * -> *) a b c d e.
Bitraversable t =>
(a -> b -> (a, c))
-> (a -> d -> (a, e)) -> a -> t b d -> (a, t c e)
bimapAccumL [f] -> e -> ([f], Maybe f)
forall a p. [a] -> p -> ([a], Maybe a)
setBranch [f] -> a -> ([f], Maybe a)
forall a a. a -> a -> (a, Maybe a)
noChange [f]
xs
  where
    setBranch :: [a] -> p -> ([a], Maybe a)
setBranch [] p
_ = ([], Maybe a
forall a. Maybe a
Nothing)
    setBranch (a
y : [a]
ys) p
_ = ([a]
ys, a -> Maybe a
forall a. a -> Maybe a
Just a
y)
    noChange :: a -> a -> (a, Maybe a)
noChange a
ys a
z = (a
ys, a -> Maybe a
forall a. a -> Maybe a
Just a
z)

-- | Set the label to a given value.
setLabel :: a -> Tree e a -> Tree e a
setLabel :: a -> Tree e a -> Tree e a
setLabel a
lb (Node e
br a
_ Forest e a
ts) = e -> a -> Forest e a -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb Forest e a
ts

-- | Change the root branch of a tree.
applyLabel :: (a -> a) -> Tree e a -> Tree e a
applyLabel :: (a -> a) -> Tree e a -> Tree e a
applyLabel a -> a
f Tree e a
t = Tree e a
t {label :: a
label = a -> a
f (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ Tree e a -> a
forall e a. Tree e a -> a
label Tree e a
t}

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

-- | 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 :: [b] -> t a -> Maybe (t b)
setLabels [b]
xs = t (Maybe b) -> Maybe (t b)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (t (Maybe b) -> Maybe (t b))
-> (t a -> t (Maybe b)) -> t a -> Maybe (t b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([b], t (Maybe b)) -> t (Maybe b)
forall a b. (a, b) -> b
snd (([b], t (Maybe b)) -> t (Maybe b))
-> (t a -> ([b], t (Maybe b))) -> t a -> t (Maybe b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([b] -> a -> ([b], Maybe b)) -> [b] -> t a -> ([b], t (Maybe b))
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL [b] -> a -> ([b], Maybe b)
forall a p. [a] -> p -> ([a], Maybe a)
setLabelM [b]
xs
  where
    setLabelM :: [a] -> p -> ([a], Maybe a)
setLabelM [] p
_ = ([], Maybe a
forall a. Maybe a
Nothing)
    setLabelM (a
y : [a]
ys) p
_ = ([a]
ys, a -> Maybe a
forall a. a -> Maybe a
Just a
y)

-- | Change the root label of a tree.
applyRoot :: (a -> a) -> Tree e a -> Tree e a
applyRoot :: (a -> a) -> Tree e a -> Tree e a
applyRoot a -> a
f Tree e a
t = Tree e a
t {label :: a
label = a -> a
f (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ Tree e a -> a
forall e a. Tree e a -> a
label Tree e a
t}

-- | Label the nodes with unique integers starting at the root with 0.
identify :: Traversable t => t a -> t Int
identify :: t a -> t Int
identify = (Int, t Int) -> t Int
forall a b. (a, b) -> b
snd ((Int, t Int) -> t Int) -> (t a -> (Int, t Int)) -> t a -> t Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> a -> (Int, Int)) -> Int -> t a -> (Int, t Int)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL (\Int
i a
_ -> (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
i)) (Int
0 :: Int)

-- | The degree of the root node.
degree :: Tree e a -> Int
degree :: Tree e a -> Int
degree = (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int) -> (Tree e a -> Int) -> Tree e a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree e a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Tree e a] -> Int) -> (Tree e a -> [Tree e a]) -> Tree e a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree e a -> [Tree e a]
forall e a. Tree e a -> Forest e a
forest

-- | Prune degree two nodes.
--
-- The information stored in 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 :: Tree e a -> Tree e a
prune t :: Tree e a
t@(Node e
_ a
_ []) = Tree e a
t
prune (Node e
paBr a
_ [Node e
daBr a
daLb [Tree e a]
daTs]) = e -> a -> [Tree e a] -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node (e
daBr e -> e -> e
forall a. Semigroup a => a -> a -> a
<> e
paBr) a
daLb [Tree e a]
daTs
prune (Node e
paBr a
paLb [Tree e a]
paTs) = e -> a -> [Tree e a] -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node e
paBr a
paLb ([Tree e a] -> Tree e a) -> [Tree e a] -> Tree e a
forall a b. (a -> b) -> a -> b
$ (Tree e a -> Tree e a) -> [Tree e a] -> [Tree e a]
forall a b. (a -> b) -> [a] -> [b]
map Tree e a -> Tree e a
forall e a. Semigroup e => Tree e a -> Tree e a
prune [Tree e a]
paTs

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

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

-- | Zip two trees with the same topology.
--
-- 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 :: (e1 -> e2 -> e)
-> (a1 -> a2 -> a) -> Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree e a)
zipTreesWith e1 -> e2 -> e
f a1 -> a2 -> a
g (Node e1
brL a1
lbL Forest e1 a1
tsL) (Node e2
brR a2
lbR Forest e2 a2
tsR) =
  if Forest e1 a1 -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest e1 a1
tsL Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Forest e2 a2 -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest e2 a2
tsR
    then -- I am proud of that :)).
      (Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree e a))
-> Forest e1 a1 -> Forest e2 a2 -> Maybe [Tree e a]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM ((e1 -> e2 -> e)
-> (a1 -> a2 -> a) -> Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree e a)
forall e1 e2 e a1 a2 a.
(e1 -> e2 -> e)
-> (a1 -> a2 -> a) -> Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree e a)
zipTreesWith e1 -> e2 -> e
f a1 -> a2 -> a
g) Forest e1 a1
tsL Forest e2 a2
tsR Maybe [Tree e a]
-> ([Tree e a] -> Maybe (Tree e a)) -> Maybe (Tree e a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Tree e a -> Maybe (Tree e a)
forall a. a -> Maybe a
Just (Tree e a -> Maybe (Tree e a))
-> ([Tree e a] -> Tree e a) -> [Tree e a] -> Maybe (Tree e a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> a -> [Tree e a] -> Tree e a
forall e a. e -> a -> Forest e a -> Tree e a
Node (e1 -> e2 -> e
f e1
brL e2
brR) (a1 -> a2 -> a
g a1
lbL a2
lbR)
    else Maybe (Tree e a)
forall a. Maybe a
Nothing

-- | Zip two trees with the same topology.
--
-- Return 'Nothing' if the topologies are different.
zipTrees :: Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree (e1, e2) (a1, a2))
zipTrees :: Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree (e1, e2) (a1, a2))
zipTrees = (e1 -> e2 -> (e1, e2))
-> (a1 -> a2 -> (a1, a2))
-> Tree e1 a1
-> Tree e2 a2
-> Maybe (Tree (e1, e2) (a1, a2))
forall e1 e2 e a1 a2 a.
(e1 -> e2 -> e)
-> (a1 -> a2 -> a) -> Tree e1 a1 -> Tree e2 a2 -> Maybe (Tree e a)
zipTreesWith (,) (,)