1use glam::DVec3;
4use rustc_hash::FxHashMap;
5use steel_utils::BlockPos;
6
7use crate::entity::ai::path::PathType;
8
9#[derive(Debug, Clone, PartialEq)]
10pub struct Node {
11 pub x: i32,
12 pub y: i32,
13 pub z: i32,
14 hash: i32,
15 pub heap_idx: i32,
16 pub g: f32,
17 pub h: f32,
18 pub f: f32,
19 pub came_from: Option<i32>,
20 pub closed: bool,
21 pub walked_distance: f32,
22 pub cost_malus: f32,
23 pub path_type: PathType,
24}
25
26impl Node {
27 #[must_use]
28 pub const fn new(x: i32, y: i32, z: i32) -> Self {
29 Self {
30 x,
31 y,
32 z,
33 hash: Self::create_hash(x, y, z),
34 heap_idx: -1,
35 g: 0.0,
36 h: 0.0,
37 f: 0.0,
38 came_from: None,
39 closed: false,
40 walked_distance: 0.0,
41 cost_malus: 0.0,
42 path_type: PathType::Blocked,
43 }
44 }
45
46 #[must_use]
47 pub const fn clone_and_move(&self, x: i32, y: i32, z: i32) -> Self {
48 let mut node = Self::new(x, y, z);
49 node.heap_idx = self.heap_idx;
50 node.g = self.g;
51 node.h = self.h;
52 node.f = self.f;
53 node.came_from = self.came_from;
54 node.closed = self.closed;
55 node.walked_distance = self.walked_distance;
56 node.cost_malus = self.cost_malus;
57 node.path_type = self.path_type;
58 node
59 }
60
61 #[must_use]
62 pub const fn create_hash(x: i32, y: i32, z: i32) -> i32 {
63 let mut hash =
64 ((y as u32) & 0xff) | (((x as u32) & 32_767) << 8) | (((z as u32) & 32_767) << 24);
65 if x < 0 {
66 hash |= 0x8000_0000;
67 }
68 if z < 0 {
69 hash |= 32_768;
70 }
71 hash as i32
72 }
73
74 #[must_use]
75 pub const fn hash(&self) -> i32 {
76 self.hash
77 }
78
79 #[must_use]
80 pub fn distance_to(&self, to: &Self) -> f32 {
81 let xd = (to.x - self.x) as f32;
82 let yd = (to.y - self.y) as f32;
83 let zd = (to.z - self.z) as f32;
84 xd.mul_add(xd, yd.mul_add(yd, zd * zd)).sqrt()
85 }
86
87 #[must_use]
88 pub fn distance_to_pos(&self, pos: BlockPos) -> f32 {
89 let xd = (pos.x() - self.x) as f32;
90 let yd = (pos.y() - self.y) as f32;
91 let zd = (pos.z() - self.z) as f32;
92 xd.mul_add(xd, yd.mul_add(yd, zd * zd)).sqrt()
93 }
94
95 #[must_use]
96 pub fn distance_manhattan_to_pos(&self, pos: BlockPos) -> f32 {
97 (pos.x() - self.x).abs() as f32
98 + (pos.y() - self.y).abs() as f32
99 + (pos.z() - self.z).abs() as f32
100 }
101
102 #[must_use]
103 pub fn distance_to_sqr(&self, to: &Self) -> f32 {
104 let xd = (to.x - self.x) as f32;
105 let yd = (to.y - self.y) as f32;
106 let zd = (to.z - self.z) as f32;
107 xd.mul_add(xd, yd.mul_add(yd, zd * zd))
108 }
109
110 #[must_use]
111 pub fn distance_manhattan(&self, to: &Self) -> f32 {
112 (to.x - self.x).abs() as f32 + (to.y - self.y).abs() as f32 + (to.z - self.z).abs() as f32
113 }
114
115 #[must_use]
116 pub const fn as_block_pos(&self) -> BlockPos {
117 BlockPos::new(self.x, self.y, self.z)
118 }
119
120 #[must_use]
121 pub fn as_vec3(&self) -> DVec3 {
122 DVec3::new(f64::from(self.x), f64::from(self.y), f64::from(self.z))
123 }
124
125 #[must_use]
126 pub const fn in_open_set(&self) -> bool {
127 self.heap_idx >= 0
128 }
129}
130
131#[derive(Debug, Clone, PartialEq)]
132pub struct Target {
133 node: Node,
134 best_heuristic: f32,
135 best_node: Option<i32>,
136 reached: bool,
137}
138
139impl Target {
140 #[must_use]
141 pub const fn new(node: Node) -> Self {
142 Self {
143 node,
144 best_heuristic: f32::MAX,
145 best_node: None,
146 reached: false,
147 }
148 }
149
150 #[must_use]
151 pub const fn node(&self) -> &Node {
152 &self.node
153 }
154
155 pub fn update_best(&mut self, heuristic: f32, node: &Node) {
156 if heuristic < self.best_heuristic {
157 self.best_heuristic = heuristic;
158 self.best_node = Some(node.hash());
159 }
160 }
161
162 #[must_use]
163 pub const fn best_node(&self) -> Option<i32> {
164 self.best_node
165 }
166
167 pub const fn set_reached(&mut self) {
168 self.reached = true;
169 }
170
171 #[must_use]
172 pub const fn is_reached(&self) -> bool {
173 self.reached
174 }
175}
176
177#[derive(Debug, Default, Clone)]
178pub struct NodeStore {
179 nodes: FxHashMap<i32, Node>,
180}
181
182impl NodeStore {
183 #[must_use]
184 pub fn new() -> Self {
185 Self::default()
186 }
187
188 pub fn clear(&mut self) {
189 self.nodes.clear();
190 }
191
192 pub fn reset_search_state(&mut self) {
193 for node in self.nodes.values_mut() {
194 node.heap_idx = -1;
195 node.g = 0.0;
196 node.h = 0.0;
197 node.f = 0.0;
198 node.came_from = None;
199 node.closed = false;
200 node.walked_distance = 0.0;
201 }
202 }
203
204 pub fn get_node(&mut self, x: i32, y: i32, z: i32) -> &mut Node {
205 let hash = Node::create_hash(x, y, z);
206 self.nodes.entry(hash).or_insert_with(|| Node::new(x, y, z))
207 }
208
209 #[must_use]
210 pub fn get(&self, hash: i32) -> Option<&Node> {
211 self.nodes.get(&hash)
212 }
213
214 pub fn get_mut(&mut self, hash: i32) -> Option<&mut Node> {
215 self.nodes.get_mut(&hash)
216 }
217
218 #[must_use]
219 pub fn len(&self) -> usize {
220 self.nodes.len()
221 }
222
223 #[must_use]
224 pub fn is_empty(&self) -> bool {
225 self.nodes.is_empty()
226 }
227}
228
229#[derive(Debug, Clone, PartialEq, Eq)]
230pub struct NodeHeap {
231 heap: Vec<i32>,
232}
233
234impl NodeHeap {
235 #[must_use]
236 pub const fn new() -> Self {
237 Self { heap: Vec::new() }
238 }
239
240 pub fn clear(&mut self, nodes: &mut NodeStore) {
241 for hash in self.heap.drain(..) {
242 if let Some(node) = nodes.get_mut(hash) {
243 node.heap_idx = -1;
244 }
245 }
246 }
247
248 #[must_use]
249 pub fn peek(&self) -> Option<i32> {
250 self.heap.first().copied()
251 }
252
253 pub fn insert(&mut self, nodes: &mut NodeStore, hash: i32) -> bool {
254 let Some(node) = nodes.get(hash) else {
255 return false;
256 };
257 if node.in_open_set() {
258 return false;
259 }
260
261 self.heap.push(hash);
262 let index = self.heap.len() - 1;
263 Self::set_heap_idx(nodes, hash, index) && self.up_heap(nodes, index)
264 }
265
266 pub fn pop(&mut self, nodes: &mut NodeStore) -> Option<i32> {
267 let popped = self.heap.first().copied()?;
268 let last = self.heap.pop()?;
269 if !self.heap.is_empty() {
270 self.heap[0] = last;
271 if !Self::set_heap_idx(nodes, last, 0) || !self.down_heap(nodes, 0) {
272 return None;
273 }
274 }
275
276 Self::set_heap_idx_to_removed(nodes, popped);
277 Some(popped)
278 }
279
280 pub fn change_cost(&mut self, nodes: &mut NodeStore, hash: i32, new_cost: f32) -> bool {
281 let Some(node) = nodes.get_mut(hash) else {
282 return false;
283 };
284 let old_cost = node.f;
285 let heap_idx = node.heap_idx;
286 node.f = new_cost;
287 if heap_idx < 0 {
288 return false;
289 }
290
291 let index = heap_idx as usize;
292 if self.heap.get(index).copied() != Some(hash) {
293 return false;
294 }
295
296 if new_cost < old_cost {
297 self.up_heap(nodes, index)
298 } else {
299 self.down_heap(nodes, index)
300 }
301 }
302
303 #[must_use]
304 pub const fn len(&self) -> usize {
305 self.heap.len()
306 }
307
308 #[must_use]
309 pub const fn is_empty(&self) -> bool {
310 self.heap.is_empty()
311 }
312
313 fn up_heap(&mut self, nodes: &mut NodeStore, mut index: usize) -> bool {
314 let Some(node_hash) = self.heap.get(index).copied() else {
315 return false;
316 };
317 let Some(cost) = nodes.get(node_hash).map(|node| node.f) else {
318 return false;
319 };
320
321 while index > 0 {
322 let parent_index = (index - 1) >> 1;
323 let Some(parent_hash) = self.heap.get(parent_index).copied() else {
324 return false;
325 };
326 let Some(parent_cost) = nodes.get(parent_hash).map(|node| node.f) else {
327 return false;
328 };
329 if cost >= parent_cost {
330 break;
331 }
332
333 self.heap[index] = parent_hash;
334 if !Self::set_heap_idx(nodes, parent_hash, index) {
335 return false;
336 }
337 index = parent_index;
338 }
339
340 self.heap[index] = node_hash;
341 Self::set_heap_idx(nodes, node_hash, index)
342 }
343
344 fn down_heap(&mut self, nodes: &mut NodeStore, mut index: usize) -> bool {
345 let Some(node_hash) = self.heap.get(index).copied() else {
346 return false;
347 };
348 let Some(cost) = nodes.get(node_hash).map(|node| node.f) else {
349 return false;
350 };
351
352 loop {
353 let left_index = 1 + (index << 1);
354 let right_index = left_index + 1;
355 if left_index >= self.heap.len() {
356 break;
357 }
358
359 let Some(left_hash) = self.heap.get(left_index).copied() else {
360 return false;
361 };
362 let Some(left_cost) = nodes.get(left_hash).map(|node| node.f) else {
363 return false;
364 };
365 let right = self
366 .heap
367 .get(right_index)
368 .and_then(|hash| nodes.get(*hash).map(|node| (*hash, node.f)));
369
370 let (child_index, child_hash, child_cost) = match right {
371 Some((right_hash, right_cost)) if right_cost <= left_cost => {
372 (right_index, right_hash, right_cost)
373 }
374 _ => (left_index, left_hash, left_cost),
375 };
376
377 if child_cost >= cost {
378 break;
379 }
380
381 self.heap[index] = child_hash;
382 if !Self::set_heap_idx(nodes, child_hash, index) {
383 return false;
384 }
385 index = child_index;
386 }
387
388 self.heap[index] = node_hash;
389 Self::set_heap_idx(nodes, node_hash, index)
390 }
391
392 fn set_heap_idx(nodes: &mut NodeStore, hash: i32, index: usize) -> bool {
393 let Ok(heap_idx) = i32::try_from(index) else {
394 return false;
395 };
396 let Some(node) = nodes.get_mut(hash) else {
397 return false;
398 };
399 node.heap_idx = heap_idx;
400 true
401 }
402
403 fn set_heap_idx_to_removed(nodes: &mut NodeStore, hash: i32) {
404 if let Some(node) = nodes.get_mut(hash) {
405 node.heap_idx = -1;
406 }
407 }
408}
409
410impl Default for NodeHeap {
411 fn default() -> Self {
412 Self::new()
413 }
414}
415
416#[cfg(test)]
417mod tests {
418 use super::{Node, NodeHeap, NodeStore, Target};
419 use crate::entity::ai::path::PathType;
420
421 #[test]
422 fn node_hash_matches_vanilla_packing_shape() {
423 assert_eq!(Node::create_hash(0, 0, 0), 0);
424 assert_eq!(Node::create_hash(1, 64, 2), 0x0200_0140);
425 assert_eq!(Node::create_hash(-1, 0, 0), 0x807f_ff00_u32 as i32);
426 assert_eq!(Node::create_hash(0, 0, -1), 0xff00_8000_u32 as i32);
427 }
428
429 #[test]
430 fn node_store_reuses_nodes_by_vanilla_hash() {
431 let mut store = NodeStore::new();
432 let hash = store.get_node(1, 64, 2).hash();
433 store.get_node(1, 64, 2).cost_malus = 4.0;
434
435 assert_eq!(store.len(), 1);
436 assert_eq!(store.get(hash).map(|node| node.cost_malus), Some(4.0));
437 }
438
439 #[test]
440 fn node_store_resets_search_state_without_losing_path_type_cost() {
441 let mut store = NodeStore::new();
442 let hash = store.get_node(1, 64, 2).hash();
443 let Some(node) = store.get_mut(hash) else {
444 panic!("node should exist");
445 };
446 node.heap_idx = 3;
447 node.g = 1.0;
448 node.h = 2.0;
449 node.f = 3.0;
450 node.came_from = Some(Node::create_hash(0, 64, 2));
451 node.closed = true;
452 node.walked_distance = 4.0;
453 node.cost_malus = 8.0;
454 node.path_type = PathType::Water;
455
456 store.reset_search_state();
457
458 let Some(node) = store.get(hash) else {
459 panic!("node should still exist");
460 };
461 assert_eq!(node.heap_idx, -1);
462 assert_eq!(node.g.to_bits(), 0.0_f32.to_bits());
463 assert_eq!(node.h.to_bits(), 0.0_f32.to_bits());
464 assert_eq!(node.f.to_bits(), 0.0_f32.to_bits());
465 assert_eq!(node.came_from, None);
466 assert!(!node.closed);
467 assert_eq!(node.walked_distance.to_bits(), 0.0_f32.to_bits());
468 assert_eq!(node.cost_malus.to_bits(), 8.0_f32.to_bits());
469 assert_eq!(node.path_type, PathType::Water);
470 }
471
472 #[test]
473 fn node_heap_pops_lowest_f_cost_first() {
474 let mut store = NodeStore::new();
475 let high = node_with_cost(&mut store, 0, 64, 0, 5.0);
476 let low = node_with_cost(&mut store, 1, 64, 0, 1.0);
477 let middle = node_with_cost(&mut store, 2, 64, 0, 3.0);
478 let mut heap = NodeHeap::new();
479
480 assert!(heap.insert(&mut store, high));
481 assert!(heap.insert(&mut store, low));
482 assert!(heap.insert(&mut store, middle));
483
484 assert_eq!(heap.peek(), Some(low));
485 assert_eq!(heap.pop(&mut store), Some(low));
486 assert_eq!(store.get(low).map(|node| node.heap_idx), Some(-1));
487 assert_eq!(heap.pop(&mut store), Some(middle));
488 assert_eq!(heap.pop(&mut store), Some(high));
489 assert!(heap.is_empty());
490 }
491
492 #[test]
493 fn node_heap_change_cost_reorders_existing_node() {
494 let mut store = NodeStore::new();
495 let high = node_with_cost(&mut store, 0, 64, 0, 5.0);
496 let low = node_with_cost(&mut store, 1, 64, 0, 1.0);
497 let mut heap = NodeHeap::new();
498
499 assert!(heap.insert(&mut store, high));
500 assert!(heap.insert(&mut store, low));
501 assert_eq!(heap.peek(), Some(low));
502
503 assert!(heap.change_cost(&mut store, high, 0.5));
504
505 assert_eq!(heap.peek(), Some(high));
506 assert_eq!(heap.pop(&mut store), Some(high));
507 }
508
509 #[test]
510 fn target_tracks_best_node_hash() {
511 let from = Node::new(0, 64, 0);
512 let better = Node::new(1, 64, 0);
513 let worse = Node::new(4, 64, 0);
514 let mut target = Target::new(Node::new(2, 64, 0));
515
516 target.update_best(10.0, &worse);
517 target.update_best(1.0, &better);
518 target.update_best(5.0, &from);
519
520 assert_eq!(target.best_node(), Some(better.hash()));
521 assert!(!target.is_reached());
522 target.set_reached();
523 assert!(target.is_reached());
524 }
525
526 fn node_with_cost(store: &mut NodeStore, x: i32, y: i32, z: i32, cost: f32) -> i32 {
527 let node = store.get_node(x, y, z);
528 node.f = cost;
529 node.hash()
530 }
531}