use radix heap in pathfinder for 10% speedup

3 files changed, 34 insertions, 3 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 9fae6684..09a8f8d4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -258,6 +258,7 @@ dependencies = [
  "num-format",
  "num-traits",
  "parking_lot",
+ "radix-heap",
  "rand 0.10.0-rc.5",
  "rustc-hash",
  "serde",
@@ -2706,6 +2707,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
 
 [[package]]
+name = "radix-heap"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59ffec9df464013295b499298811e6a3de31bf8128092135826517db12dee601"
+
+[[package]]
 name = "rand"
 version = "0.8.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/azalea/Cargo.toml b/azalea/Cargo.toml
index f82a344d..9a038a8e 100644
--- a/azalea/Cargo.toml
+++ b/azalea/Cargo.toml
@@ -36,6 +36,7 @@ nohash-hasher.workspace = true
 num-format.workspace = true
 num-traits.workspace = true
 parking_lot.workspace = true
+radix-heap = "0.4.2"
 rustc-hash.workspace = true
 serde = { workspace = true, optional = true }
 tokio.workspace = true
diff --git a/azalea/src/pathfinder/astar.rs b/azalea/src/pathfinder/astar.rs
index d5dc86bd..719036f0 100644
--- a/azalea/src/pathfinder/astar.rs
+++ b/azalea/src/pathfinder/astar.rs
@@ -1,5 +1,5 @@
 use std::{
-    cmp::{self},
+    cmp::{self, Reverse},
     collections::BinaryHeap,
     fmt::{self, Debug},
     hash::{BuildHasherDefault, Hash},
@@ -8,6 +8,7 @@ use std::{
 
 use indexmap::IndexMap;
 use num_format::ToFormattedString;
+use radix_heap::RadixHeapMap;
 use rustc_hash::FxHasher;
 use tracing::{debug, trace, warn};
 
@@ -285,6 +286,11 @@ impl<P: Hash + Copy + Clone, M: Clone> Clone for Movement<P, M> {
 #[derive(Default)]
 struct PathfinderHeap {
     binary_heap: BinaryHeap<WeightedNode>,
+    /// Key is f_score.to_bits(), value is (g_score, index)
+    ///
+    /// As long as the f_score is positive, comparing it as bits is fine. Also,
+    /// it has to be `Reverse`d to make it a min-heap.
+    radix_heap: RadixHeapMap<Reverse<u32>, (f32, usize)>,
 }
 impl PathfinderHeap {
     pub fn new() -> Self {
@@ -292,10 +298,27 @@ impl PathfinderHeap {
     }
 
     pub fn push(&mut self, item: WeightedNode) {
-        self.binary_heap.push(item);
+        if let Some(top) = self.radix_heap.top() {
+            // this can happen when the heuristic isn't optimal, so just fall back to a
+            // binary heap in those cases
+            if item.f_score < f32::from_bits(top.0) {
+                self.binary_heap.push(item);
+                return;
+            }
+        }
+        self.radix_heap
+            .push(Reverse(item.f_score.to_bits()), (item.g_score, item.index))
     }
     pub fn pop(&mut self) -> Option<WeightedNode> {
-        self.binary_heap.pop()
+        self.binary_heap.pop().or_else(|| {
+            self.radix_heap
+                .pop()
+                .map(|(f_score, (g_score, index))| WeightedNode {
+                    f_score: f32::from_bits(f_score.0),
+                    g_score,
+                    index,
+                })
+        })
     }
 }