use std::collections::{BTreeSet, HashMap};
use crate::components::RiderPhase;
use crate::entity::EntityId;
use crate::world::World;
type Partition = HashMap<EntityId, BTreeSet<EntityId>>;
#[derive(Debug, Clone, Default)]
pub struct RiderIndex {
waiting: Partition,
residents: Partition,
abandoned: Partition,
}
static EMPTY: std::sync::LazyLock<BTreeSet<EntityId>> = std::sync::LazyLock::new(BTreeSet::new);
fn insert(partition: &mut Partition, stop: EntityId, rider: EntityId) {
partition.entry(stop).or_default().insert(rider);
}
fn remove(partition: &mut Partition, stop: EntityId, rider: EntityId) {
if let Some(set) = partition.get_mut(&stop) {
set.remove(&rider);
if set.is_empty() {
partition.remove(&stop);
}
}
}
impl RiderIndex {
pub(crate) fn insert_waiting(&mut self, stop: EntityId, rider: EntityId) {
insert(&mut self.waiting, stop, rider);
}
pub(crate) fn insert_resident(&mut self, stop: EntityId, rider: EntityId) {
insert(&mut self.residents, stop, rider);
}
pub(crate) fn insert_abandoned(&mut self, stop: EntityId, rider: EntityId) {
insert(&mut self.abandoned, stop, rider);
}
pub(crate) fn remove_waiting(&mut self, stop: EntityId, rider: EntityId) {
remove(&mut self.waiting, stop, rider);
}
pub(crate) fn remove_resident(&mut self, stop: EntityId, rider: EntityId) {
remove(&mut self.residents, stop, rider);
}
pub(crate) fn remove_abandoned(&mut self, stop: EntityId, rider: EntityId) {
remove(&mut self.abandoned, stop, rider);
}
pub(crate) fn waiting_at(&self, stop: EntityId) -> &BTreeSet<EntityId> {
self.waiting.get(&stop).unwrap_or(&EMPTY)
}
pub(crate) fn residents_at(&self, stop: EntityId) -> &BTreeSet<EntityId> {
self.residents.get(&stop).unwrap_or(&EMPTY)
}
pub(crate) fn abandoned_at(&self, stop: EntityId) -> &BTreeSet<EntityId> {
self.abandoned.get(&stop).unwrap_or(&EMPTY)
}
pub(crate) fn waiting_count_at(&self, stop: EntityId) -> usize {
self.waiting.get(&stop).map_or(0, BTreeSet::len)
}
pub(crate) fn resident_count_at(&self, stop: EntityId) -> usize {
self.residents.get(&stop).map_or(0, BTreeSet::len)
}
pub(crate) fn abandoned_count_at(&self, stop: EntityId) -> usize {
self.abandoned.get(&stop).map_or(0, BTreeSet::len)
}
pub(crate) fn rebuild(&mut self, world: &World) {
self.waiting.clear();
self.residents.clear();
self.abandoned.clear();
for (id, rider) in world.iter_riders() {
if let Some(stop) = rider.current_stop() {
match rider.phase() {
RiderPhase::Waiting => self.insert_waiting(stop, id),
RiderPhase::Resident => self.insert_resident(stop, id),
RiderPhase::Abandoned => self.insert_abandoned(stop, id),
_ => {}
}
}
}
}
}