wolf-graph 0.1.0

use std::sync::{Arc, RwLock};
use std::borrow::Cow;
use std::collections::BTreeMap;

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

use crate::{Edges, MutableGraph, Nodes, VisitableGraph};
use crate::{
    ids::{EdgeID, NodeID},
    Error,
    details::{edge::Edge, node::Node}
};

/// A general graph data structure with value semantics.
///
/// The graph is parameterized by three types:
/// - `GData` is the type of the graph data.
/// - `NData` is the type of the node data.
/// - `EData` is the type of the edge data.
///
/// The graph is implemented as a set of nodes and edges. Nodes are identified by
/// unique `NodeID` values, and edges are identified by unique `EdgeID` values.
/// The graph and its nodes and edges can have associated data, or they can be
/// the unit type `()`.
#[derive(Debug, Clone)]
pub struct Graph<GData, NData, EData>
where
    GData: Clone + 'static,
    NData: Clone + 'static,
    EData: Clone + 'static,
{
    nodes: Arc<RwLock<BTreeMap<NodeID, Cow<'static, Node<NData>>>>>,
    edges: Arc<RwLock<BTreeMap<EdgeID, Cow<'static, Edge<EData>>>>>,
    data: Cow<'static, GData>,
}

/// A convenience type for a graph with no data.
pub type BlankGraph = Graph<(), (), ()>;

// Internal Utilities
impl<GData, NData, EData> Graph<GData, NData, EData>
where
    GData: Clone + 'static,
    NData: Clone + 'static,
    EData: Clone + 'static,
{
    fn _node(&self, id: &NodeID) -> Result<Cow<'static, Node<NData>>, Error> {
        self.nodes.read().unwrap().get(id).cloned()
            .ok_or(Error::NodeNotFound(id.clone()))
    }

    fn _edge(&self, id: &EdgeID) -> Result<Cow<'static, Edge<EData>>, Error> {
        self.edges.read().unwrap().get(id).cloned()
            .ok_or(Error::EdgeNotFound(id.clone()))
    }

    fn _update_nodes(&mut self, f: impl FnOnce(&mut BTreeMap<NodeID, Cow<'static, Node<NData>>>)) {
        let mut nodes = (*self.nodes.read().unwrap()).clone();
        f(&mut nodes);
        self.nodes = Arc::new(RwLock::new(nodes));
    }

    fn _update_edges(&mut self, f: impl FnOnce(&mut BTreeMap<EdgeID, Cow<'static, Edge<EData>>>)) {
        let mut edges = (*self.edges.read().unwrap()).clone();
        f(&mut edges);
        self.edges = Arc::new(RwLock::new(edges));
    }

    fn _remove_edge(&mut self, edge: &mut Edge<EData>) {
        let source = edge.source.clone();
        let target = edge.target.clone();
        self._update_edges(|edges| {
            edges.remove(&edge.id);
        });
        self._update_nodes(|nodes| {
            if source == target {
                // Self-loop
                let node = nodes.get_mut(&source).unwrap()
                    .removing_out_edge(&edge.id)
                    .removing_in_edge(&edge.id);
                nodes.insert(source, Cow::Owned(node));
            } else {
                let source_node = nodes.get_mut(&source).unwrap().removing_out_edge(&edge.id);
                let target_node = nodes.get_mut(&target).unwrap().removing_in_edge(&edge.id);
                nodes.insert(source, Cow::Owned(source_node));
                nodes.insert(target, Cow::Owned(target_node));
            }
        });
    }

    fn _remove_node(&mut self, id: &NodeID) {
        self._update_nodes(|nodes| {
            nodes.remove(id);
        });
    }

    fn _clear_edges(&mut self, node: &mut Node<NData>) {
        for edge_id in self.incident_edges(&node.id).unwrap() {
            let mut edge = self._edge(&edge_id).unwrap().into_owned();
            self._remove_edge(&mut edge);
        }
    }
}

// Constructor where `GData` is specified directly or is not `Default`
impl<GData, NData, EData> Graph<GData, NData, EData>
where
    GData: Clone + 'static,
    NData: Clone + 'static,
    EData: Clone + 'static,
{
    pub fn new_with_data(data: GData) -> Self {
        Graph {
            nodes: Arc::new(RwLock::new(BTreeMap::new())),
            edges: Arc::new(RwLock::new(BTreeMap::new())),
            data: Cow::Owned(data),
        }
    }
}

// Constructor where `GData` implements `Default`
impl<GData, NData, EData> Graph<GData, NData, EData>
where
    GData: Clone + Default + 'static,
    NData: Clone + 'static,
    EData: Clone + 'static,
{
    pub fn new() -> Self {
        Graph {
            nodes: Arc::new(RwLock::new(BTreeMap::new())),
            edges: Arc::new(RwLock::new(BTreeMap::new())),
            data: Cow::Owned(GData::default()),
        }
    }
}

impl<GData, NData, EData> Default for Graph<GData, NData, EData>
where
    GData: Clone + Default + 'static,
    NData: Clone + 'static,
    EData: Clone + 'static,
{
    fn default() -> Self {
        Graph::new()
    }
}

impl<GData, NData, EData> MutableGraph for Graph<GData, NData, EData>
where
    GData: Clone + 'static,
    NData: Clone + 'static,
    EData: Clone + 'static,
{
    fn add_node_with_data(&mut self, id: impl AsRef<NodeID>, data: NData) -> Result<()> {
        let id = id.as_ref();
        if self.nodes.read().unwrap().contains_key(id) {
            bail!(Error::DuplicateNode(id.clone()))
        } else {
            let node = Node::new(id.clone(), data);
            self._update_nodes(|nodes| {
                nodes.insert(id.clone(), Cow::Owned(node));
            });
            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.edges.read().unwrap().contains_key(id) {
            bail!(Error::DuplicateEdge(id.clone()))
        } else if !self.nodes.read().unwrap().contains_key(source) {
            bail!(Error::NodeNotFound(source.clone()))
        } else if !self.nodes.read().unwrap().contains_key(target) {
            bail!(Error::NodeNotFound(target.clone()))
        } else {
            self._update_edges(|edges| {
                edges.insert(id.clone(), Cow::Owned(Edge::new(id.clone(), source.clone(), target.clone(), data)));
            });

            self._update_nodes(|nodes| {
                if source == target {
                    // Self-loop
                    let node = nodes.get_mut(source).unwrap()
                        .inserting_out_edge(id.clone())
                        .inserting_in_edge(id.clone());
                    nodes.insert(source.clone(), Cow::Owned(node));
                } else {
                    let source_node = nodes.get_mut(source).unwrap().inserting_out_edge(id.clone());
                    let target_node = nodes.get_mut(target).unwrap().inserting_in_edge(id.clone());
                    nodes.insert(source.clone(), Cow::Owned(source_node));
                    nodes.insert(target.clone(), Cow::Owned(target_node));
                }
            });

            Ok(())
        }
    }

    fn remove_edge(&mut self, id: impl AsRef<EdgeID>) -> Result<()> {
        let edge = self._edge(id.as_ref())?;
        self._remove_edge(&mut edge.into_owned());
        Ok(())
    }

    fn clear_edges(&mut self, id: impl AsRef<NodeID>) -> Result<()> {
        let mut node = self._node(id.as_ref())?.into_owned();
        self._clear_edges(&mut node);
        Ok(())
    }

    fn remove_node(&mut self, id: impl AsRef<NodeID>) -> Result<()> {
        let id = id.as_ref();
        let mut node = self._node(id)?.into_owned();
        self._clear_edges(&mut node);
        self._remove_node(id);
        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();
        self._node(new_source)?;
        self._node(new_target)?;
        let mut edge = self._edge(id.as_ref())?.into_owned();
        let old_source = edge.source.clone();
        let old_target = edge.target.clone();
        if &old_source != new_source || &old_target != new_target {
            edge.source = new_source.clone();
            edge.target = new_target.clone();
            self._update_edges(|edges| {
                edges.insert(id.clone(), Cow::Owned(edge));
            });
            self._update_nodes(|nodes| {
                if &old_source != new_source && &old_target != new_target {
                    let old_source_node = nodes.get_mut(&old_source).unwrap().removing_out_edge(id);
                    let old_target_node = nodes.get_mut(&old_target).unwrap().removing_in_edge(id);
                    nodes.insert(old_source.clone(), Cow::Owned(old_source_node));
                    nodes.insert(old_target.clone(), Cow::Owned(old_target_node));
                    let new_source_node = nodes.get_mut(new_source).unwrap().inserting_out_edge(id.clone());
                    let new_target_node = nodes.get_mut(new_target).unwrap().inserting_in_edge(id.clone());
                    nodes.insert(new_source.clone(), Cow::Owned(new_source_node));
                    nodes.insert(new_target.clone(), Cow::Owned(new_target_node));
                } else if &old_source != new_source {
                    let old_source_node = nodes.get_mut(&old_source).unwrap().removing_out_edge(id);
                    nodes.insert(old_source.clone(), Cow::Owned(old_source_node));
                    let new_source_node = nodes.get_mut(new_source).unwrap().inserting_out_edge(id.clone());
                    nodes.insert(new_source.clone(), Cow::Owned(new_source_node));
                } else if &old_target != new_target {
                    let old_target_node = nodes.get_mut(&old_target).unwrap().removing_in_edge(id);
                    nodes.insert(old_target.clone(), Cow::Owned(old_target_node));
                    let new_target_node = nodes.get_mut(new_target).unwrap().inserting_in_edge(id.clone());
                    nodes.insert(new_target.clone(), Cow::Owned(new_target_node));
                }
            });
        }
        Ok(())
    }

    fn set_data(&mut self, data: GData) {
        self.data = Cow::Owned(data);
    }

    fn set_node_data(&mut self, id: impl AsRef<NodeID>, data: NData) -> Result<()> {
        let id = id.as_ref();
        self._node(id)?;
        self._update_nodes(|nodes| {
            let node = nodes.get_mut(id).unwrap().setting_data(data);
            nodes.insert(id.clone(), Cow::Owned(node));
        });
        Ok(())
    }

    fn set_edge_data(&mut self, id: impl AsRef<EdgeID>, data: EData) -> Result<()> {
        let id = id.as_ref();
        self._edge(id)?;
        self._update_edges(|edges| {
            let edge = edges.get_mut(id).unwrap().setting_data(data);
            edges.insert(id.clone(), Cow::Owned(edge));
        });
        Ok(())
    }

    fn with_data(&mut self, transform: &dyn Fn(&mut Self::GData)) {
        let mut data = self.data.clone().into_owned();
        transform(&mut data);
        self.set_data(data);
    }

    fn with_node_data(&mut self, id: impl AsRef<NodeID>, transform: &dyn Fn(&mut Self::NData)) -> Result<()> {
        let id = id.as_ref();
        let mut data = self.node_data(id)?.into_owned();
        transform(&mut data);
        self.set_node_data(id, data)?;
        Ok(())
    }

    fn with_edge_data(&mut self, id: impl AsRef<EdgeID>, transform: &dyn Fn(&mut Self::EData)) -> Result<()> {
        let id = id.as_ref();
        let mut data = self.edge_data(id)?.into_owned();
        transform(&mut data);
        self.set_edge_data(id, data)?;
        Ok(())
    }
}

// Queries
impl<GData, NData, EData> VisitableGraph for Graph<GData, NData, EData>
where
    GData: Clone + 'static,
    NData: Clone + 'static,
    EData: Clone + 'static,
{
    type GData = GData;
    type NData = NData;
    type EData = EData;

    fn is_empty(&self) -> bool {
        self.nodes.read().unwrap().is_empty()
    }

    fn node_count(&self) -> usize {
        self.nodes.read().unwrap().len()
    }

    fn edge_count(&self) -> usize {
        self.edges.read().unwrap().len()
    }

    fn all_nodes(&self) -> Nodes {
        self.nodes.read().unwrap().keys().cloned().collect()
    }

    fn all_edges(&self) -> Edges {
        self.edges.read().unwrap().keys().cloned().collect()
    }

    fn has_node(&self, id: impl AsRef<NodeID>) -> bool {
        self.nodes.read().unwrap().contains_key(id.as_ref())
    }

    fn has_edge(&self, id: impl AsRef<EdgeID>) -> bool {
        self.edges.read().unwrap().contains_key(id.as_ref())
    }

    fn has_edge_from_to(&self, source: impl AsRef<NodeID>, target: impl AsRef<NodeID>) -> bool {
        self.edges.read().unwrap()
            .values().any(|edge| &edge.source == source.as_ref() && &edge.target == target.as_ref())
    }

    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.has_edge_from_to(a, b) || self.has_edge_from_to(b, a)
    }

    fn source(&self, id: impl AsRef<EdgeID>) -> Result<NodeID> {
        Ok(self._edge(id.as_ref())?.source.clone())
    }

    fn target(&self, id: impl AsRef<EdgeID>) -> Result<NodeID> {
        Ok(self._edge(id.as_ref())?.target.clone())
    }

    fn endpoints(&self, id: impl AsRef<EdgeID>) -> Result<(NodeID, NodeID)> {
        let edge = self._edge(id.as_ref())?;
        Ok((edge.source.clone(), edge.target.clone()))
    }

    fn out_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        Ok(self._node(id.as_ref())?.out_edges.read().unwrap().iter().cloned().collect())
    }

    fn in_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        Ok(self._node(id.as_ref())?.in_edges.read().unwrap().iter().cloned().collect())
    }

    fn incident_edges(&self, id: impl AsRef<NodeID>) -> Result<Edges> {
        let id = id.as_ref();
        let mut edges = self.out_edges(id)?;
        edges.extend(self.in_edges(id)?);
        Ok(edges)
    }

    fn out_degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        Ok(self._node(id.as_ref())?.out_edges.read().unwrap().len())
    }

    fn in_degree(&self, id: impl AsRef<NodeID>) -> Result<usize> {
        Ok(self._node(id.as_ref())?.in_edges.read().unwrap().len())
    }

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

    fn successors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        Ok(self.out_edges(id)?.iter()
            .map(|edge| self.target(edge).unwrap()).collect())
    }

    fn predecessors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        Ok(self.in_edges(id)?.iter()
            .map(|edge| self.source(edge).unwrap()).collect())
    }

    fn neighbors(&self, id: impl AsRef<NodeID>) -> Result<Nodes> {
        let id = id.as_ref();
        let mut neighbors = self.successors(id)?;
        neighbors.extend(self.predecessors(id)?);
        Ok(neighbors.into_iter().collect())
    }

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

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

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

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

    fn node_data(&self, id: impl AsRef<NodeID>) -> Result<Cow<'static, NData>> {
        Ok(self._node(id.as_ref())?.data.clone())
    }

    fn edge_data(&self, id: impl AsRef<EdgeID>) -> Result<Cow<'static, EData>> {
        Ok(self._edge(id.as_ref())?.data.clone())
    }

    fn all_roots(&self) -> Nodes {
        self.nodes.read().unwrap().iter()
            .filter(|(_, node)| node.in_edges.read().unwrap().is_empty())
            .map(|(id, _)| id).cloned().collect()
    }

    fn all_leaves(&self) -> Nodes {
        self.nodes.read().unwrap().iter()
            .filter(|(_, node)| node.out_edges.read().unwrap().is_empty())
            .map(|(id, _)| id).cloned().collect()
    }

    fn non_roots(&self) -> Nodes {
        self.nodes.read().unwrap().iter()
            .filter(|(_, node)| !node.in_edges.read().unwrap().is_empty())
            .map(|(id, _)| id).cloned().collect()
    }

    fn non_leaves(&self) -> Nodes {
        self.nodes.read().unwrap().iter()
            .filter(|(_, node)| !node.out_edges.read().unwrap().is_empty())
            .map(|(id, _)| id).cloned().collect()
    }

    fn all_internals(&self) -> Nodes {
        self.nodes.read().unwrap().iter()
            .filter(|(_, node)| {
                !node.in_edges.read().unwrap().is_empty() && !node.out_edges.read().unwrap().is_empty()
            })
            .map(|(id, _)| id).cloned().collect()
    }

    fn is_leaf(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        Ok(self._node(id.as_ref())?.out_edges.read().unwrap().is_empty())
    }

    fn is_root(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        Ok(self._node(id.as_ref())?.in_edges.read().unwrap().is_empty())
    }

    fn is_internal(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        let node = self._node(id.as_ref())?;
        Ok(!node.in_edges.read().unwrap().is_empty() && !node.out_edges.read().unwrap().is_empty())
    }
}

impl<GData, NData, EData> PartialEq for Graph<GData, NData, EData>
where
    GData: Clone + PartialEq + 'static,
    NData: Clone + PartialEq + 'static,
    EData: Clone + PartialEq + 'static,
{
    fn eq(&self, other: &Self) -> bool {
        self.nodes.read().unwrap().iter().eq(other.nodes.read().unwrap().iter())
            && self.edges.read().unwrap().iter().eq(other.edges.read().unwrap().iter())
            && self.data == other.data
    }
}

#[cfg(feature = "serde")]
impl<GData, NData, EData> Serialize for Graph<GData, NData, EData>
where
    GData: Clone + Serialize + 'static,
    NData: Clone + Serialize + 'static,
    EData: Clone + Serialize + 'static,
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let nodes: Vec<Node<NData>> = self.nodes.read().unwrap()
            .values()
            .map(|n| n.clone().into_owned())
            .collect();

        let edges: Vec<Edge<EData>> = self.edges.read().unwrap()
            .values()
            .map(|e| e.clone().into_owned())
            .collect();

        if std::mem::size_of::<GData>() > 0 {
            (&nodes, &edges, &self.data).serialize(serializer)
        } else {
            (&nodes, &edges).serialize(serializer)
        }
    }
}

#[cfg(feature = "serde")]
impl<'de, GData, NData, EData> Deserialize<'de> for Graph<GData, NData, EData>
where
    GData: Clone + Deserialize<'de> + 'static + Default,
    NData: Clone + Deserialize<'de> + 'static + Default,
    EData: Clone + Deserialize<'de> + 'static + Default,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let mut graph: Graph<GData, NData, EData>;
        let nodes: Vec<Node<NData>>;
        let edges: Vec<Edge<EData>>;

        if std::mem::size_of::<GData>() > 0 {
            let (n, e, data): (Vec<Node<NData>>, Vec<Edge<EData>>, GData) = Deserialize::deserialize(deserializer)?;
            graph = Graph::new_with_data(data);
            nodes = n;
            edges = e;
        } else {
            let (n, e): (Vec<Node<NData>>, Vec<Edge<EData>>) = Deserialize::deserialize(deserializer)?;
            graph = Graph::new();
            nodes = n;
            edges = e;
        }

        for node in nodes {
            let node_data: NData = node.data.into_owned();
            graph.add_node_with_data(&node.id, node_data).map_err(de::Error::custom)?;
        }

        for edge in edges {
            let edge_data: EData = edge.data.into_owned();
            graph.add_edge_with_data(&edge.id, &edge.source, &edge.target, edge_data).map_err(de::Error::custom)?;
        }

        Ok(graph)
    }
}

// If Serde and SerdeJSON are both present, add convenience to serialize a Graph
// to JSON.
#[cfg(all(feature = "serde", feature = "serde_json"))]
impl<GData, NData, EData> Graph<GData, NData, EData>
where
    GData: Clone + Serialize + 'static,
    NData: Clone + Serialize + 'static,
    EData: Clone + Serialize + 'static,
{
    pub fn to_json(&self) -> String {
        serde_json::to_string(self).unwrap()
    }
}

// If Serde and SerdeJSON are both present, add convenience to deserialize a
// Graph from JSON.
#[cfg(all(feature = "serde", feature = "serde_json"))]
impl<'de, GData, NData, EData> Graph<GData, NData, EData>
where
    GData: Clone + Default + Deserialize<'de> + 'static,
    NData: Clone + Default + Deserialize<'de> + 'static,
    EData: Clone + Default + Deserialize<'de> + 'static,
{
    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<GData, NData, EData> std::fmt::Display for Graph<GData, NData, EData>
where
    GData: Clone + Serialize + 'static,
    NData: Clone + Serialize + 'static,
    EData: Clone + Serialize + 'static,
{
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.to_json())
    }
}