sectorsync_core/
spatial_index.rs

1//! Low-level cell index for station-local AOI candidate queries.
2
3use std::collections::{HashMap, HashSet};
4
5use crate::ids::EntityHandle;
6use crate::spatial::{Aabb3, Bounds, CellCoord3, GridSpec, Position3};
7
8const MAX_DENSE_DEDUP_SLOTS: usize = 262_144;
9
10/// Occupancy count for one non-empty cell.
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub struct CellOccupancy {
13    /// Cell coordinate.
14    pub cell: CellCoord3,
15    /// Number of indexed entity handles in the cell.
16    pub entities: usize,
17}
18
19/// Strategy used by the last scratch-backed cell query.
20#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
21pub enum CellQueryStrategy {
22    /// Probe every cell touched by the query bounds.
23    #[default]
24    Grid,
25    /// Scan non-empty cells when the query covers a larger sparse volume.
26    OccupiedCells,
27}
28
29/// Work counters from the last scratch-backed cell query.
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct CellQueryStats {
32    /// Query strategy selected from query volume and current index occupancy.
33    pub strategy: CellQueryStrategy,
34    /// Grid cells probed directly by a grid query.
35    pub grid_cells_probed: usize,
36    /// Non-empty cells inspected by an occupied-cell scan.
37    pub occupied_cells_scanned: usize,
38    /// Non-empty cells overlapping the query bounds.
39    pub matched_cells: usize,
40    /// Unique candidate handles produced by the query.
41    pub candidate_handles: usize,
42}
43
44/// Reusable scratch storage for allocation-aware cell queries.
45#[derive(Clone, Debug, Default)]
46pub struct CellQueryScratch {
47    seen_dense: Vec<u64>,
48    seen_collisions: Vec<u32>,
49    seen_sparse: HashSet<EntityHandle>,
50    query_epoch: u32,
51    handles: Vec<EntityHandle>,
52    matching_cells: Vec<CellCoord3>,
53    stats: CellQueryStats,
54}
55
56impl CellQueryScratch {
57    /// Clears retained query results while keeping allocated capacity.
58    pub fn clear(&mut self) {
59        self.begin_query();
60        self.handles.clear();
61        self.matching_cells.clear();
62        self.stats = CellQueryStats::default();
63    }
64
65    /// Returns handles produced by the last query.
66    pub fn handles(&self) -> &[EntityHandle] {
67        &self.handles
68    }
69
70    /// Number of handles produced by the last query.
71    pub fn len(&self) -> usize {
72        self.handles.len()
73    }
74
75    /// Returns whether the last query produced no handles.
76    pub fn is_empty(&self) -> bool {
77        self.handles.is_empty()
78    }
79
80    /// Work counters produced by the last query.
81    pub const fn stats(&self) -> CellQueryStats {
82        self.stats
83    }
84
85    /// Capacity retained for unique candidate handles.
86    pub fn handle_capacity(&self) -> usize {
87        self.handles.capacity()
88    }
89
90    /// Capacity retained by the candidate deduplication set.
91    pub fn dedup_capacity(&self) -> usize {
92        self.seen_dense
93            .capacity()
94            .saturating_add(self.seen_collisions.capacity())
95            .saturating_add(self.seen_sparse.capacity())
96    }
97
98    /// Capacity retained for occupied cells matched by sparse queries.
99    pub fn matching_cell_capacity(&self) -> usize {
100        self.matching_cells.capacity()
101    }
102
103    fn begin_query(&mut self) {
104        self.query_epoch = self.query_epoch.wrapping_add(1);
105        if self.query_epoch == 0 {
106            self.seen_dense.fill(0);
107            self.seen_collisions.fill(0);
108            self.query_epoch = 1;
109        }
110        self.seen_sparse.clear();
111    }
112
113    fn insert_seen(&mut self, handle: EntityHandle) -> bool {
114        let Ok(index) = usize::try_from(handle.index()) else {
115            return self.seen_sparse.insert(handle);
116        };
117        if index >= MAX_DENSE_DEDUP_SLOTS {
118            return self.seen_sparse.insert(handle);
119        }
120        if index >= self.seen_dense.len() {
121            self.seen_dense.resize(index + 1, 0);
122            self.seen_collisions.resize(index + 1, 0);
123        }
124        let marker = (u64::from(self.query_epoch) << 32) | u64::from(handle.generation());
125        let previous = self.seen_dense[index];
126        let previous_epoch = u32::try_from(previous >> 32).expect("marker epoch fits u32");
127        if previous_epoch != self.query_epoch {
128            self.seen_dense[index] = marker;
129            true
130        } else if self.seen_collisions[index] == self.query_epoch {
131            self.seen_sparse.insert(handle)
132        } else if previous == marker {
133            false
134        } else {
135            self.seen_collisions[index] = self.query_epoch;
136            let previous_generation =
137                u32::try_from(previous & u64::from(u32::MAX)).expect("marker generation fits u32");
138            self.seen_sparse
139                .insert(EntityHandle::new(handle.index(), previous_generation));
140            self.seen_sparse.insert(handle)
141        }
142    }
143}
144
145/// Station-local 3D cell index.
146#[derive(Clone, Debug)]
147pub struct CellIndex {
148    grid: GridSpec,
149    cells: HashMap<CellCoord3, Vec<EntityHandle>>,
150    entity_cells: HashMap<EntityHandle, Vec<CellCoord3>>,
151}
152
153impl CellIndex {
154    /// Creates an empty cell index.
155    pub fn new(grid: GridSpec) -> Self {
156        Self {
157            grid,
158            cells: HashMap::new(),
159            entity_cells: HashMap::new(),
160        }
161    }
162
163    /// Returns the grid spec.
164    pub const fn grid(&self) -> GridSpec {
165        self.grid
166    }
167
168    /// Inserts or updates an entity in all cells touched by its bounds.
169    pub fn upsert(&mut self, handle: EntityHandle, position: Position3, bounds: Bounds) {
170        self.remove(handle);
171        let cells = self.grid.cells_for_bounds(position, bounds);
172        for cell in &cells {
173            self.cells.entry(*cell).or_default().push(handle);
174        }
175        self.entity_cells.insert(handle, cells);
176    }
177
178    /// Removes an entity from the index.
179    pub fn remove(&mut self, handle: EntityHandle) -> bool {
180        let Some(cells) = self.entity_cells.remove(&handle) else {
181            return false;
182        };
183
184        for cell in cells {
185            let remove_cell = if let Some(handles) = self.cells.get_mut(&cell) {
186                handles.retain(|candidate| *candidate != handle);
187                handles.is_empty()
188            } else {
189                false
190            };
191            if remove_cell {
192                self.cells.remove(&cell);
193            }
194        }
195
196        true
197    }
198
199    /// Queries candidate handles overlapping an AABB.
200    pub fn query_aabb(&self, aabb: Aabb3) -> Vec<EntityHandle> {
201        let mut scratch = CellQueryScratch::default();
202        self.query_aabb_into(aabb, &mut scratch);
203        scratch.handles
204    }
205
206    /// Queries candidate handles overlapping an AABB using caller scratch.
207    pub fn query_aabb_into<'a>(
208        &self,
209        aabb: Aabb3,
210        scratch: &'a mut CellQueryScratch,
211    ) -> &'a [EntityHandle] {
212        scratch.clear();
213        let min = self.grid.cell_at(aabb.min);
214        let max = self.grid.cell_at(aabb.max);
215
216        let grid_cells = query_cell_volume(min, max);
217        if grid_cells <= self.cells.len() {
218            scratch.stats.strategy = CellQueryStrategy::Grid;
219            scratch.stats.grid_cells_probed = grid_cells;
220            for x in min.x..=max.x {
221                for y in min.y..=max.y {
222                    for z in min.z..=max.z {
223                        self.collect_cell(CellCoord3::new(x, y, z), scratch);
224                    }
225                }
226            }
227        } else {
228            scratch.stats.strategy = CellQueryStrategy::OccupiedCells;
229            scratch.stats.occupied_cells_scanned = self.cells.len();
230            scratch.matching_cells.extend(
231                self.cells
232                    .keys()
233                    .copied()
234                    .filter(|cell| cell_in_range(*cell, min, max)),
235            );
236            scratch.matching_cells.sort_unstable();
237            for index in 0..scratch.matching_cells.len() {
238                self.collect_cell(scratch.matching_cells[index], scratch);
239            }
240        }
241
242        scratch.stats.candidate_handles = scratch.handles.len();
243
244        scratch.handles()
245    }
246
247    /// Queries candidate handles inside cells touched by a sphere.
248    pub fn query_sphere(&self, center: Position3, radius: f32) -> Vec<EntityHandle> {
249        self.query_aabb(Bounds::Sphere { radius }.to_aabb(center))
250    }
251
252    /// Queries candidate handles inside cells touched by a sphere using caller scratch.
253    pub fn query_sphere_into<'a>(
254        &self,
255        center: Position3,
256        radius: f32,
257        scratch: &'a mut CellQueryScratch,
258    ) -> &'a [EntityHandle] {
259        self.query_aabb_into(Bounds::Sphere { radius }.to_aabb(center), scratch)
260    }
261
262    fn collect_cell(&self, cell: CellCoord3, scratch: &mut CellQueryScratch) {
263        if let Some(handles) = self.cells.get(&cell) {
264            scratch.stats.matched_cells = scratch.stats.matched_cells.saturating_add(1);
265            for handle in handles {
266                if scratch.insert_seen(*handle) {
267                    scratch.handles.push(*handle);
268                }
269            }
270        }
271    }
272
273    /// Returns handles indexed directly in one cell.
274    pub fn handles_in_cell(&self, cell: CellCoord3) -> Vec<EntityHandle> {
275        self.cells.get(&cell).cloned().unwrap_or_default()
276    }
277
278    /// Returns handles indexed directly in one cell without allocating.
279    pub fn handles_in_cell_slice(&self, cell: CellCoord3) -> &[EntityHandle] {
280        self.cells.get(&cell).map_or(&[], Vec::as_slice)
281    }
282
283    /// Returns cells currently occupied by one entity handle.
284    pub fn cells_for_handle(&self, handle: EntityHandle) -> Option<&[CellCoord3]> {
285        self.entity_cells.get(&handle).map(Vec::as_slice)
286    }
287
288    /// Number of indexed entities.
289    pub fn entity_count(&self) -> usize {
290        self.entity_cells.len()
291    }
292
293    /// Number of non-empty cells.
294    pub fn occupied_cell_count(&self) -> usize {
295        self.cells.len()
296    }
297
298    /// Returns deterministic occupancy counts for all non-empty cells.
299    pub fn cell_occupancy(&self) -> Vec<CellOccupancy> {
300        let mut cells = self
301            .cells
302            .iter()
303            .map(|(cell, handles)| CellOccupancy {
304                cell: *cell,
305                entities: handles.len(),
306            })
307            .collect::<Vec<_>>();
308        cells.sort_by_key(|occupancy| occupancy.cell);
309        cells
310    }
311}
312
313fn query_cell_volume(min: CellCoord3, max: CellCoord3) -> usize {
314    fn axis_cells(min: i32, max: i32) -> usize {
315        if max < min {
316            return 0;
317        }
318        usize::try_from(i64::from(max) - i64::from(min) + 1).unwrap_or(usize::MAX)
319    }
320
321    axis_cells(min.x, max.x)
322        .saturating_mul(axis_cells(min.y, max.y))
323        .saturating_mul(axis_cells(min.z, max.z))
324}
325
326const fn cell_in_range(cell: CellCoord3, min: CellCoord3, max: CellCoord3) -> bool {
327    cell.x >= min.x
328        && cell.x <= max.x
329        && cell.y >= min.y
330        && cell.y <= max.y
331        && cell.z >= min.z
332        && cell.z <= max.z
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338
339    #[test]
340    fn index_exposes_handles_by_cell() {
341        let grid = GridSpec::new(10.0).expect("valid grid");
342        let mut index = CellIndex::new(grid);
343        let handle = EntityHandle::new(1, 0);
344        index.upsert(handle, Position3::new(1.0, 2.0, 3.0), Bounds::Point);
345        let cell = grid.cell_at(Position3::new(1.0, 2.0, 3.0));
346
347        assert_eq!(index.handles_in_cell(cell), vec![handle]);
348        assert_eq!(index.handles_in_cell_slice(cell), &[handle]);
349        assert!(
350            index
351                .handles_in_cell_slice(CellCoord3::new(99, 99, 99))
352                .is_empty()
353        );
354        assert_eq!(index.cells_for_handle(handle), Some([cell].as_slice()));
355    }
356
357    #[test]
358    fn scratch_query_deduplicates_and_reuses_storage() {
359        let grid = GridSpec::new(10.0).expect("valid grid");
360        let mut index = CellIndex::new(grid);
361        let handle = EntityHandle::new(1, 0);
362        index.upsert(
363            handle,
364            Position3::new(9.0, 0.0, 0.0),
365            Bounds::Sphere { radius: 2.0 },
366        );
367        let mut scratch = CellQueryScratch::default();
368
369        let first = index.query_aabb_into(
370            Bounds::Sphere { radius: 4.0 }.to_aabb(Position3::new(10.0, 0.0, 0.0)),
371            &mut scratch,
372        );
373        assert_eq!(first, &[handle]);
374        assert_eq!(scratch.len(), 1);
375        assert_eq!(scratch.stats().strategy, CellQueryStrategy::Grid);
376        assert_eq!(scratch.stats().candidate_handles, 1);
377        assert!(scratch.handle_capacity() >= 1);
378        assert!(scratch.dedup_capacity() >= 1);
379
380        let second = index.query_aabb_into(
381            Bounds::Point.to_aabb(Position3::new(100.0, 0.0, 0.0)),
382            &mut scratch,
383        );
384        assert!(second.is_empty());
385        assert!(scratch.is_empty());
386    }
387
388    #[test]
389    fn dense_dedup_preserves_generations_and_bounds_sparse_handles() {
390        let grid = GridSpec::new(10.0).expect("valid grid");
391        let mut index = CellIndex::new(grid);
392        let old = EntityHandle::new(7, 1);
393        let current = EntityHandle::new(7, 2);
394        let sparse = EntityHandle::new(u32::MAX, 0);
395        let spanning = Bounds::Sphere { radius: 2.0 };
396        index.upsert(old, Position3::new(9.0, 0.0, 0.0), spanning);
397        index.upsert(current, Position3::new(9.0, 0.0, 0.0), spanning);
398        index.upsert(sparse, Position3::new(9.0, 0.0, 0.0), Bounds::Point);
399        let mut scratch = CellQueryScratch::default();
400
401        let handles = index.query_aabb_into(
402            spanning.to_aabb(Position3::new(10.0, 0.0, 0.0)),
403            &mut scratch,
404        );
405
406        assert_eq!(handles, &[old, current, sparse]);
407        assert!(scratch.dedup_capacity() < MAX_DENSE_DEDUP_SLOTS);
408    }
409
410    #[test]
411    fn sparse_large_query_scans_occupied_cells_deterministically() {
412        let grid = GridSpec::new(10.0).expect("valid grid");
413        let mut index = CellIndex::new(grid);
414        let high = EntityHandle::new(2, 0);
415        let low = EntityHandle::new(1, 0);
416        index.upsert(high, Position3::new(95.0, 0.0, 0.0), Bounds::Point);
417        index.upsert(low, Position3::new(-95.0, 0.0, 0.0), Bounds::Point);
418        let mut scratch = CellQueryScratch::default();
419
420        let handles = index.query_aabb_into(
421            Aabb3::new(
422                Position3::new(-100.0, -100.0, -100.0),
423                Position3::new(100.0, 100.0, 100.0),
424            ),
425            &mut scratch,
426        );
427
428        assert_eq!(handles, &[low, high]);
429        assert_eq!(scratch.stats().strategy, CellQueryStrategy::OccupiedCells);
430        assert_eq!(scratch.stats().occupied_cells_scanned, 2);
431        assert_eq!(scratch.stats().matched_cells, 2);
432        assert_eq!(scratch.stats().candidate_handles, 2);
433        assert!(scratch.matching_cell_capacity() >= 2);
434    }
435}
sectorsync_core/spatial_index.rs

sectorsync_core/
spatial_index.rs
