wolf-graph 0.1.0

use std::borrow::Cow;

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

use crate::{EdgeID, Edges, Error, MutableForest, MutableTree, NodeID, Nodes, VisitableForest, VisitableGraph};

/// A forest is a graph that is a collection of trees.
///
/// Unlike a tree, a forest can have multiple roots.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Forest<InnerTree>
where
    InnerTree: MutableTree<GData = (), NData = (), EData = ()>,
{
    tree: InnerTree,
}

/// A convenience type for a forest with no additional data on nodes or edges.
pub type BlankForest = Forest<crate::BlankTree>;

impl<Inner> Forest<Inner>
where
    Inner: MutableTree<GData = (), NData = (), EData = ()>,
{
    pub fn new_with_tree(tree: Inner) -> Self {
        Forest { tree }
    }
}

impl<Inner> Forest<Inner>
where
    Inner: MutableTree<GData = (), NData = (), EData = ()> + Default,
{
    pub fn new() -> Self {
        Self::new_with_tree(Inner::default())
    }
}

impl<Inner> Forest<Inner>
where
    Inner: MutableTree<GData = (), NData = (), EData = ()>,
{
    fn is_tree_root(&self, node: &NodeID) -> bool {
        node == &self.tree.root()
    }

    fn check_not_tree_root(&self, node: &NodeID) -> Result<()> {
        if self.is_tree_root(node) {
            return Err(Error::NodeNotFound(node.clone()).into());
        }
        Ok(())
    }
}

impl<Inner> Default for Forest<Inner>
where
    Inner: MutableTree<GData = (), NData = (), EData = ()> + Default,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<Inner> VisitableGraph for Forest<Inner>
where
    Inner: VisitableGraph + MutableTree<GData = (), NData = (), EData = ()>,
{
    type GData = Inner::GData;
    type NData = Inner::NData;
    type EData = Inner::EData;

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

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

    fn edge_data(&self, id: impl AsRef<EdgeID>) -> Result<Cow<'static, Self::EData>> {
        let id = id.as_ref();
        self.check_not_tree_root(&self.tree.source(id)?)?;
        self.tree.edge_data(id)
    }

    fn is_empty(&self) -> bool {
        self.tree.node_count() < 1
    }

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

    fn edge_count(&self) -> usize {
        self.tree.edge_count() - self.tree.out_degree(&self.tree.root()).unwrap()
    }

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

    fn all_edges(&self) -> Edges {
        self.tree.all_edges()
            .into_iter().filter(|edge| {
                let source = self.tree.source(edge).unwrap();
                !self.is_tree_root(&source)
            })
            .collect()
    }

    fn has_node(&self, node: impl AsRef<NodeID>) -> bool {
        let node = node.as_ref();
        !self.is_tree_root(node) && self.tree.has_node(node)
    }

    fn has_edge(&self, edge: impl AsRef<EdgeID>) -> bool {
        self.source(edge).map_or(false, |source| !self.is_tree_root(&source))
    }

    fn has_edge_from_to(&self, source: impl AsRef<NodeID>, target: impl AsRef<NodeID>) -> bool {
        let source = source.as_ref();
        let target = target.as_ref();
        !self.is_tree_root(source) && self.tree.has_edge_from_to(source, target)
    }

    fn has_edge_between(&self, a: impl AsRef<NodeID>, b: impl AsRef<NodeID>) -> bool {
        let a = a.as_ref();
        let b = b.as_ref();
        !self.is_tree_root(a) && !self.is_tree_root(b) && self.tree.has_edge_between(a, b)
    }

    fn source(&self, id: impl AsRef<EdgeID>) -> Result<NodeID> {
        let id = id.as_ref();
        let source = self.tree.source(id)?;
        if self.is_tree_root(&source) {
            return Err(Error::EdgeNotFound(id.clone()).into());
        }
        Ok(source)
    }

    fn target(&self, id: impl AsRef<EdgeID>) -> Result<NodeID> {
        let id = id.as_ref();
        let source = self.tree.source(id)?;
        if self.is_tree_root(&source) {
            return Err(Error::EdgeNotFound(id.clone()).into());
        }
        self.tree.target(id)
    }

    fn endpoints(&self, id: impl AsRef<EdgeID>) -> Result<(NodeID, NodeID)> {
        let id = id.as_ref();
        let (source, target) = self.tree.endpoints(id)?;
        if self.is_tree_root(&source) {
            return Err(Error::EdgeNotFound(id.clone()).into());
        }
        Ok((source, target))
    }

    fn out_degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        self.tree.out_degree(id)
    }

    fn in_degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        if self.is_root(id)? {
            return Ok(0);
        }
        self.tree.in_degree(id)
    }

    fn degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        Ok(self.in_degree(id)? + self.out_degree(id)?)
    }

    fn out_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        self.tree.out_edges(id)
    }

    fn in_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        if self.is_root(id)? {
            return Ok(Edges::default());
        }
        self.tree.in_edges(id)
    }

    fn incident_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        let incident_edges = self.tree.incident_edges(id)?;
        let root = self.tree.root();
        Ok(incident_edges.into_iter().filter(|edge| self.tree.source(edge).unwrap() != root).collect())
    }

    fn has_successors(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        self.tree.has_successors(id)
    }

    fn has_predecessors(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        if self.is_root(id)? {
            return Ok(false);
        }
        self.tree.has_predecessors(id)
    }

    fn has_neighbors(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        let id = id.as_ref();
        Ok(self.has_successors(id)? || self.has_predecessors(id)?)
    }

    fn successors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        self.tree.successors(id)
    }

    fn predecessors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        if self.is_root(id)? {
            return Ok(Nodes::default());
        }
        self.tree.predecessors(id)
    }

    fn neighbors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        Ok(self.tree.neighbors(id)?.into_iter().filter(|node| *node != self.tree.root()).collect())
    }

    fn all_roots(&self) -> Nodes {
        self.tree.successors(&self.tree.root()).unwrap_or_default()
    }

    fn all_leaves(&self) -> Nodes {
        self.tree.all_leaves().into_iter().filter(|node| *node != self.tree.root()).collect()
    }

    fn non_roots(&self) -> Nodes {
        let roots = self.all_roots();
        self.all_nodes().into_iter().filter(|node| !roots.contains(node)).collect()
    }

    fn non_leaves(&self) -> Nodes {
        let leaves = self.all_leaves();
        self.all_nodes().into_iter().filter(|node| !leaves.contains(node)).collect()
    }

    fn all_internals(&self) -> Nodes {
        let leaves = self.all_leaves();
        self.all_nodes().into_iter().filter(|node| !leaves.contains(node)).collect()
    }

    fn is_leaf(&self, node: impl AsRef<NodeID>) -> Result<bool> {
        let node = node.as_ref();
        self.check_not_tree_root(node)?;
        self.tree.is_leaf(node)
    }

    fn is_root(&self, node: impl AsRef<NodeID>) -> Result<bool> {
        let node = node.as_ref();
        if node == &self.tree.root() {
            return Err(Error::NodeNotFound(node.clone()).into());
        }
        let in_edge = self.tree.in_edge(node)?.unwrap();
        Ok(self.tree.source(&in_edge)? == self.tree.root())
    }

    fn is_internal(&self, node: impl AsRef<NodeID>) -> Result<bool> {
        let node = node.as_ref();
        self.check_not_tree_root(node)?;
        if self.is_root(node)? {
            return Ok(false);
        }
        self.tree.is_internal(node)
    }
}

impl<Inner> VisitableForest for Forest<Inner>
where
    Inner: VisitableForest + MutableTree<GData = (), NData = (), EData = ()>,
{
    fn in_edge(&self, node: impl AsRef<NodeID>) -> Result<Option<EdgeID>> {
        let node = node.as_ref();
        self.check_not_tree_root(node)?;
        self.tree.in_edge(node)
    }

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

    fn parent(&self, node: impl AsRef<NodeID>) -> Result<Option<NodeID>> {
        let node = node.as_ref();
        self.check_not_tree_root(node)?;
        if self.is_root(node)? {
            return Ok(None);
        }
        self.tree.parent(node)
    }

    fn children(&self, node: Option<impl AsRef<NodeID>>) -> Result<Nodes> {
        let node = node.as_ref();
        if let Some(node) = node {
            self.check_not_tree_root(node.as_ref())?;
        }
        self.tree.children(node)
    }

    fn has_children(&self, node: impl AsRef<NodeID>) -> Result<bool> {
        let node = node.as_ref();
        self.check_not_tree_root(node)?;
        self.tree.has_children(node)
    }

    fn child_count(&self, node: impl AsRef<NodeID>) -> Result<usize> {
        let node = node.as_ref();
        self.check_not_tree_root(node)?;
        self.tree.child_count(node)
    }
}

// MutableForest implementation
impl<Inner> MutableForest for Forest<Inner>
where
    Inner: MutableForest + MutableTree<GData = (), NData = (), EData = ()>,
{
    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<()> {
        self.tree.add_node_with_node_and_edge_data(node, parent, edge, node_data, edge_data)
    }

    fn remove_node_ungrouping(&mut self, id: impl AsRef<NodeID>) -> Result<()> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        self.tree.remove_node_ungrouping(id)
    }

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

    fn remove_children(&mut self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        self.tree.remove_children(id)
    }

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

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

    fn set_node_data(&mut self, id: impl AsRef<NodeID>, data: Self::NData) -> Result<()> {
        let id = id.as_ref();
        self.check_not_tree_root(id)?;
        self.tree.set_node_data(id, data)
    }

    fn set_edge_data(&mut self, id: impl AsRef<EdgeID>, data: Self::EData) -> Result<()> {
        let id = id.as_ref();
        let source = self.tree.source(id)?;
        self.check_not_tree_root(&source)?;
        self.tree.set_edge_data(id, data)
    }

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

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

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

#[cfg(feature = "serde")]
impl<Inner> Serialize for Forest<Inner>
where
    Inner: Serialize + MutableTree<GData = (), NData = (), EData = ()>
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::ser::Serializer,
    {
        self.tree.serialize(serializer)
    }
}

#[cfg(feature = "serde")]
impl<'de, Inner> Deserialize<'de> for Forest<Inner>
where
    Inner: Deserialize<'de> + MutableTree<GData = (), NData = (), EData = ()>
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::de::Deserializer<'de>,
    {
        Inner::deserialize(deserializer).map(Self::new_with_tree)
    }
}

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

// If Serde and SerdeJSON are both present, add conveniences to serialize a Tree
// to JSON.
#[cfg(all(feature = "serde", feature = "serde_json"))]
impl<'de, Inner> Forest<Inner>
where
    Inner: MutableTree<GData = (), NData = (), EData = ()> + 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 Forest<Inner>
where
    Inner: MutableTree<GData = (), NData = (), EData = ()> + Serialize,
{
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.to_json().unwrap())
    }
}