use hashbrown::HashMap;
use std::cell::{Ref, RefCell, RefMut};
use crate::{agent::Agent, types::AgentId};
pub trait AgentStore<A: Agent> {
fn get(&self, id: AgentId) -> Option<Ref<'_, A>>;
fn get_mut(&self, id: AgentId) -> Option<RefMut<'_, A>>;
fn insert(&mut self, agent: A);
fn remove(&mut self, id: AgentId) -> Option<A>;
fn contains(&self, id: AgentId) -> bool {
self.get(id).is_some()
}
fn iter_ids(&self) -> Vec<AgentId>;
fn iter_ids_into(&self, buf: &mut Vec<AgentId>) {
buf.extend(self.iter_ids());
}
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[derive(Debug, Default)]
pub struct HashMapStore<A: Agent> {
agents: HashMap<AgentId, RefCell<A>>,
}
impl<A: Agent> HashMapStore<A> {
pub fn new() -> Self {
Self {
agents: HashMap::new(),
}
}
}
impl<A: Agent> AgentStore<A> for HashMapStore<A> {
fn get(&self, id: AgentId) -> Option<Ref<'_, A>> {
self.agents.get(&id).map(|cell| cell.borrow())
}
fn get_mut(&self, id: AgentId) -> Option<RefMut<'_, A>> {
self.agents.get(&id).map(|cell| cell.borrow_mut())
}
fn insert(&mut self, agent: A) {
self.agents.insert(agent.id(), RefCell::new(agent));
}
fn remove(&mut self, id: AgentId) -> Option<A> {
self.agents.remove(&id).map(|cell| cell.into_inner())
}
fn iter_ids(&self) -> Vec<AgentId> {
self.agents.keys().copied().collect()
}
fn iter_ids_into(&self, buf: &mut Vec<AgentId>) {
buf.extend(self.agents.keys().copied());
}
fn len(&self) -> usize {
self.agents.len()
}
}
#[derive(Debug, Default)]
pub struct VecStore<A: Agent> {
agents: Vec<Option<RefCell<A>>>,
count: usize,
}
impl<A: Agent> VecStore<A> {
pub fn new() -> Self {
Self {
agents: Vec::new(),
count: 0,
}
}
fn ensure_slot(&mut self, id: AgentId) {
let index = id as usize;
if index >= self.agents.len() {
self.agents.resize_with(index + 1, || None);
}
}
}
impl<A: Agent> AgentStore<A> for VecStore<A> {
fn get(&self, id: AgentId) -> Option<Ref<'_, A>> {
self.agents
.get(id as usize)
.and_then(|slot| slot.as_ref().map(|cell| cell.borrow()))
}
fn get_mut(&self, id: AgentId) -> Option<RefMut<'_, A>> {
self.agents
.get(id as usize)
.and_then(|slot| slot.as_ref().map(|cell| cell.borrow_mut()))
}
fn insert(&mut self, agent: A) {
let id = agent.id();
self.ensure_slot(id);
let slot = &mut self.agents[id as usize];
if slot.is_none() {
self.count += 1;
}
*slot = Some(RefCell::new(agent));
}
fn remove(&mut self, id: AgentId) -> Option<A> {
if (id as usize) < self.agents.len() {
let taken = self.agents[id as usize]
.take()
.map(|cell| cell.into_inner());
if taken.is_some() {
self.count -= 1;
}
taken
} else {
None
}
}
fn iter_ids(&self) -> Vec<AgentId> {
self.agents
.iter()
.enumerate()
.filter_map(|(idx, slot)| slot.as_ref().map(|_| idx as AgentId))
.collect()
}
fn iter_ids_into(&self, buf: &mut Vec<AgentId>) {
for (idx, slot) in self.agents.iter().enumerate() {
if slot.is_some() {
buf.push(idx as AgentId);
}
}
}
fn len(&self) -> usize {
self.count
}
}