improve Instance::find_block

4 files changed, 79 insertions, 17 deletions

diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
index d831770c..f0f053fa 100755
--- a/azalea-world/src/chunk_storage.rs
+++ b/azalea-world/src/chunk_storage.rs
@@ -33,7 +33,7 @@ pub struct PartialChunkStorage {
 pub struct ChunkStorage {
     pub height: u32,
     pub min_y: i32,
-    pub chunks: HashMap<ChunkPos, Weak<RwLock<Chunk>>>,
+    pub map: HashMap<ChunkPos, Weak<RwLock<Chunk>>>,
 }
 
 /// A single chunk in a world (16*?*16 blocks). This only contains the blocks
@@ -188,7 +188,7 @@ impl PartialChunkStorage {
         chunk_storage: &mut ChunkStorage,
     ) {
         if let Some(chunk) = &chunk {
-            chunk_storage.chunks.insert(*pos, Arc::downgrade(chunk));
+            chunk_storage.map.insert(*pos, Arc::downgrade(chunk));
         } else {
             // don't remove it from the shared storage, since it'll be removed
             // automatically if this was the last reference
@@ -203,12 +203,12 @@ impl ChunkStorage {
         ChunkStorage {
             height,
             min_y,
-            chunks: HashMap::new(),
+            map: HashMap::new(),
         }
     }
 
     pub fn get(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
-        self.chunks.get(pos).and_then(|chunk| chunk.upgrade())
+        self.map.get(pos).and_then(|chunk| chunk.upgrade())
     }
 
     pub fn get_block_state(&self, pos: &BlockPos) -> Option<BlockState> {
diff --git a/azalea-world/src/iterators.rs b/azalea-world/src/iterators.rs
index 4054bfe0..c7290781 100644
--- a/azalea-world/src/iterators.rs
+++ b/azalea-world/src/iterators.rs
@@ -181,7 +181,7 @@ pub struct ChunkIterator {
     pub max_distance: u32,
     pub start: ChunkPos,
     pub pos: ChunkPos,
-    pub layer: i32,
+    pub layer: u32,
     pub leg: i32,
 }
 impl ChunkIterator {
@@ -235,7 +235,7 @@ impl Iterator for ChunkIterator {
                     self.pos.x += 1;
                     self.leg = 0;
                     self.layer += 1;
-                    if self.layer == self.max_distance as i32 {
+                    if self.layer == self.max_distance {
                         return None;
                     }
                 }
diff --git a/azalea-world/src/world.rs b/azalea-world/src/world.rs
index 8f4251bc..6e169939 100644
--- a/azalea-world/src/world.rs
+++ b/azalea-world/src/world.rs
@@ -106,7 +106,7 @@ impl Instance {
 
     /// Find the coordinates of a block in the world.
     ///
-    /// Note that this is sorted by `x+y+z` and not `x^2+y^2+z^2`, for
+    /// Note that this is sorted by `x+y+z` and not `x^2+y^2+z^2` for
     /// optimization purposes.
     ///
     /// ```
@@ -124,13 +124,21 @@ impl Instance {
 
         let nearest_to: BlockPos = nearest_to.into();
         let start_chunk: ChunkPos = (&nearest_to).into();
-        let iter = ChunkIterator::new(start_chunk, 32);
+        let mut iter = ChunkIterator::new(start_chunk, 32);
 
-        for chunk_pos in iter {
-            let chunk = self.chunks.get(&chunk_pos).unwrap();
+        let mut nearest_found_pos: Option<BlockPos> = None;
+        let mut nearest_found_distance = 0;
 
-            let mut nearest_found_pos: Option<BlockPos> = None;
-            let mut nearest_found_distance = 0;
+        // we do `while` instead of `for` so we can access iter later
+        while let Some(chunk_pos) = iter.next() {
+            let Some(chunk) = self.chunks.get(&chunk_pos) else {
+                // if the chunk isn't loaded then we skip it.
+                // we don't just return since it *could* cause issues if there's a random
+                // unloaded chunk and then more that are loaded.
+                // unlikely but still something to consider, and it's not like this slows it
+                // down much anyways.
+                continue;
+            };
 
             for (section_index, section) in chunk.read().sections.iter().enumerate() {
                 let maybe_has_block = match &section.states.palette {
@@ -171,13 +179,31 @@ impl Instance {
                 }
             }
 
-            // if we found the position, return it
-            if nearest_found_pos.is_some() {
-                return nearest_found_pos;
+            if let Some(nearest_found_pos) = nearest_found_pos {
+                // this is required because find_block searches chunk-by-chunk, which can cause
+                // us to find blocks first that aren't actually the closest
+                let required_chunk_distance = u32::max(
+                    u32::max(
+                        (chunk_pos.x - start_chunk.x).unsigned_abs(),
+                        (chunk_pos.z - start_chunk.z).unsigned_abs(),
+                    ),
+                    ((nearest_to.y - nearest_found_pos.y).unsigned_abs()).div_ceil(16),
+                );
+                let nearest_chunk_distance = iter.layer;
+
+                // if we found the position and there's no chance there's something closer,
+                // return it
+                if nearest_chunk_distance >= required_chunk_distance {
+                    return Some(nearest_found_pos);
+                }
             }
         }
 
-        None
+        if nearest_found_pos.is_some() {
+            nearest_found_pos
+        } else {
+            None
+        }
     }
 }
 
@@ -214,3 +240,39 @@ impl From<ChunkStorage> for Instance {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use azalea_registry::Block;
+
+    use crate::Chunk;
+
+    use super::*;
+
+    #[test]
+    fn find_block() {
+        let mut instance = Instance::default();
+
+        let chunk_storage = &mut instance.chunks;
+        let mut partial_chunk_storage = PartialChunkStorage::default();
+
+        // block at (17, 0, 0) and (0, 18, 0)
+
+        partial_chunk_storage.set(
+            &ChunkPos { x: 0, z: 0 },
+            Some(Chunk::default()),
+            chunk_storage,
+        );
+        partial_chunk_storage.set(
+            &ChunkPos { x: 1, z: 0 },
+            Some(Chunk::default()),
+            chunk_storage,
+        );
+
+        chunk_storage.set_block_state(&BlockPos { x: 17, y: 0, z: 0 }, Block::Stone.into());
+        chunk_storage.set_block_state(&BlockPos { x: 0, y: 18, z: 0 }, Block::Stone.into());
+
+        let pos = instance.find_block(BlockPos { x: 0, y: 0, z: 0 }, &Block::Stone.into());
+        assert_eq!(pos, Some(BlockPos { x: 17, y: 0, z: 0 }));
+    }
+}
diff --git a/azalea/examples/testbot.rs b/azalea/examples/testbot.rs
index 6fb50964..e2f1c436 100644
--- a/azalea/examples/testbot.rs
+++ b/azalea/examples/testbot.rs
@@ -306,7 +306,7 @@ async fn swarm_handle(
                 for (name, world) in &swarm.instance_container.read().worlds {
                     println!("world name: {name}");
                     if let Some(w) = world.upgrade() {
-                        for chunk_pos in w.read().chunks.chunks.values() {
+                        for chunk_pos in w.read().chunks.map.values() {
                             println!("chunk: {chunk_pos:?}");
                         }
                     } else {