wolf-graph 0.1.0

use std::borrow::Cow;

use anyhow::{Result, bail};
#[cfg(feature = "serde")]
use serde::{ser::{Serialize, Serializer}, de::{self, Deserialize, Deserializer}};

use crate::{nid, EdgeID, Edges, Error, IsTree, MutableForest, MutableGraph, MutableTree, NodeID, Nodes, PathExists, BlankGraph, TopologicalSort, VisitableForest, VisitableGraph, VisitableTree};

/// A tree is a directed graph in which there is exactly one path between any two nodes.
///
/// A tree always has a root node, which is the only node that has no incoming edges.
#[derive(Debug, Clone)]
pub struct Tree<Inner>
where
    Inner: MutableGraph,
{
    root: NodeID,
    graph: Inner,
}

/// A convenience type for a tree with no additional data on nodes or edges.
pub type BlankTree = Tree<BlankGraph>;

impl<Inner> Tree<Inner>
where
    Inner: MutableGraph,
{
    pub fn new_unchecked(root: NodeID, graph: Inner) -> Self {
        Self { root, graph }
    }

    pub fn new_with_root_and_graph(root: NodeID, graph: Inner) -> Result<Self> {
        graph.check_is_tree(&root)?;
        Ok(Self::new_unchecked(root, graph))
    }

    pub fn graph(&self) -> &Inner {
        &self.graph
    }
}

impl<Inner> Tree<Inner>
where
    Inner: MutableGraph + Default + Clone,
    Inner::NData: Default,
{
    pub fn new_with_root(root: NodeID) -> Self {
        Self::new_unchecked(root, Inner::default().adding_node(&nid!("root")).unwrap())
    }

    pub fn new() -> Self {
        Self::new_with_root(nid!("root"))
    }
}

impl<Inner> Default for Tree<Inner>
where
    Inner: MutableGraph + Default + Clone,
    Inner::NData: Default,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<Inner> VisitableGraph for Tree<Inner>
where
    Inner: MutableGraph,
{
    type GData = Inner::GData;
    type NData = Inner::NData;
    type EData = Inner::EData;

    fn data(&self) -> &Self::GData {
        self.graph.data()
    }

    fn node_data(&self, id: impl AsRef<NodeID>) -> Result<Cow<'static, Self::NData>> {
        self.graph.node_data(id)
    }

    fn edge_data(&self, id: impl AsRef<EdgeID>) -> Result<Cow<'static, Self::EData>> {
        self.graph.edge_data(id)
    }

    fn is_empty(&self) -> bool {
        self.graph.is_empty()
    }

    fn node_count(&self) -> usize {
        self.graph.node_count()
    }

    fn edge_count(&self) -> usize {
        self.graph.edge_count()
    }

    fn all_nodes(&self) -> Nodes {
        self.graph.all_nodes()
    }

    fn all_edges(&self) -> Edges {
        self.graph.all_edges()
    }

    fn has_node(&self, id: impl AsRef<NodeID>) -> bool {
        self.graph.has_node(id)
    }

    fn has_edge(&self, id: impl AsRef<EdgeID>) -> bool {
        self.graph.has_edge(id)
    }

    fn has_edge_from_to(&self, source: impl AsRef<NodeID>, target: impl AsRef<NodeID>) -> bool {
        self.graph.has_edge_from_to(source, target)
    }

    fn has_edge_between(&self, a: impl AsRef<NodeID>, b: impl AsRef<NodeID>) -> bool {
        self.graph.has_edge_between(a, b)
    }

    fn source(&self, id: impl AsRef<EdgeID>) -> Result<NodeID> {
        self.graph.source(id)
    }

    fn target(&self, id: impl AsRef<EdgeID>) -> Result<NodeID> {
        self.graph.target(id)
    }

    fn endpoints(&self, id: impl AsRef<EdgeID>) -> Result<(NodeID, NodeID)> {
        self.graph.endpoints(id)
    }

    fn out_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        self.graph.out_edges(id)
    }

    fn in_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        self.graph.in_edges(id)
    }

    fn incident_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        self.graph.incident_edges(id)
    }

    fn out_degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        self.graph.out_degree(id)
    }

    fn in_degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        self.graph.in_degree(id)
    }

    fn degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        self.graph.degree(id)
    }

    fn successors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        self.graph.successors(id)
    }

    fn predecessors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        self.graph.predecessors(id)
    }

    fn neighbors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        self.graph.neighbors(id)
    }

    fn has_successors(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        self.graph.has_successors(id)
    }

    fn has_predecessors(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        self.graph.has_predecessors(id)
    }

    fn has_neighbors(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        self.graph.has_neighbors(id)
    }

    fn all_roots(&self) -> Nodes {
        vec![self.root.clone()].into_iter().collect()
    }

    fn all_leaves(&self) -> Nodes {
        self.graph.all_leaves()
    }

    fn non_roots(&self) -> Nodes {
        self.all_nodes().into_iter().filter(|n| n != &self.root).collect()
    }

    fn non_leaves(&self) -> Nodes {
        self.graph.non_leaves()
    }

    fn all_internals(&self) -> Nodes {
        self.graph.all_internals()
    }

    fn is_leaf(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        self.graph.is_leaf(id)
    }

    fn is_root(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        self.graph.is_root(id)
    }

    fn is_internal(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        self.graph.is_internal(id)
    }
}

impl<Inner> VisitableTree for Tree<Inner>
where
    Inner: MutableGraph,
{
    fn root(&self) -> NodeID {
        self.root.clone()
    }
}

impl<Inner> VisitableForest for Tree<Inner>
where
    Inner: MutableGraph,
{
    fn in_edge(&self, node: impl AsRef<NodeID>) -> Result<Option<EdgeID>> {
        Ok(self.in_edges(node)?.first().cloned())
    }

    fn in_edge_with_root(&self, node: impl AsRef<NodeID>) -> Result<Option<EdgeID>> {
        self.in_edge(node)
    }

    fn parent(&self, node: impl AsRef<NodeID>) -> Result<Option<NodeID>> {
        Ok(self.in_edge(node)?.map(|edge| self.source(&edge).unwrap()))
    }

    fn children(&self, node: Option<impl AsRef<NodeID>>) -> Result<Nodes> {
        if let Some(node) = node {
            self.successors(node)
        } else {
            self.successors(self.root())
        }
    }

    fn has_children(&self, node: impl AsRef<NodeID>) -> Result<bool> {
        self.has_successors(node)
    }

    fn child_count(&self, node: impl AsRef<NodeID>) -> Result<usize> {
        self.out_degree(node)
    }
}

impl<Inner> MutableTree for Tree<Inner>
where
    Inner: MutableGraph,
{
    fn set_root(&mut self, root: impl AsRef<NodeID>) -> Result<()> {
        let root = root.as_ref();
        self.graph.check_is_tree(root)?;
        self.root = root.as_ref().clone();
        Ok(())
    }
}

impl<Inner> MutableForest for Tree<Inner>
where
    Inner: MutableGraph,
{
    fn add_node_with_node_and_edge_data(
        &mut self,
        node: impl AsRef<NodeID>,
        parent: Option<impl AsRef<NodeID>>,
        edge: impl AsRef<EdgeID>,
        node_data: Self::NData,
        edge_data: Self::EData,
    ) -> Result<()> {
        let node = node.as_ref();
        self.graph.add_node_with_data(node, node_data)?;
        let parent = parent.map(|p| p.as_ref().clone()).unwrap_or_else(|| self.root());
        self.graph.add_edge_with_data(edge, parent, node, edge_data)?;
        Ok(())
    }

    fn remove_node_ungrouping(&mut self, id: impl AsRef<NodeID>) -> Result<()> {
        let id = id.as_ref();
        if id == &self.root {
            let children = self.children(Some(id))?;
            if children.len() != 1 {
                bail!(Error::NotATree);
            }
            let new_root = children.into_iter().next().unwrap();
            self.graph.remove_node(id)?;
            self.set_root(&new_root)?;
        } else {
            let new_parent = self.parent(id)?.unwrap();
            let children = self.children(Some(id))?;
            for child in children {
                self.move_node(&child, Some(&new_parent))?;
            }
            self.graph.remove_node(id)?;
        }
        Ok(())
    }

    fn remove_node_and_children(&mut self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        let id = id.as_ref();

        // Can't remove root
        if id == &self.root {
            bail!(Error::NotATree);
        }

        // Remove child nodes in reverse-topological sort order (most distant from the target first).
        let to_remove = self.topological_sort_opt(&Nodes::from([id.clone()]), true)?;
        for node in to_remove.iter() {
            self.graph.remove_node(node)?;
        }
        Ok(to_remove.into_iter().collect())
    }

    fn remove_children(&mut self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        let id = id.as_ref();

        // Remove child nodes in reverse-topological sort order (most distant from the target first).
        let children = self.children(Some(id))?;
        let to_remove = self.topological_sort_opt(&children, true)?;
        for node in to_remove.iter() {
            self.graph.remove_node(node)?;
        }
        Ok(to_remove.into_iter().collect())
    }

    fn move_node(&mut self, id: impl AsRef<NodeID>, new_parent: Option<impl AsRef<NodeID>>) -> Result<()> {
        let id = id.as_ref();

        // Can't move root
        if id == &self.root {
            bail!(Error::NotATree);
        }

        let edge = self.in_edge(id)?.unwrap();
        let root = self.root();
        let new_parent = new_parent.map(|p| p.as_ref().clone()).unwrap_or_else(|| root.clone());
        let new_parent = new_parent.as_ref();
        if !self.graph.can_move_dag_edge(&edge, new_parent, id)? {
            bail!(Error::NotATree);
        }
        self.graph.move_edge(&edge, new_parent, id)?;
        Ok(())
    }

    fn set_data(&mut self, data: Self::GData) {
        self.graph.set_data(data);
    }

    fn set_node_data(&mut self, id: impl AsRef<NodeID>, data: Self::NData) -> Result<()> {
        self.graph.set_node_data(id, data)
    }

    fn set_edge_data(&mut self, id: impl AsRef<EdgeID>, data: Self::EData) -> Result<()> {
        self.graph.set_edge_data(id, data)
    }

    fn with_data(&mut self, transform: &dyn Fn(&mut Self::GData)) {
        self.graph.with_data(transform);
    }

    fn with_node_data(&mut self, id: impl AsRef<NodeID>, transform: &dyn Fn(&mut Self::NData)) -> Result<()> {
        self.graph.with_node_data(id, transform)
    }

    fn with_edge_data(&mut self, id: impl AsRef<EdgeID>, transform: &dyn Fn(&mut Self::EData)) -> Result<()> {
        self.graph.with_edge_data(id, transform)
    }
}

impl<Inner> PartialEq for Tree<Inner>
where
    Inner: MutableGraph + PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.root == other.root && self.graph == other.graph
    }
}

impl<Inner> Eq for Tree<Inner>
where
    Inner: MutableGraph + Eq,
{
}

#[cfg(feature = "serde")]
impl<Inner> Serialize for Tree<Inner>
where
    Inner: MutableGraph + Serialize,
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        (&self.root, &self.graph).serialize(serializer)
    }
}

#[cfg(feature = "serde")]
impl<'de, Inner> Deserialize<'de> for Tree<Inner>
where
    Inner: MutableGraph + Deserialize<'de>,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let (root, graph) = <(NodeID, Inner)>::deserialize(deserializer)?;
        Tree::new_with_root_and_graph(root, graph).map_err(de::Error::custom)
    }
}

// If Serde and SerdeJSON are both present, add conveniences to serialize a Tree
// to JSON.
#[cfg(all(feature = "serde", feature = "serde_json"))]
impl<Inner> Tree<Inner>
where
    Inner: MutableGraph + Serialize,
{
    pub fn to_json(&self) -> String {
        serde_json::to_string(self).unwrap()
    }
}

// If Serde and SerdeJSON are both present, add conveniences to deserialize a Tree
// from JSON.
#[cfg(all(feature = "serde", feature = "serde_json"))]
impl<'de, Inner> Tree<Inner>
where
    Inner: MutableGraph + Deserialize<'de>,
{
    pub fn from_json(json: &'de str) -> Result<Self, serde_json::Error> {
        serde_json::from_str(json)
    }
}

#[cfg(all(feature = "serde", feature = "serde_json"))]
impl<Inner> std::fmt::Display for Tree<Inner>
where
    Inner: MutableGraph + Serialize,
{
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.to_json())
    }
}