wolf-derivation-graph 0.1.0

use std::{cell::RefCell, collections::HashMap, rc::Rc};
use anyhow::{anyhow, Result};

use wolf_graph::prelude::*;

struct NodeData<T> {
    value: Option<T>,
    label: String,
}

impl<T> NodeData<T> {
    fn new() -> Self {
        Self {
            value: None,
            label: String::new(),
        }
    }

    fn value(&self) -> Option<&T> {
        self.value.as_ref()
    }

    fn set_value(&mut self, value: T) {
        self.value = Some(value);
    }

    fn clear_value(&mut self) {
        self.value = None;
    }

    fn label(&self) -> &str {
        &self.label
    }

    fn set_label(&mut self, label: impl Into<String>) {
        self.label = label.into();
    }
}

type NodeDataRef<T> = Rc<RefCell<NodeData<T>>>;
type G<T> = Graph<(), NodeDataRef<T>, String>;
type D<T> = DAG::<G<T>>;
pub type OpInputs<T> = HashMap<String, T>;
pub type Operation<T> = Rc<dyn Fn(OpInputs<T>) -> Result<T>>;

pub struct Deriver<T: Clone + 'static> {
    node: NodeID,
    op: Operation<T>,
}

impl<T: Eq + Clone + 'static> Deriver<T> {
    pub fn new(node: impl AsRef<NodeID>, op: Operation<T>) -> Self {
        Self {
            node: node.as_ref().clone(),
            op,
        }
    }

    pub fn node(&self) -> &NodeID {
        &self.node
    }

    pub fn op(&self) -> &Operation<T> {
        &self.op
    }
}

pub struct DerivationGraph<T: Eq + Clone + 'static>
{
    graph: D<T>,
    derivations: HashMap<NodeID, Deriver<T>>,
}

impl<T: Eq + Clone + 'static> DerivationGraph<T> {
    pub fn new() -> Self {
        Self {
            graph: D::new(),
            derivations: HashMap::new(),
        }
    }

    pub fn set_node_label(&mut self, id: impl AsRef<NodeID>, label: impl Into<String>) -> Result<()> {
        self.graph.node_data(id)?.borrow_mut().set_label(label);
        Ok(())
    }

    pub fn node_label(&self, id: impl AsRef<NodeID>) -> Result<String> {
        Ok(self.graph.node_data(id)?.borrow().label().to_string())
    }

    pub fn set_edge_label(&mut self, id: impl AsRef<EdgeID>, label: impl Into<String>) -> Result<()> {
        self.graph.set_edge_data(id, label.into())
    }

    pub fn edge_label(&self, id: impl AsRef<EdgeID>) -> Result<String> {
        Ok(self.graph.edge_data(id)?.to_string())
    }

    pub fn add_node(
        &mut self,
        id: impl AsRef<NodeID>,
        label: impl Into<String>,
    ) -> Result<()> {
        self.graph.add_node_with_data(&id, Rc::new(RefCell::new(NodeData::new())))
            .map(|_| self.set_node_label(&id, label))
            .map(|_| ())
    }

    pub fn add_node_with_value(
        &mut self,
        id: impl AsRef<NodeID>,
        label: impl Into<String>,
        value: impl Into<T>,
    ) -> Result<()> {
        let id = id.as_ref();
        self.add_node(id, label)?;
        self.set_node_value(id, value)
    }

    pub fn add_node_with_operation(
        &mut self,
        id: impl AsRef<NodeID>,
        label: impl Into<String>,
        op: Operation<T>,
    ) -> Result<()> {
        let id = id.as_ref();
        self.add_node(id, label)?;
        self.set_operation(id, op)
    }

    pub fn set_operation(
        &mut self,
        node: impl AsRef<NodeID>,
        op: Operation<T>,
    ) -> Result<()> {
        let node = node.as_ref();
        let deriver = Deriver::new(node, op);
        self.derivations.insert(node.clone(), deriver);
        self.clear_node_values_transitively(node)?;
        Ok(())
    }

    pub fn add_edge(
        &mut self,
        id: impl AsRef<EdgeID>,
        source: impl AsRef<NodeID>,
        target: impl AsRef<NodeID>,
        label: impl Into<String>,
    ) -> Result<()> {
        self.graph.add_edge_with_data(id, source, target, label.into())
    }

    pub fn node_value(&self, id: impl AsRef<NodeID>) -> Result<Option<T>> {
        Ok(self.graph.node_data(id)?.borrow().value().cloned())
    }

    pub fn has_node_value(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        Ok(self.node_value(id)?.is_some())
    }

    pub fn set_node_value(&self, id: impl AsRef<NodeID>, value: impl Into<T>) -> Result<()> {
        let id = id.as_ref();
        let value = value.into();

        // Do nothing if the node already has the same data.
        if let Some(existing_value) = self.node_value(id)? {
            if existing_value == value {
                return Ok(());
            }
        }

        // Clear all the transitive downstream nodes of the node.
        self.clear_node_values_transitively(id)?;

        // Set the data of the node.
        let node_data = self.graph.node_data(id)?;
        node_data.borrow_mut().set_value(value);

        Ok(())
    }

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

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

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

        // If the node already has data, return it.
        if let Some(data) = self.node_value(id)? {
            return Ok(data.clone());
        }

        // If the node has no deriver, return an error.
        let deriver = self.derivations.get(id).ok_or_else(|| anyhow!("no deriver for node '{id}'"))?;

        // Accumulate the derived data from the source of each in_edge into `inputs`.
        let mut inputs = HashMap::new();
        let in_edges = self.graph.in_edges(id)?;
        for in_edge in in_edges {
            let source = self.source(&in_edge)?;
            let data = self.derived_node_value(source)?;
            let label = self.graph.edge_data(&in_edge)?.into_owned();
            inputs.insert(label, data);
        }

        let op = deriver.op();
        let data = op(inputs)?;
        self.set_node_value(id, data.clone())?;
        Ok(data)
    }

    pub fn clear_all_node_values(&self) {
        for node in self.graph.all_nodes() {
            self.clear_node_value(&node).unwrap();
        }
    }
}

// Private methods.
impl<T: Eq + Clone + 'static> DerivationGraph<T> {
    // Clears all downstream nodes of the given node transitively, stopping at
    // nodes that are already clear.
    fn clear_node_values_transitively(&self, id: impl AsRef<NodeID>) -> Result<()> {
        let mut stack = vec![id.as_ref().clone()];
        while let Some(id) = stack.pop() {
            if self.clear_node_value(&id)? {
                for successor in self.graph.successors(&id)? {
                    if self.has_node_value(&successor)? {
                        stack.push(successor);
                    }
                }
            }
        }
        Ok(())
    }

    // Returns true if the node had data to clear.
    fn clear_node_value(&self, id: impl AsRef<NodeID>) -> Result<bool> {
        let id = id.as_ref();
        let value = self.node_value(id)?;
        let had_value = value.is_some();
        if had_value {
            self.graph.node_data(id)?.borrow_mut().clear_value();
        }
        Ok(had_value)
    }
}

impl<T: Eq + Clone + 'static> Default for DerivationGraph<T> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
}