pe_graph/

graph.rs

//! Graph definition — `StateGraph<S>` builder and validation.
//!
//! A `StateGraph` defines the topology: nodes (async functions) and edges
//! (fixed or conditional transitions). Call `compile()` to validate the
//! structure and produce a `CompiledGraph` ready for execution.

use crate::compiled::CompiledGraph;
use pe_core::error::PeError;
use pe_core::node::NodeFn;
use pe_core::state::State;
use pe_core::types::{END, START};
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::Arc;

/// Fixed edge: always transitions from one node to another.
#[derive(Debug, Clone)]
pub(crate) struct Edge {
    pub from: String,
    pub to: String,
}

/// Type alias for a routing function that reads state and returns next node names.
pub(crate) type RouterFn<S> = Arc<dyn Fn(&S) -> Vec<String> + Send + Sync>;

/// Conditional edge: router function reads state and returns target node names.
pub(crate) struct ConditionalEdge<S: State> {
    pub from: String,
    pub router: RouterFn<S>,
}

/// Declarative graph definition with typed nodes and edges.
///
/// Build a graph using the fluent API, then call `compile()` to validate
/// and freeze it into a `CompiledGraph` for execution.
///
/// # Example
///
/// ```ignore
/// use pe_graph::{StateGraph, START, END};
///
/// let graph = StateGraph::new()
///     .add_node("prepare", prepare_node)
///     .add_node("chat", chat_node)
///     .add_edge(START, "prepare")
///     .add_edge("prepare", "chat")
///     .add_edge("chat", END)
///     .compile()?;
/// ```
pub struct StateGraph<S: State> {
    pub(crate) nodes: HashMap<String, Arc<dyn NodeFn<S>>>,
    pub(crate) edges: Vec<Edge>,
    pub(crate) conditional_edges: Vec<ConditionalEdge<S>>,
}

impl<S: State> StateGraph<S> {
    /// Create an empty graph definition.
    pub fn new() -> Self {
        Self {
            nodes: HashMap::new(),
            edges: Vec::new(),
            conditional_edges: Vec::new(),
        }
    }

    /// Add a named node to the graph.
    ///
    /// Reserved names (`START`, `END`) are rejected at `compile()` time,
    /// not here — so the builder stays infallible and chainable.
    pub fn add_node(mut self, name: impl Into<String>, node: impl NodeFn<S> + 'static) -> Self {
        let name = name.into();
        self.nodes.insert(name, Arc::new(node));
        self
    }

    /// Add a fixed edge: `from` always transitions to `to` after completing.
    ///
    /// Use `START` for the first node, `END` for terminal nodes.
    pub fn add_edge(mut self, from: impl Into<String>, to: impl Into<String>) -> Self {
        self.edges.push(Edge {
            from: from.into(),
            to: to.into(),
        });
        self
    }

    /// Add a conditional edge: router reads state and returns next node names.
    ///
    /// Return `vec![END.to_string()]` to terminate. Return multiple names
    /// for parallel execution in the next superstep.
    pub fn add_conditional_edge(
        mut self,
        from: impl Into<String>,
        router: impl Fn(&S) -> Vec<String> + Send + Sync + 'static,
    ) -> Self {
        self.conditional_edges.push(ConditionalEdge {
            from: from.into(),
            router: Arc::new(router),
        });
        self
    }

    /// Get the fixed-edge successor node names for a given node.
    ///
    /// Returns successors reachable via fixed edges (not conditional).
    /// Used by `resume()` and `Command::Update` to skip re-running the
    /// interrupted node and continue to its successors.
    ///
    /// Includes END as a successor — the Pregel engine handles END
    /// gracefully (no node found → empty results → graph completes).
    /// Without this, a node whose only edge is `node → END` would
    /// return empty, causing the caller to re-run the interrupted node.
    pub(crate) fn fixed_successors(&self, node_name: &str) -> Vec<String> {
        self.edges
            .iter()
            .filter(|e| e.from == node_name)
            .map(|e| e.to.clone())
            .collect()
    }

    /// Validate the graph structure and freeze it into a `CompiledGraph`.
    ///
    /// Validation checks:
    /// - START has at least one outgoing edge
    /// - All fixed edge targets/sources exist as node names (or START/END)
    /// - No conditional edges from START
    /// - All nodes are reachable from START
    #[must_use = "CompiledGraph is the executable form — don't discard it"]
    pub fn compile(self) -> Result<CompiledGraph<S>, PeError> {
        self.validate()?;
        Ok(CompiledGraph::new(Arc::new(self)))
    }

    /// Validate graph structure. Returns error on first violation found.
    fn validate(&self) -> Result<(), PeError> {
        let node_names: HashSet<&str> = self.nodes.keys().map(|s| s.as_str()).collect();

        // No nodes with reserved names
        if node_names.contains(START) || node_names.contains(END) {
            return Err(PeError::GraphValue {
                details: format!(
                    "Cannot use reserved names '{}' or '{}' as node names",
                    START, END
                ),
            });
        }

        // START must have at least one outgoing fixed edge
        let has_start_edge = self.edges.iter().any(|e| e.from == START);
        if !has_start_edge {
            return Err(PeError::GraphValue {
                details: "START has no outgoing edges — add at least one edge from START".into(),
            });
        }

        // No conditional edges from START (no state to route on yet)
        if self.conditional_edges.iter().any(|ce| ce.from == START) {
            return Err(PeError::GraphValue {
                details: "Conditional edges from START are not allowed — use a router node instead"
                    .into(),
            });
        }

        // All fixed edge targets must be known nodes or END
        // All fixed edge sources must be known nodes or START
        for edge in &self.edges {
            if edge.from != START && !node_names.contains(edge.from.as_str()) {
                return Err(PeError::GraphValue {
                    details: format!("Edge source '{}' is not a known node", edge.from),
                });
            }
            if edge.to != END && !node_names.contains(edge.to.as_str()) {
                return Err(PeError::GraphValue {
                    details: format!("Edge target '{}' is not a known node", edge.to),
                });
            }
        }

        // Conditional edge sources must be known nodes
        for ce in &self.conditional_edges {
            if !node_names.contains(ce.from.as_str()) {
                return Err(PeError::GraphValue {
                    details: format!("Conditional edge source '{}' is not a known node", ce.from),
                });
            }
        }

        // All nodes must be reachable from START via BFS
        // Conditional edges can reach any node, so treat them as wildcards
        let reachable = self.reachable_from_start();
        for name in node_names {
            if !reachable.contains(name) {
                return Err(PeError::UnreachableNode {
                    node: name.to_string(),
                });
            }
        }

        Ok(())
    }

    /// BFS from START. Fixed edges are followed directly.
    /// Conditional edges from a reachable node make ALL nodes reachable
    /// (since we can't evaluate the router at compile time).
    fn reachable_from_start(&self) -> HashSet<String> {
        let mut reachable = HashSet::new();
        let mut queue = VecDeque::new();

        // Seed with START's fixed edge targets
        for edge in &self.edges {
            if edge.from == START && edge.to != END {
                queue.push_back(edge.to.clone());
            }
        }

        // If START has a conditional edge, all nodes are reachable
        // (already rejected in validate, but defensive)
        let cond_sources: HashSet<&str> = self
            .conditional_edges
            .iter()
            .map(|ce| ce.from.as_str())
            .collect();

        while let Some(node) = queue.pop_front() {
            if !reachable.insert(node.clone()) {
                continue; // already visited
            }

            // If this node has a conditional edge, all nodes are potentially reachable
            if cond_sources.contains(node.as_str()) {
                for name in self.nodes.keys() {
                    queue.push_back(name.clone());
                }
            }

            // Follow fixed edges from this node
            for edge in &self.edges {
                if edge.from == node && edge.to != END {
                    queue.push_back(edge.to.clone());
                }
            }
        }

        reachable
    }
}

impl<S: State> Default for StateGraph<S> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tests::{AppendNode, TestState};

    #[test]
    fn test_valid_linear_graph_compiles() {
        let graph = StateGraph::<TestState>::new()
            .add_node("a", AppendNode::new("a", "hello"))
            .add_node("b", AppendNode::new("b", "world"))
            .add_edge(START, "a")
            .add_edge("a", "b")
            .add_edge("b", END)
            .compile();

        assert!(graph.is_ok());
    }

    #[test]
    fn test_no_start_edge_rejected() {
        let result = StateGraph::<TestState>::new()
            .add_node("a", AppendNode::new("a", "hello"))
            .add_edge("a", END)
            .compile();

        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("START"), "Error should mention START: {err}");
    }

    #[test]
    fn test_unknown_edge_target_rejected() {
        let result = StateGraph::<TestState>::new()
            .add_node("a", AppendNode::new("a", "hello"))
            .add_edge(START, "a")
            .add_edge("a", "nonexistent")
            .compile();

        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("nonexistent"),
            "Error should mention missing node: {err}"
        );
    }

    #[test]
    fn test_unreachable_node_rejected() {
        let result = StateGraph::<TestState>::new()
            .add_node("a", AppendNode::new("a", "hello"))
            .add_node("orphan", AppendNode::new("orphan", "lost"))
            .add_edge(START, "a")
            .add_edge("a", END)
            .compile();

        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("orphan"),
            "Error should mention unreachable node: {err}"
        );
    }

    #[test]
    fn test_conditional_from_start_rejected() {
        let result = StateGraph::<TestState>::new()
            .add_node("a", AppendNode::new("a", "hello"))
            .add_edge(START, "a")
            .add_conditional_edge(START, |_: &TestState| vec!["a".into()])
            .add_edge("a", END)
            .compile();

        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("Conditional edges from START"),
            "Error should explain restriction: {err}"
        );
    }

    #[test]
    fn test_conditional_edge_makes_nodes_reachable() {
        // "b" is only reachable through conditional edge from "a"
        let graph = StateGraph::<TestState>::new()
            .add_node("a", AppendNode::new("a", "hello"))
            .add_node("b", AppendNode::new("b", "world"))
            .add_edge(START, "a")
            .add_conditional_edge("a", |_: &TestState| vec![END.into()])
            .add_edge("b", END)
            .compile();

        assert!(
            graph.is_ok(),
            "Nodes reachable via conditional edges should be valid"
        );
    }

    #[test]
    fn test_reserved_name_rejected_at_compile() {
        let err = StateGraph::<TestState>::new()
            .add_node(START, AppendNode::new(START, "bad"))
            .add_edge(START, START)
            .compile()
            .unwrap_err()
            .to_string();

        assert!(
            err.contains("reserved name"),
            "Error should mention reserved name: {err}"
        );
    }
}