wolf-graph 0.1.0

use std::borrow::Cow;

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

use crate::{EdgeID, Edges, Error, MutableGraph, NodeID, Nodes, PathExists, BlankGraph, TopologicalSort, VisitableGraph};

/// A DAG (Directed Acyclic Graph) data structure.
#[derive(Debug, Clone)]
pub struct DAG<Inner>
where
    Inner: MutableGraph,
{
    graph: Inner,
    roots: Nodes,
    leaves: Nodes,
}

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

impl<Inner> DAG<Inner>
where
    Inner: MutableGraph,
{
    pub fn new_with_graph(graph: Inner) -> Result<Self> {
        if !graph.is_dag() {
            return Err(Error::NotADAG.into());
        }

        let roots = graph.all_roots().clone();
        let leaves = graph.all_leaves().clone();

        Ok(Self {
            graph,
            roots,
            leaves,
        })
    }

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

impl<Inner> DAG<Inner>
where
    Inner: MutableGraph + Default,
{
    pub fn new() -> Self {
        Self::new_with_graph(Inner::default()).unwrap()
    }

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

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

impl<Inner> VisitableGraph for DAG<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 {
        self.roots.iter().cloned().collect()
    }

    fn all_leaves(&self) -> Nodes {
        self.leaves.iter().cloned().collect()
    }

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

    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> {
        Ok(self.leaves.contains(id.as_ref()))
    }

    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> MutableGraph for DAG<Inner>
where
    Inner: MutableGraph,
{
    fn add_node_with_data(&mut self, id: impl AsRef<NodeID>, data: Self::NData) -> Result<()> {
        let id = id.as_ref();
        self.graph.add_node_with_data(id, data)?;
        self.roots.insert(id.clone());
        self.leaves.insert(id.clone());
        Ok(())
    }

    fn add_edge_with_data(&mut self, id: impl AsRef<EdgeID>, source: impl AsRef<NodeID>, target: impl AsRef<NodeID>, data: Self::EData) -> Result<()> {
        let id = id.as_ref();
        let source = source.as_ref();
        let target = target.as_ref();
        if !self.graph.can_add_dag_edge(source, target, None::<EdgeID>)? {
            return Err(Error::NotADAG.into());
        }
        self.graph.add_edge_with_data(id, source, target, data)?;
        self.roots.remove(target.as_ref());
        self.leaves.remove(source.as_ref());
        Ok(())
    }

    fn remove_edge(&mut self, id: impl AsRef<EdgeID>) -> Result<()> {
        let id = id.as_ref();
        let (old_source, old_target) = self.graph.endpoints(id)?;

        self.graph.remove_edge(id)?;

        if !self.graph.has_successors(&old_source)? {
            self.leaves.insert(old_source);
        }

        if !self.graph.has_predecessors(&old_target)? {
            self.roots.insert(old_target);
        }

        Ok(())
    }

    fn clear_edges(&mut self, id: impl AsRef<NodeID>) -> Result<()> {
        for edge in self.graph.incident_edges(id)? {
            self.remove_edge(&edge)?;
        }
        Ok(())
    }

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

        self.clear_edges(id)?;
        self.graph.remove_node(id)?;
        self.roots.remove(id.as_ref());
        self.leaves.remove(id.as_ref());
        Ok(())
    }

    fn move_edge(&mut self, id: impl AsRef<EdgeID>, new_source: impl AsRef<NodeID>, new_target: impl AsRef<NodeID>) -> Result<()> {
        let id = id.as_ref();
        let new_source = new_source.as_ref();
        let new_target = new_target.as_ref();

        if !self.graph.can_add_dag_edge(new_source, new_target, None::<EdgeID>)? {
            return Err(Error::NotADAG.into());
        }

        let (old_source, old_target) = self.graph.endpoints(id)?;

        self.graph.move_edge(id, new_source, new_target)?;

        if &old_source != new_source {
            if !self.graph.has_successors(&old_source)? {
                self.leaves.insert(old_source);
            }
            self.leaves.remove(new_source);
        }

        if &old_target != new_target {
            if !self.graph.has_predecessors(&old_target)? {
                self.roots.insert(old_target);
            }
            self.roots.remove(new_target);
        }

        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 DAG<Inner>
where
    Inner: MutableGraph + PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.graph == other.graph
    }
}

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

// Implement builder methods
impl<Inner> DAG<Inner>
where
    Inner: MutableGraph + Clone,
{
    pub fn adding_node_with_data(&self, id: impl AsRef<NodeID>, data: Inner::NData) -> Result<Self> {
        let mut dag = self.clone();
        dag.add_node_with_data(id, data)?;
        Ok(dag)
    }

    pub fn adding_edge_with_data(&self, id: impl AsRef<EdgeID>, source: impl AsRef<NodeID>, target: impl AsRef<NodeID>, data: Inner::EData) -> Result<Self> {
        let mut dag = self.clone();
        dag.add_edge_with_data(id, source, target, data)?;
        Ok(dag)
    }

    pub fn removing_edge(&self, id: impl AsRef<EdgeID>) -> Result<Self> {
        let mut dag = self.clone();
        dag.remove_edge(id)?;
        Ok(dag)
    }

    pub fn removing_node(&self, id: impl AsRef<NodeID>) -> Result<Self> {
        let mut dag = self.clone();
        dag.remove_node(id)?;
        Ok(dag)
    }

    pub fn moving_edge(&self, id: impl AsRef<EdgeID>, new_source: impl AsRef<NodeID>, new_target: impl AsRef<NodeID>) -> Result<Self> {
        let mut dag = self.clone();
        dag.move_edge(id, new_source, new_target)?;
        Ok(dag)
    }
}

// Implement default methods
impl<Inner> DAG<Inner>
where
    Inner: MutableGraph + Clone,
    Inner::NData: Default,
{
    pub fn add_node(&mut self, id: impl AsRef<NodeID>) -> Result<()> {
        self.add_node_with_data(id, Inner::NData::default())
    }

    pub fn adding_node(&self, id: impl AsRef<NodeID>) -> Result<Self> {
        self.adding_node_with_data(id, Inner::NData::default())
    }
}

impl<Inner> DAG<Inner>
where
    Inner: MutableGraph + Clone,
    Inner::EData: Default,
{
    pub fn add_edge(&mut self, id: impl AsRef<EdgeID>, source: impl AsRef<NodeID>, target: impl AsRef<NodeID>) -> Result<()> {
        self.add_edge_with_data(id, source, target, Inner::EData::default())
    }

    pub fn adding_edge(&self, id: impl AsRef<EdgeID>, source: impl AsRef<NodeID>, target: impl AsRef<NodeID>) -> Result<Self> {
        self.adding_edge_with_data(id, source, target, Inner::EData::default())
    }
}

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

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

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