bevy_a5 0.1.2

//! Cell → entity spatial index.
//!
//! Maintains a `HashMap<u64, SmallVec<[Entity; 4]>>` keyed by raw A5 cell ID
//! so apps can look up "what entities are in cell X?" in O(1).
//!
//! ```rust,ignore
//! use bevy_a5::prelude::*;
//! use bevy_a5::hash::{CellHashPlugin, CellEntityIndex};
//!
//! App::new().add_plugins((BevyA5Plugin, CellHashPlugin));
//!
//! fn neighbours_in_cell(index: Res<CellEntityIndex>, cell: GeoCell) {
//!     for &entity in index.entities_in(cell.raw()) {
//!         // ...
//!     }
//! }
//! ```

use bevy_app::prelude::*;
use bevy_ecs::prelude::*;
use bevy_platform::collections::HashMap;
use smallvec::SmallVec;

use crate::cell::GeoCell;

/// Inline capacity for each cell's entity bucket. At realistic densities
/// most cells hold 0–4 entities, so this lets the common-case bucket live
/// in the `HashMap`'s value slot with no heap allocation.
const BUCKET_INLINE: usize = 4;

type Bucket = SmallVec<[Entity; BUCKET_INLINE]>;

/// Resource: maps raw A5 cell index → entities currently in that cell.
///
/// Maintained by [`update_cell_entity_index`]; treat the returned slices as
/// read-only — modifying the index from user code will fight the system.
///
/// Each cell's entity list is a [`SmallVec`] inline-sized for `BUCKET_INLINE`
/// entries, so the common case (sparse cells, ≤4 entities) avoids the
/// per-bucket heap allocation a plain `Vec<Entity>` would pay.
#[derive(Resource, Default, Debug)]
pub struct CellEntityIndex {
    map: HashMap<u64, Bucket>,
    last_cell: HashMap<Entity, u64>,
}

impl CellEntityIndex {
    /// Borrow the entities currently in `cell`. Empty slice if none.
    pub fn entities_in(&self, cell: u64) -> &[Entity] {
        self.map.get(&cell).map(SmallVec::as_slice).unwrap_or(&[])
    }

    /// Iterate every (cell, entities-in-cell) pair currently tracked. The
    /// inner slice is guaranteed to be non-empty (empty buckets are pruned
    /// on remove).
    pub fn iter(&self) -> impl Iterator<Item = (&u64, &[Entity])> {
        self.map.iter().map(|(c, b)| (c, b.as_slice()))
    }

    /// Total number of distinct cells with at least one entity.
    pub fn occupied_cell_count(&self) -> usize {
        self.map.len()
    }

    fn insert(&mut self, entity: Entity, cell: u64) {
        // Record the new cell, getting back the previously-recorded one.
        if let Some(prev) = self.last_cell.insert(entity, cell) {
            if prev == cell {
                // Entity was already in this cell — nothing to do.
                return;
            }
            // Pull the entity out of its previous bucket.
            if let Some(bucket) = self.map.get_mut(&prev) {
                bucket.retain(|e| *e != entity);
                if bucket.is_empty() {
                    self.map.remove(&prev);
                }
            }
        }
        // Either this is the entity's first cell, or it just changed.
        self.map.entry(cell).or_default().push(entity);
    }

    fn remove(&mut self, entity: Entity) {
        if let Some(prev) = self.last_cell.remove(&entity)
            && let Some(bucket) = self.map.get_mut(&prev)
        {
            bucket.retain(|e| *e != entity);
            if bucket.is_empty() {
                self.map.remove(&prev);
            }
        }
    }
}

/// Plugin: registers [`CellEntityIndex`] and the system that maintains it.
pub struct CellHashPlugin;

impl Plugin for CellHashPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<CellEntityIndex>();
        app.add_systems(PostUpdate, update_cell_entity_index);
    }
}

/// Walks every entity with a `GeoCell`, syncs additions / mutations into the
/// index, and prunes entities whose `GeoCell` was removed.
///
/// Runs in `PostUpdate`; cheap per frame because it only iterates `Changed`
/// entries plus a removal sweep. Pair with the [`CellTracker`](crate::cell::CellTracker)
/// helper if your code positions entities in world space — the tracker
/// avoids redundant `GeoCell` writes that would otherwise re-trigger this
/// system every frame.
pub fn update_cell_entity_index(
    mut index: ResMut<CellEntityIndex>,
    changed: Query<(Entity, &GeoCell), Changed<GeoCell>>,
    mut removed: RemovedComponents<GeoCell>,
) {
    for entity in removed.read() {
        index.remove(entity);
    }
    for (entity, cell) in changed.iter() {
        index.insert(entity, cell.raw());
    }
}