Skip to main content

steel_core/poi/
poi_storage.rs

1//! World-level POI storage manager.
2//!
3//! Tracks special blocks (beds, workstations, bells, nether portals, etc.)
4//! so game systems can efficiently query for nearby points of interest
5//! without scanning every block. Organized by chunk column for efficient
6//! load/unload and spatial queries.
7
8use rustc_hash::FxHashMap;
9use steel_registry::{REGISTRY, RegistryExt};
10use steel_utils::{BlockPos, BlockStateId, ChunkPos, PackedSectionBlockPos, SectionPos};
11
12use super::poi_instance::PointOfInterest;
13use super::poi_set::PointOfInterestSet;
14use crate::chunk::section::ChunkSection;
15
16/// Filter for POI queries based on ticket availability.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum OccupationStatus {
19    /// Only POIs with at least one free ticket.
20    Free,
21    /// Only POIs with zero free tickets.
22    Occupied,
23    /// All POIs regardless of ticket status.
24    Any,
25}
26
27impl OccupationStatus {
28    /// Returns `true` if the given POI matches this status filter.
29    #[must_use]
30    pub const fn matches(self, poi: &PointOfInterest, max_tickets: u32) -> bool {
31        match self {
32            Self::Any => true,
33            Self::Free => poi.has_space(),
34            Self::Occupied => poi.is_occupied(max_tickets),
35        }
36    }
37}
38
39/// Column of POI sets indexed by section Y coordinate.
40type PoiColumn = FxHashMap<i32, PointOfInterestSet>;
41
42/// World-level storage for all points of interest.
43///
44/// Organized as a two-level map: `ChunkPos -> section_y -> PointOfInterestSet`.
45/// This structure mirrors chunk lifecycle (load/unload per column) and provides
46/// efficient spatial queries by narrowing to relevant columns first.
47pub struct PointOfInterestStorage {
48    columns: FxHashMap<ChunkPos, PoiColumn>,
49}
50
51impl Default for PointOfInterestStorage {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57#[inline]
58const fn resolve_pos(pos: BlockPos) -> (ChunkPos, i32, PackedSectionBlockPos) {
59    let section_pos = SectionPos::from_block_pos(pos);
60    let chunk_pos = ChunkPos::new(section_pos.x(), section_pos.z());
61    let packed = PackedSectionBlockPos::from_block_pos(pos);
62    (chunk_pos, section_pos.y(), packed)
63}
64
65fn max_tickets_for(type_id: usize) -> u32 {
66    REGISTRY
67        .poi_types
68        .by_id(type_id)
69        .map_or(0, |t| t.ticket_count)
70}
71
72fn distance_sq(a: BlockPos, b: BlockPos) -> i64 {
73    let dx = i64::from(a.0.x - b.0.x);
74    let dy = i64::from(a.0.y - b.0.y);
75    let dz = i64::from(a.0.z - b.0.z);
76    dx * dx + dy * dy + dz * dz
77}
78
79impl PointOfInterestStorage {
80    /// Creates an empty POI storage.
81    #[must_use]
82    pub fn new() -> Self {
83        Self {
84            columns: FxHashMap::default(),
85        }
86    }
87
88    fn get_or_create_set(
89        &mut self,
90        chunk_pos: ChunkPos,
91        section_y: i32,
92    ) -> &mut PointOfInterestSet {
93        self.columns
94            .entry(chunk_pos)
95            .or_default()
96            .entry(section_y)
97            .or_default()
98    }
99
100    /// Adds a POI at the given block position.
101    pub fn add(&mut self, pos: BlockPos, poi_type_id: usize, max_tickets: u32) {
102        let (chunk_pos, section_y, packed) = resolve_pos(pos);
103        let set = self.get_or_create_set(chunk_pos, section_y);
104        set.add(packed, PointOfInterest::new(pos, poi_type_id, max_tickets));
105    }
106
107    /// Removes the POI at the given block position.
108    pub fn remove(&mut self, pos: BlockPos) {
109        let (chunk_pos, section_y, packed) = resolve_pos(pos);
110        let Some(column) = self.columns.get_mut(&chunk_pos) else {
111            return;
112        };
113        let Some(set) = column.get_mut(&section_y) else {
114            return;
115        };
116
117        set.remove(packed);
118        if set.is_empty() {
119            column.remove(&section_y);
120            if column.is_empty() {
121                self.columns.remove(&chunk_pos);
122            }
123        }
124    }
125
126    /// Returns the POI type ID at the given position, if any.
127    #[must_use]
128    pub fn get_type(&self, pos: BlockPos) -> Option<usize> {
129        let (chunk_pos, section_y, packed) = resolve_pos(pos);
130        self.columns
131            .get(&chunk_pos)?
132            .get(&section_y)?
133            .get(packed)
134            .map(|poi| poi.poi_type_id)
135    }
136
137    /// Returns `true` if the POI at the given position has all tickets reserved.
138    #[must_use]
139    pub fn is_occupied(&self, pos: BlockPos) -> bool {
140        let (chunk_pos, section_y, packed) = resolve_pos(pos);
141        let Some(column) = self.columns.get(&chunk_pos) else {
142            return false;
143        };
144        let Some(set) = column.get(&section_y) else {
145            return false;
146        };
147        let Some(poi) = set.get(packed) else {
148            return false;
149        };
150        poi.is_occupied(max_tickets_for(poi.poi_type_id))
151    }
152
153    /// Reserves a ticket at the given position. Returns `true` if successful.
154    #[must_use]
155    pub fn reserve_ticket(&mut self, pos: BlockPos) -> bool {
156        let (chunk_pos, section_y, packed) = resolve_pos(pos);
157        let Some(set) = self
158            .columns
159            .get_mut(&chunk_pos)
160            .and_then(|c| c.get_mut(&section_y))
161        else {
162            return false;
163        };
164        let Some(poi) = set.get_mut(packed) else {
165            return false;
166        };
167        poi.reserve_ticket()
168    }
169
170    /// Releases a ticket at the given position. Returns `true` if successful.
171    #[must_use]
172    pub fn release_ticket(&mut self, pos: BlockPos) -> bool {
173        let (chunk_pos, section_y, packed) = resolve_pos(pos);
174        let Some(set) = self
175            .columns
176            .get_mut(&chunk_pos)
177            .and_then(|c| c.get_mut(&section_y))
178        else {
179            return false;
180        };
181        let Some(poi) = set.get_mut(packed) else {
182            return false;
183        };
184        poi.release_ticket(max_tickets_for(poi.poi_type_id))
185    }
186
187    /// Returns all matching POIs in a specific chunk column.
188    #[must_use]
189    pub fn get_in_chunk(
190        &self,
191        type_predicate: &impl Fn(usize) -> bool,
192        chunk_x: i32,
193        chunk_z: i32,
194        status: OccupationStatus,
195    ) -> Vec<(BlockPos, usize)> {
196        let chunk_pos = ChunkPos::new(chunk_x, chunk_z);
197        let Some(column) = self.columns.get(&chunk_pos) else {
198            return Vec::new();
199        };
200
201        let mut results = Vec::new();
202        for set in column.values() {
203            for poi in set.get_matching(type_predicate, status, &max_tickets_for) {
204                results.push((poi.pos, poi.poi_type_id));
205            }
206        }
207        results
208    }
209
210    /// Returns all matching POIs within a cubic region centered on `center`.
211    #[must_use]
212    pub fn get_in_square(
213        &self,
214        type_predicate: &impl Fn(usize) -> bool,
215        center: BlockPos,
216        radius: i32,
217        status: OccupationStatus,
218    ) -> Vec<(BlockPos, usize)> {
219        let min_section = SectionPos::from_block_pos(BlockPos::new(
220            center.0.x - radius,
221            center.0.y - radius,
222            center.0.z - radius,
223        ));
224        let max_section = SectionPos::from_block_pos(BlockPos::new(
225            center.0.x + radius,
226            center.0.y + radius,
227            center.0.z + radius,
228        ));
229
230        let mut results = Vec::new();
231
232        for cx in min_section.x()..=max_section.x() {
233            for cz in min_section.z()..=max_section.z() {
234                let chunk_pos = ChunkPos::new(cx, cz);
235                let Some(column) = self.columns.get(&chunk_pos) else {
236                    continue;
237                };
238
239                for section_y in min_section.y()..=max_section.y() {
240                    let Some(set) = column.get(&section_y) else {
241                        continue;
242                    };
243
244                    for poi in set.get_matching(type_predicate, status, &max_tickets_for) {
245                        let dx = (poi.pos.0.x - center.0.x).abs();
246                        let dy = (poi.pos.0.y - center.0.y).abs();
247                        let dz = (poi.pos.0.z - center.0.z).abs();
248
249                        if dx <= radius && dy <= radius && dz <= radius {
250                            results.push((poi.pos, poi.poi_type_id));
251                        }
252                    }
253                }
254            }
255        }
256
257        results
258    }
259
260    /// Returns all matching POIs within a vanilla horizontal square centered on `center`.
261    ///
262    /// Mirrors `PoiManager.getInSquare`: X/Z are constrained by `radius`, while Y is not.
263    #[must_use]
264    pub fn get_in_horizontal_square(
265        &self,
266        type_predicate: &impl Fn(usize) -> bool,
267        center: BlockPos,
268        radius: i32,
269        status: OccupationStatus,
270    ) -> Vec<(BlockPos, usize)> {
271        let center_chunk = ChunkPos::from_block_pos(center);
272        let chunk_radius = radius.div_euclid(16) + 1;
273        let mut results = Vec::new();
274
275        for cx in center_chunk.0.x - chunk_radius..=center_chunk.0.x + chunk_radius {
276            for cz in center_chunk.0.y - chunk_radius..=center_chunk.0.y + chunk_radius {
277                let chunk_pos = ChunkPos::new(cx, cz);
278                let Some(column) = self.columns.get(&chunk_pos) else {
279                    continue;
280                };
281
282                for set in column.values() {
283                    for poi in set.get_matching(type_predicate, status, &max_tickets_for) {
284                        let dx = (poi.pos.0.x - center.0.x).abs();
285                        let dz = (poi.pos.0.z - center.0.z).abs();
286                        if dx <= radius && dz <= radius {
287                            results.push((poi.pos, poi.poi_type_id));
288                        }
289                    }
290                }
291            }
292        }
293
294        results
295    }
296
297    /// Counts matching POIs within a spherical region.
298    #[must_use]
299    pub fn count(
300        &self,
301        type_predicate: &impl Fn(usize) -> bool,
302        pos: BlockPos,
303        radius: i32,
304        status: OccupationStatus,
305    ) -> usize {
306        let radius_sq = i64::from(radius) * i64::from(radius);
307        self.count_in_square(type_predicate, pos, radius, status, &|candidate| {
308            distance_sq(candidate, pos) <= radius_sq
309        })
310    }
311
312    /// Counts matching POIs within a cubic region, filtered by an additional predicate.
313    fn count_in_square(
314        &self,
315        type_predicate: &impl Fn(usize) -> bool,
316        center: BlockPos,
317        radius: i32,
318        status: OccupationStatus,
319        filter: &impl Fn(BlockPos) -> bool,
320    ) -> usize {
321        let min_section = SectionPos::from_block_pos(BlockPos::new(
322            center.0.x - radius,
323            center.0.y - radius,
324            center.0.z - radius,
325        ));
326        let max_section = SectionPos::from_block_pos(BlockPos::new(
327            center.0.x + radius,
328            center.0.y + radius,
329            center.0.z + radius,
330        ));
331
332        let mut count = 0;
333
334        for cx in min_section.x()..=max_section.x() {
335            for cz in min_section.z()..=max_section.z() {
336                let chunk_pos = ChunkPos::new(cx, cz);
337                let Some(column) = self.columns.get(&chunk_pos) else {
338                    continue;
339                };
340
341                for section_y in min_section.y()..=max_section.y() {
342                    let Some(set) = column.get(&section_y) else {
343                        continue;
344                    };
345
346                    for poi in set.get_matching(type_predicate, status, &max_tickets_for) {
347                        let dx = (poi.pos.0.x - center.0.x).abs();
348                        let dy = (poi.pos.0.y - center.0.y).abs();
349                        let dz = (poi.pos.0.z - center.0.z).abs();
350
351                        if dx <= radius && dy <= radius && dz <= radius && filter(poi.pos) {
352                            count += 1;
353                        }
354                    }
355                }
356            }
357        }
358
359        count
360    }
361
362    /// Scans a chunk section for POI block states and populates the storage.
363    ///
364    /// # Panics
365    /// Panics if the POI type registry contains an inconsistent state-to-type mapping.
366    pub fn scan_and_populate(&mut self, section: &ChunkSection, section_pos: SectionPos) {
367        let registry = &REGISTRY.poi_types;
368        let chunk_pos = ChunkPos::new(section_pos.x(), section_pos.z());
369        let set = self.get_or_create_set(chunk_pos, section_pos.y());
370
371        for y in 0..16u8 {
372            for z in 0..16u8 {
373                for x in 0..16u8 {
374                    let state_id = section.states.get(x as usize, y as usize, z as usize);
375
376                    let Some(poi_type_id) = registry.type_id_for_state(state_id) else {
377                        continue;
378                    };
379                    let poi_type = registry
380                        .by_id(poi_type_id)
381                        .expect("POI type ID from state lookup must be valid");
382                    let block_pos = BlockPos::new(
383                        (section_pos.x() << 4) + i32::from(x),
384                        (section_pos.y() << 4) + i32::from(y),
385                        (section_pos.z() << 4) + i32::from(z),
386                    );
387                    let packed = PackedSectionBlockPos::from_block_pos(block_pos);
388                    set.add(
389                        packed,
390                        PointOfInterest::new(block_pos, poi_type_id, poi_type.ticket_count),
391                    );
392                }
393            }
394        }
395    }
396
397    /// Updates POI storage when a block state changes.
398    ///
399    /// # Panics
400    /// Panics if the POI type registry contains an inconsistent state-to-type mapping.
401    pub fn on_block_state_change(
402        &mut self,
403        pos: BlockPos,
404        old_state: BlockStateId,
405        new_state: BlockStateId,
406    ) {
407        let registry = &REGISTRY.poi_types;
408        let old_poi = registry.type_id_for_state(old_state);
409        let new_poi = registry.type_id_for_state(new_state);
410
411        if old_poi == new_poi {
412            return;
413        }
414
415        if old_poi.is_some() {
416            self.remove(pos);
417        }
418
419        if let Some(type_id) = new_poi {
420            let poi_type = registry
421                .by_id(type_id)
422                .expect("POI type ID from state lookup must be valid");
423            self.add(pos, type_id, poi_type.ticket_count);
424        }
425    }
426
427    /// Collects all POI data in a chunk column for persistence.
428    ///
429    /// Returns `(BlockPos, free_tickets)` for each POI.
430    #[must_use]
431    pub fn collect_for_chunk(&self, chunk_pos: ChunkPos) -> Vec<(BlockPos, u32)> {
432        let Some(column) = self.columns.get(&chunk_pos) else {
433            return Vec::new();
434        };
435        let mut results = Vec::new();
436        for set in column.values() {
437            for (_, poi) in set.iter() {
438                results.push((poi.pos, poi.free_tickets));
439            }
440        }
441        results
442    }
443
444    /// Restores ticket state for POIs after loading from disk.
445    ///
446    /// Called after `scan_and_populate` has created fresh POIs from block states.
447    /// Applies saved `free_tickets` values to matching positions.
448    pub fn restore_tickets(&mut self, chunk_pos: ChunkPos, tickets: &[(BlockPos, u32)]) {
449        let Some(column) = self.columns.get_mut(&chunk_pos) else {
450            return;
451        };
452        for &(pos, free_tickets) in tickets {
453            let section_y = SectionPos::block_to_section_coord(pos.0.y);
454            let packed = PackedSectionBlockPos::from_block_pos(pos);
455            if let Some(set) = column.get_mut(&section_y)
456                && let Some(poi) = set.get_mut(packed)
457            {
458                poi.free_tickets = free_tickets;
459            }
460        }
461    }
462
463    /// Removes all POI data for a chunk column. Called during chunk unload.
464    pub fn remove_chunk(&mut self, chunk_pos: ChunkPos) {
465        self.columns.remove(&chunk_pos);
466    }
467}
468
469#[cfg(test)]
470mod tests {
471    use super::{OccupationStatus, PointOfInterestStorage};
472    use steel_registry::init_vanilla_registry;
473    use steel_utils::BlockPos;
474
475    fn sorted_positions(mut positions: Vec<(BlockPos, usize)>) -> Vec<BlockPos> {
476        positions.sort_by_key(|(pos, _)| (pos.x(), pos.y(), pos.z()));
477        positions.into_iter().map(|(pos, _)| pos).collect()
478    }
479
480    #[test]
481    fn horizontal_square_query_matches_y_unbounded_vanilla_search() {
482        init_vanilla_registry();
483        let mut storage = PointOfInterestStorage::new();
484        storage.add(BlockPos::new(0, -64, 0), 7, 0);
485        storage.add(BlockPos::new(0, 320, 0), 7, 0);
486        storage.add(BlockPos::new(2, 64, 0), 7, 0);
487        storage.add(BlockPos::new(0, 64, 2), 7, 0);
488
489        let positions = sorted_positions(storage.get_in_horizontal_square(
490            &|type_id| type_id == 7,
491            BlockPos::new(0, 64, 0),
492            1,
493            OccupationStatus::Any,
494        ));
495
496        assert_eq!(
497            positions,
498            vec![BlockPos::new(0, -64, 0), BlockPos::new(0, 320, 0)]
499        );
500    }
501}