1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum OccupationStatus {
19 Free,
21 Occupied,
23 Any,
25}
26
27impl OccupationStatus {
28 #[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
39type PoiColumn = FxHashMap<i32, PointOfInterestSet>;
41
42pub 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 #[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 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 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(§ion_y) else {
114 return;
115 };
116
117 set.remove(packed);
118 if set.is_empty() {
119 column.remove(§ion_y);
120 if column.is_empty() {
121 self.columns.remove(&chunk_pos);
122 }
123 }
124 }
125
126 #[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(§ion_y)?
133 .get(packed)
134 .map(|poi| poi.poi_type_id)
135 }
136
137 #[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(§ion_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 #[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(§ion_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 #[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(§ion_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 #[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 #[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(§ion_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 #[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 #[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 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(§ion_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 pub fn scan_and_populate(&mut self, section: &ChunkSection, section_pos: SectionPos) {
367 let registry = ®ISTRY.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 pub fn on_block_state_change(
402 &mut self,
403 pos: BlockPos,
404 old_state: BlockStateId,
405 new_state: BlockStateId,
406 ) {
407 let registry = ®ISTRY.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 #[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 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(§ion_y)
456 && let Some(poi) = set.get_mut(packed)
457 {
458 poi.free_tickets = free_tickets;
459 }
460 }
461 }
462
463 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}