improve Instance::find_block

1 files changed, 72 insertions, 10 deletions

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 }));
+    }
+}