use crate::compiled::CompiledGraph;
use pe_core::state::State;
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
pub struct GraphRegistry {
graphs: HashMap<String, Arc<dyn Any + Send + Sync>>,
}
impl GraphRegistry {
pub fn new() -> Self {
Self {
graphs: HashMap::new(),
}
}
pub fn register<S: State + 'static>(
&mut self,
name: impl Into<String>,
graph: CompiledGraph<S>,
) {
self.graphs.insert(name.into(), Arc::new(graph));
}
pub fn get<S: State + 'static>(&self, name: &str) -> Option<&CompiledGraph<S>> {
self.graphs
.get(name)
.and_then(|arc| arc.downcast_ref::<CompiledGraph<S>>())
}
pub fn list(&self) -> Vec<&str> {
self.graphs.keys().map(|s| s.as_str()).collect()
}
pub fn contains(&self, name: &str) -> bool {
self.graphs.contains_key(name)
}
pub fn len(&self) -> usize {
self.graphs.len()
}
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"));
}
}