pe-graph 0.1.0

Graph execution engine for Potential Expectations — state graphs, Pregel model, ReAct topology, and builder DSL
Documentation
//! Graph registry — named storage for compiled graphs.
//!
//! Type-erased storage allows registering graphs with different state types
//! under string names, then retrieving them with type-safe downcasting.

use crate::compiled::CompiledGraph;
use pe_core::state::State;
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;

/// Named registry for compiled graphs.
///
/// Stores graphs as type-erased `Arc<dyn Any>` values. Retrieval requires
/// specifying the correct state type `S` — a type mismatch returns `None`.
///
/// # Example
///
/// ```ignore
/// let mut registry = GraphRegistry::new();
/// registry.register("react-agent", compiled_graph);
///
/// let graph: Option<&CompiledGraph<MyState>> = registry.get("react-agent");
/// ```
pub struct GraphRegistry {
    graphs: HashMap<String, Arc<dyn Any + Send + Sync>>,
}

impl GraphRegistry {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self {
            graphs: HashMap::new(),
        }
    }

    /// Register a compiled graph under a name.
    ///
    /// Overwrites any existing graph with the same name.
    pub fn register<S: State + 'static>(
        &mut self,
        name: impl Into<String>,
        graph: CompiledGraph<S>,
    ) {
        self.graphs.insert(name.into(), Arc::new(graph));
    }

    /// Retrieve a graph by name, downcasting to the expected state type.
    ///
    /// Returns `None` if the name doesn't exist or the type doesn't match.
    pub fn get<S: State + 'static>(&self, name: &str) -> Option<&CompiledGraph<S>> {
        self.graphs
            .get(name)
            .and_then(|arc| arc.downcast_ref::<CompiledGraph<S>>())
    }

    /// List all registered graph names.
    pub fn list(&self) -> Vec<&str> {
        self.graphs.keys().map(|s| s.as_str()).collect()
    }

    /// Check if a graph is registered under the given name.
    pub fn contains(&self, name: &str) -> bool {
        self.graphs.contains_key(name)
    }

    /// Number of registered graphs.
    pub fn len(&self) -> usize {
        self.graphs.len()
    }

    /// Whether the registry is empty.
    pub fn is_empty(&self) -> bool {
        self.graphs.is_empty()
    }
}

impl Default for GraphRegistry {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::StateGraph;
    use crate::tests::{AppendNode, TestState};
    use pe_core::types::{END, START};

    fn make_test_graph() -> CompiledGraph<TestState> {
        StateGraph::<TestState>::new()
            .add_node("a", AppendNode::new("a", "hello"))
            .add_edge(START, "a")
            .add_edge("a", END)
            .compile()
            .unwrap()
    }

    #[test]
    fn test_register_and_get() {
        let mut reg = GraphRegistry::new();
        reg.register("test", make_test_graph());

        let got: Option<&CompiledGraph<TestState>> = reg.get("test");
        assert!(got.is_some());
    }

    #[test]
    fn test_get_nonexistent_returns_none() {
        let reg = GraphRegistry::new();
        let got: Option<&CompiledGraph<TestState>> = reg.get("nope");
        assert!(got.is_none());
    }

    #[test]
    fn test_list_and_contains() {
        let mut reg = GraphRegistry::new();
        reg.register("alpha", make_test_graph());
        reg.register("beta", make_test_graph());

        assert_eq!(reg.len(), 2);
        assert!(reg.contains("alpha"));
        assert!(reg.contains("beta"));
        assert!(!reg.contains("gamma"));

        let names = reg.list();
        assert!(names.contains(&"alpha"));
        assert!(names.contains(&"beta"));
    }
}