From 074bb6198393e0ad87a2239f1c563bf41a1fbe59 Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Wed, 4 Oct 2023 22:54:07 -0500
Subject: [PATCH] improve pathfinder heuristics

---
 azalea-core/src/position.rs        |  1 +
 azalea/src/pathfinder/astar.rs     |  9 +++++--
 azalea/src/pathfinder/goals.rs     | 40 +++++++++++++++++++++++++++---
 azalea/src/pathfinder/moves/mod.rs | 10 ++++----
 4 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/azalea-core/src/position.rs b/azalea-core/src/position.rs
index e1993b18..f6bc4157 100755
--- a/azalea-core/src/position.rs
+++ b/azalea-core/src/position.rs
@@ -259,6 +259,7 @@ impl ChunkSectionPos {
         block >> 4
     }
 }
+
 /// The coordinates of a block inside a chunk.
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 pub struct ChunkBlockPos {
diff --git a/azalea/src/pathfinder/astar.rs b/azalea/src/pathfinder/astar.rs
index bc7b2309..4e1d1039 100644
--- a/azalea/src/pathfinder/astar.rs
+++ b/azalea/src/pathfinder/astar.rs
@@ -6,7 +6,7 @@ use std::{
     time::{Duration, Instant},
 };
 
-use log::{trace, warn};
+use log::{debug, trace, warn};
 use priority_queue::PriorityQueue;
 
 pub struct Path<P, M>
@@ -55,8 +55,12 @@ where
     let mut best_paths: [P; 7] = [start; 7];
     let mut best_path_scores: [f32; 7] = [heuristic(start); 7];
 
+    let mut num_nodes = 0;
+
     while let Some((current_node, _)) = open_set.pop() {
+        num_nodes += 1;
         if success(current_node) {
+            debug!("Nodes considered: {num_nodes}");
             return Path {
                 movements: reconstruct_path(nodes, current_node),
                 partial: false,
@@ -99,7 +103,8 @@ where
             }
         }
 
-        if start_time.elapsed() > timeout {
+        // check for timeout every ~1ms
+        if num_nodes % 1000 == 0 && start_time.elapsed() > timeout {
             // timeout, just return the best path we have so far
             trace!("A* couldn't find a path in time, returning best path");
             break;
diff --git a/azalea/src/pathfinder/goals.rs b/azalea/src/pathfinder/goals.rs
index 4e278c74..c8e2c211 100644
--- a/azalea/src/pathfinder/goals.rs
+++ b/azalea/src/pathfinder/goals.rs
@@ -1,6 +1,11 @@
+use std::f32::consts::SQRT_2;
+
 use azalea_core::position::{BlockPos, Vec3};
 
-use super::Goal;
+use super::{
+    costs::{FALL_ONE_BLOCK_COST, JUMP_ONE_BLOCK_COST},
+    Goal,
+};
 
 pub struct BlockPosGoal(pub BlockPos);
 impl Goal for BlockPosGoal {
@@ -8,13 +13,32 @@ impl Goal for BlockPosGoal {
         let dx = (self.0.x - n.x) as f32;
         let dy = (self.0.y - n.y) as f32;
         let dz = (self.0.z - n.z) as f32;
-        dx * dx + dy * dy + dz * dz
+
+        xz_heuristic(dx, dz) + y_heuristic(dy)
     }
     fn success(&self, n: BlockPos) -> bool {
         n == self.0
     }
 }
 
+fn xz_heuristic(dx: f32, dz: f32) -> f32 {
+    let x = dx.abs();
+    let z = dz.abs();
+
+    let diagonal;
+    let straight;
+
+    if x < z {
+        straight = z - x;
+        diagonal = x;
+    } else {
+        straight = x - z;
+        diagonal = z;
+    }
+
+    diagonal * SQRT_2 + straight
+}
+
 pub struct XZGoal {
     pub x: i32,
     pub z: i32,
@@ -23,20 +47,28 @@ impl Goal for XZGoal {
     fn heuristic(&self, n: BlockPos) -> f32 {
         let dx = (self.x - n.x) as f32;
         let dz = (self.z - n.z) as f32;
-        dx * dx + dz * dz
+        xz_heuristic(dx, dz)
     }
     fn success(&self, n: BlockPos) -> bool {
         n.x == self.x && n.z == self.z
     }
 }
 
+fn y_heuristic(dy: f32) -> f32 {
+    if dy < 0.0 {
+        FALL_ONE_BLOCK_COST * -dy
+    } else {
+        *JUMP_ONE_BLOCK_COST * dy
+    }
+}
+
 pub struct YGoal {
     pub y: i32,
 }
 impl Goal for YGoal {
     fn heuristic(&self, n: BlockPos) -> f32 {
         let dy = (self.y - n.y) as f32;
-        dy * dy
+        y_heuristic(dy)
     }
     fn success(&self, n: BlockPos) -> bool {
         n.y == self.y
diff --git a/azalea/src/pathfinder/moves/mod.rs b/azalea/src/pathfinder/moves/mod.rs
index 81a184aa..1714f8bc 100644
--- a/azalea/src/pathfinder/moves/mod.rs
+++ b/azalea/src/pathfinder/moves/mod.rs
@@ -25,6 +25,11 @@ type Edge = astar::Edge<BlockPos, MoveData>;
 
 pub type SuccessorsFn = fn(&mut Vec<Edge>, &PathfinderCtx, BlockPos);
 
+pub fn default_move(edges: &mut Vec<Edge>, ctx: &PathfinderCtx, node: BlockPos) {
+    basic::basic_move(edges, ctx, node);
+    parkour::parkour_move(edges, ctx, node);
+}
+
 #[derive(Clone)]
 pub struct MoveData {
     /// Use the context to determine what events should be sent to complete this
@@ -272,11 +277,6 @@ pub struct IsReachedCtx<'a> {
     pub physics: &'a azalea_entity::Physics,
 }
 
-pub fn default_move(edges: &mut Vec<Edge>, ctx: &PathfinderCtx, node: BlockPos) {
-    basic::basic_move(edges, ctx, node);
-    parkour::parkour_move(edges, ctx, node);
-}
-
 /// Returns whether the entity is at the node and should start going to the
 /// next node.
 #[must_use]