use std::collections::{HashMap, HashSet};
use crate::engine::manager::ECSReference;
use crate::{AgentTemplateId, Entity};
use super::error::{AgentError, AgentResult};
use super::template::AgentTemplate;
pub struct AgentRegistry {
templates: HashMap<String, AgentTemplate>,
ids: HashMap<String, AgentTemplateId>,
names_by_id: Vec<String>,
entities_by_template: HashMap<AgentTemplateId, Vec<Entity>>,
template_by_entity: HashMap<Entity, AgentTemplateId>,
sealed: bool,
pending_spawn_hooks: Vec<(String, Entity)>,
pending_despawn_hooks: Vec<(String, Entity)>,
pending_spawn_batch_hooks: Vec<(AgentTemplateId, Vec<Entity>)>,
pending_despawn_batch_hooks: Vec<(AgentTemplateId, Vec<Entity>)>,
}
impl AgentRegistry {
pub fn new() -> Self {
Self {
templates: HashMap::new(),
ids: HashMap::new(),
names_by_id: Vec::new(),
entities_by_template: HashMap::new(),
template_by_entity: HashMap::new(),
sealed: false,
pending_spawn_hooks: Vec::new(),
pending_despawn_hooks: Vec::new(),
pending_spawn_batch_hooks: Vec::new(),
pending_despawn_batch_hooks: Vec::new(),
}
}
pub fn register(&mut self, mut template: AgentTemplate) -> AgentResult<()> {
if self.sealed {
return Err(AgentError::RegistrySealed);
}
if self.templates.contains_key(&template.name) {
return Err(AgentError::DuplicateTemplate(template.name.clone()));
}
let id = AgentTemplateId(self.names_by_id.len() as u32);
template.id = Some(id);
self.ids.insert(template.name.clone(), id);
self.names_by_id.push(template.name.clone());
self.entities_by_template.entry(id).or_default();
self.templates.insert(template.name.clone(), template);
Ok(())
}
pub fn get(&self, name: &str) -> AgentResult<&AgentTemplate> {
self.templates
.get(name)
.ok_or_else(|| AgentError::TemplateNotFound(name.to_owned()))
}
pub fn id(&self, name: &str) -> AgentResult<AgentTemplateId> {
self.ids
.get(name)
.copied()
.ok_or_else(|| AgentError::TemplateNotFound(name.to_owned()))
}
pub fn get_by_id(&self, id: AgentTemplateId) -> AgentResult<&AgentTemplate> {
let name = self
.names_by_id
.get(id.0 as usize)
.ok_or_else(|| AgentError::TemplateNotFound(format!("#{}", id.0)))?;
self.get(name)
}
pub fn entities(&self, id: AgentTemplateId) -> &[Entity] {
self.entities_by_template
.get(&id)
.map(Vec::as_slice)
.unwrap_or(&[])
}
pub fn entity_template(&self, entity: Entity) -> Option<AgentTemplateId> {
self.template_by_entity.get(&entity).copied()
}
pub fn seal(&mut self) {
self.sealed = true;
}
#[inline]
pub fn is_sealed(&self) -> bool {
self.sealed
}
#[inline]
pub fn len(&self) -> usize {
self.templates.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.templates.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &AgentTemplate)> {
self.templates.iter().map(|(k, v)| (k.as_str(), v))
}
pub fn enqueue_spawn_hook(&mut self, template_name: impl Into<String>, entity: Entity) {
let template_name = template_name.into();
if let Some(id) = self.ids.get(&template_name).copied() {
self.template_by_entity.insert(entity, id);
self.entities_by_template
.entry(id)
.or_default()
.push(entity);
}
self.pending_spawn_hooks.push((template_name, entity));
}
pub fn enqueue_spawn_batch_hook(
&mut self,
template_id: AgentTemplateId,
entities: Vec<Entity>,
) -> AgentResult<()> {
self.get_by_id(template_id)?;
for entity in &entities {
self.template_by_entity.insert(*entity, template_id);
}
self.entities_by_template
.entry(template_id)
.or_default()
.extend(entities.iter().copied());
self.pending_spawn_batch_hooks.push((template_id, entities));
Ok(())
}
pub fn enqueue_spawn_hooks_by_id(
&mut self,
template_id: AgentTemplateId,
entities: &[Entity],
) -> AgentResult<()> {
let name = self
.names_by_id
.get(template_id.0 as usize)
.ok_or_else(|| AgentError::TemplateNotFound(format!("#{}", template_id.0)))?;
self.pending_spawn_hooks.extend(
entities
.iter()
.copied()
.map(|entity| (name.clone(), entity)),
);
Ok(())
}
pub fn enqueue_despawn_hook(&mut self, template_name: impl Into<String>, entity: Entity) {
let template_name = template_name.into();
if let Some(id) = self.ids.get(&template_name).copied() {
self.template_by_entity.remove(&entity);
if let Some(live) = self.entities_by_template.get_mut(&id) {
live.retain(|candidate| *candidate != entity);
}
}
self.pending_despawn_hooks.push((template_name, entity));
}
pub fn enqueue_despawn_batch_hook(
&mut self,
template_id: AgentTemplateId,
entities: Vec<Entity>,
) -> AgentResult<()> {
self.get_by_id(template_id)?;
for entity in &entities {
self.template_by_entity.remove(entity);
}
if let Some(live) = self.entities_by_template.get_mut(&template_id) {
let despawned: HashSet<Entity> = entities.iter().copied().collect();
live.retain(|entity| !despawned.contains(entity));
}
self.pending_despawn_batch_hooks
.push((template_id, entities));
Ok(())
}
pub fn enqueue_despawn_hooks_by_id(
&mut self,
template_id: AgentTemplateId,
entities: &[Entity],
) -> AgentResult<()> {
let name = self
.names_by_id
.get(template_id.0 as usize)
.ok_or_else(|| AgentError::TemplateNotFound(format!("#{}", template_id.0)))?;
self.pending_despawn_hooks.extend(
entities
.iter()
.copied()
.map(|entity| (name.clone(), entity)),
);
Ok(())
}
pub fn flush_spawn_hooks(&mut self, ecs: ECSReference<'_>) {
let pending_batches = std::mem::take(&mut self.pending_spawn_batch_hooks);
for (id, entities) in pending_batches {
if let Ok(template) = self.get_by_id(id) {
if let Some(hook) = template.on_spawn_batch() {
hook(ecs, &entities);
}
}
}
let pending = std::mem::take(&mut self.pending_spawn_hooks);
for (name, entity) in pending {
if let Some(template) = self.templates.get(&name) {
if let Some(hook) = template.on_spawn() {
hook(ecs, entity);
}
}
}
}
pub fn flush_despawn_hooks(&mut self, ecs: ECSReference<'_>) {
let pending_batches = std::mem::take(&mut self.pending_despawn_batch_hooks);
for (id, entities) in pending_batches {
if let Ok(template) = self.get_by_id(id) {
if let Some(hook) = template.on_despawn_batch() {
hook(ecs, &entities);
}
}
}
let pending = std::mem::take(&mut self.pending_despawn_hooks);
for (name, entity) in pending {
if let Some(template) = self.templates.get(&name) {
if let Some(hook) = template.on_despawn() {
hook(ecs, entity);
}
}
}
}
}
impl Default for AgentRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agents::template::AgentTemplate;
fn make_template(name: &str) -> AgentTemplate {
AgentTemplate::builder(name).build()
}
#[test]
fn register_and_get() {
let mut reg = AgentRegistry::new();
reg.register(make_template("Wolf")).unwrap();
let tmpl = reg.get("Wolf").unwrap();
assert_eq!(tmpl.name(), "Wolf");
}
#[test]
fn duplicate_registration_is_rejected() {
let mut reg = AgentRegistry::new();
reg.register(make_template("Fox")).unwrap();
let err = reg.register(make_template("Fox")).unwrap_err();
assert!(matches!(err, AgentError::DuplicateTemplate(name) if name == "Fox"));
}
#[test]
fn sealed_registry_rejects_registration() {
let mut reg = AgentRegistry::new();
reg.seal();
let err = reg.register(make_template("Deer")).unwrap_err();
assert_eq!(err, AgentError::RegistrySealed);
}
#[test]
fn get_unknown_template_returns_error() {
let reg = AgentRegistry::new();
assert_eq!(
reg.get("Ghost").unwrap_err(),
AgentError::TemplateNotFound("Ghost".into()),
);
}
#[test]
fn seal_prevents_further_registration() {
let mut reg = AgentRegistry::new();
reg.register(make_template("A")).unwrap();
reg.seal();
assert!(reg.is_sealed());
assert_eq!(
reg.register(make_template("B")).unwrap_err(),
AgentError::RegistrySealed,
);
assert!(reg.get("A").is_ok());
}
#[test]
fn len_and_is_empty() {
let mut reg = AgentRegistry::new();
assert!(reg.is_empty());
reg.register(make_template("A")).unwrap();
assert_eq!(reg.len(), 1);
assert!(!reg.is_empty());
reg.register(make_template("B")).unwrap();
assert_eq!(reg.len(), 2);
}
#[test]
fn enqueue_and_check_pending() {
let mut reg = AgentRegistry::new();
reg.register(make_template("Sheep")).unwrap();
let entity: Entity = unsafe { std::mem::transmute(0u64) };
reg.enqueue_spawn_hook("Sheep", entity);
assert_eq!(reg.pending_spawn_hooks.len(), 1);
}
#[test]
fn enqueue_despawn_hook_tracks_pending() {
let mut reg = AgentRegistry::new();
reg.register(make_template("Sheep")).unwrap();
let entity: Entity = Entity::from_raw(0);
reg.enqueue_despawn_hook("Sheep", entity);
assert_eq!(reg.pending_despawn_hooks.len(), 1);
}
#[test]
fn flush_spawn_hooks_clears_queue() {
let mut reg = AgentRegistry::new();
reg.register(make_template("Rabbit")).unwrap();
let entity: Entity = unsafe { std::mem::transmute(0u64) };
reg.enqueue_spawn_hook("Rabbit", entity);
let _ = std::mem::take(&mut reg.pending_spawn_hooks);
assert!(reg.pending_spawn_hooks.is_empty());
}
#[test]
fn batch_enqueue_rejects_unknown_template_without_mutating_indexes() {
let mut reg = AgentRegistry::new();
let entity = Entity::from_raw(1);
let invalid = AgentTemplateId(99);
assert_eq!(
reg.enqueue_spawn_batch_hook(invalid, vec![entity])
.unwrap_err(),
AgentError::TemplateNotFound("#99".into())
);
assert!(reg.entity_template(entity).is_none());
assert!(reg.entities(invalid).is_empty());
assert!(reg.pending_spawn_batch_hooks.is_empty());
}
#[test]
fn batch_despawn_removes_entities_with_set_lookup() {
let mut reg = AgentRegistry::new();
reg.register(make_template("Sheep")).unwrap();
let template_id = reg.id("Sheep").unwrap();
let live: Vec<Entity> = (1..=5).map(Entity::from_raw).collect();
reg.enqueue_spawn_batch_hook(template_id, live.clone())
.unwrap();
reg.enqueue_despawn_batch_hook(template_id, vec![live[1], live[3], live[1]])
.unwrap();
assert_eq!(reg.entities(template_id), &[live[0], live[2], live[4]]);
assert!(reg.entity_template(live[1]).is_none());
assert!(reg.entity_template(live[3]).is_none());
}
}