pe_graph/
registry.rs

1//! Graph registry — named storage for compiled graphs.
2//!
3//! Type-erased storage allows registering graphs with different state types
4//! under string names, then retrieving them with type-safe downcasting.
5
6use crate::compiled::CompiledGraph;
7use pe_core::state::State;
8use std::any::Any;
9use std::collections::HashMap;
10use std::sync::Arc;
11
12/// Named registry for compiled graphs.
13///
14/// Stores graphs as type-erased `Arc<dyn Any>` values. Retrieval requires
15/// specifying the correct state type `S` — a type mismatch returns `None`.
16///
17/// # Example
18///
19/// ```ignore
20/// let mut registry = GraphRegistry::new();
21/// registry.register("react-agent", compiled_graph);
22///
23/// let graph: Option<&CompiledGraph<MyState>> = registry.get("react-agent");
24/// ```
25pub struct GraphRegistry {
26    graphs: HashMap<String, Arc<dyn Any + Send + Sync>>,
27}
28
29impl GraphRegistry {
30    /// Create an empty registry.
31    pub fn new() -> Self {
32        Self {
33            graphs: HashMap::new(),
34        }
35    }
36
37    /// Register a compiled graph under a name.
38    ///
39    /// Overwrites any existing graph with the same name.
40    pub fn register<S: State + 'static>(
41        &mut self,
42        name: impl Into<String>,
43        graph: CompiledGraph<S>,
44    ) {
45        self.graphs.insert(name.into(), Arc::new(graph));
46    }
47
48    /// Retrieve a graph by name, downcasting to the expected state type.
49    ///
50    /// Returns `None` if the name doesn't exist or the type doesn't match.
51    pub fn get<S: State + 'static>(&self, name: &str) -> Option<&CompiledGraph<S>> {
52        self.graphs
53            .get(name)
54            .and_then(|arc| arc.downcast_ref::<CompiledGraph<S>>())
55    }
56
57    /// List all registered graph names.
58    pub fn list(&self) -> Vec<&str> {
59        self.graphs.keys().map(|s| s.as_str()).collect()
60    }
61
62    /// Check if a graph is registered under the given name.
63    pub fn contains(&self, name: &str) -> bool {
64        self.graphs.contains_key(name)
65    }
66
67    /// Number of registered graphs.
68    pub fn len(&self) -> usize {
69        self.graphs.len()
70    }
71
72    /// Whether the registry is empty.
73    pub fn is_empty(&self) -> bool {
74        self.graphs.is_empty()
75    }
76}
77
78impl Default for GraphRegistry {
79    fn default() -> Self {
80        Self::new()
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87    use crate::graph::StateGraph;
88    use crate::tests::{AppendNode, TestState};
89    use pe_core::types::{END, START};
90
91    fn make_test_graph() -> CompiledGraph<TestState> {
92        StateGraph::<TestState>::new()
93            .add_node("a", AppendNode::new("a", "hello"))
94            .add_edge(START, "a")
95            .add_edge("a", END)
96            .compile()
97            .unwrap()
98    }
99
100    #[test]
101    fn test_register_and_get() {
102        let mut reg = GraphRegistry::new();
103        reg.register("test", make_test_graph());
104
105        let got: Option<&CompiledGraph<TestState>> = reg.get("test");
106        assert!(got.is_some());
107    }
108
109    #[test]
110    fn test_get_nonexistent_returns_none() {
111        let reg = GraphRegistry::new();
112        let got: Option<&CompiledGraph<TestState>> = reg.get("nope");
113        assert!(got.is_none());
114    }
115
116    #[test]
117    fn test_list_and_contains() {
118        let mut reg = GraphRegistry::new();
119        reg.register("alpha", make_test_graph());
120        reg.register("beta", make_test_graph());
121
122        assert_eq!(reg.len(), 2);
123        assert!(reg.contains("alpha"));
124        assert!(reg.contains("beta"));
125        assert!(!reg.contains("gamma"));
126
127        let names = reg.list();
128        assert!(names.contains(&"alpha"));
129        assert!(names.contains(&"beta"));
130    }
131}
pe_graph/registry.rs

pe_graph/
registry.rs
