fix sorting

azalea-rs · mat-1 · Mar 8, 2023 · Feb 27, 2023 · Mar 7, 2023 · Mar 8, 2023
commit 32ca2429f309bdde944539c0e3f30d52e9ecb17e
diff --git a/azalea-core/src/position.rs b/azalea-core/src/position.rs
@@ -12,6 +12,8 @@ macro_rules! vec3_impl {
                 Self { x, y, z }
             }
 
+            /// Get the distance of this vector to the origin by doing `x^2 + y^2 +
+            /// z^2`.
             pub fn length_sqr(&self) -> $type {
                 self.x * self.x + self.y * self.y + self.z * self.z
             }
@@ -139,6 +141,11 @@ impl BlockPos {
             z: self.z as f64 + 0.5,
         }
     }
+
+    /// Get the distance of this vector from the origin by doing `x + y + z`.
+    pub fn length_manhattan(&self) -> u32 {
+        (self.x.abs() + self.y.abs() + self.z.abs()) as u32
+    }
 }
 
 /// Chunk coordinates are used to represent where a chunk is in the world. You

diff --git a/azalea-world/src/iterators.rs b/azalea-world/src/iterators.rs
@@ -77,18 +77,19 @@ impl Iterator for BlockIterator {
     }
 }
 
-/// A spiral iterator, useful for iterating over chunks in a world.
+/// A spiral iterator, useful for iterating over chunks in a world. Use
+/// `ChunkIterator` to sort by x+y+z (Manhattan) distance.
 ///
 /// ```
 /// # use azalea_core::ChunkPos;
-/// # use azalea_world::iterators::ChunkIterator;
+/// # use azalea_world::iterators::SquareChunkIterator;
 ///
-/// let mut iter = ChunkIterator::new(ChunkPos::default(), 4);
+/// let mut iter = SquareChunkIterator::new(ChunkPos::default(), 4);
 /// for chunk_pos in iter {
 ///   println!("{:?}", chunk_pos);
 /// }
 /// ```
-pub struct ChunkIterator {
+pub struct SquareChunkIterator {
     start: ChunkPos,
     number_of_points: u32,
 
@@ -99,7 +100,7 @@ pub struct ChunkIterator {
     segment_passed: u32,
     current_iter: u32,
 }
-impl ChunkIterator {
+impl SquareChunkIterator {
     pub fn new(start: ChunkPos, max_distance: u32) -> Self {
         Self {
             start,
@@ -133,7 +134,7 @@ impl ChunkIterator {
         self.number_of_points = u32::pow(max_distance * 2 - 1, 2);
     }
 }
-impl Iterator for ChunkIterator {
+impl Iterator for SquareChunkIterator {
     type Item = ChunkPos;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -164,3 +165,83 @@ impl Iterator for ChunkIterator {
         Some(output)
     }
 }
+
+/// A diagonal spiral iterator, useful for iterating over chunks in a world.
+///
+/// ```
+/// # use azalea_core::ChunkPos;
+/// # use azalea_world::iterators::ChunkIterator;
+///
+/// let mut iter = ChunkIterator::new(ChunkPos::default(), 4);
+/// for chunk_pos in iter {
+///   println!("{:?}", chunk_pos);
+/// }
+/// ```
+pub struct ChunkIterator {
+    pub max_distance: u32,
+    pub start: ChunkPos,
+    pub pos: ChunkPos,
+    pub layer: i32,
+    pub leg: i32,
+}
+impl ChunkIterator {
+    pub fn new(start: ChunkPos, max_distance: u32) -> Self {
+        Self {
+            max_distance,
+            start,
+            pos: ChunkPos { x: 2, z: -1 },
+            layer: 1,
+            leg: -1,
+        }
+    }
+}
+impl Iterator for ChunkIterator {
+    type Item = ChunkPos;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.leg {
+            -1 => {
+                self.leg = 0;
+                return Some(self.start);
+            }
+            0 => {
+                if self.max_distance == 1 {
+                    return None;
+                }
+                self.pos.x -= 1;
+                self.pos.z += 1;
+                if self.pos.x == 0 {
+                    self.leg = 1;
+                }
+            }
+            1 => {
+                self.pos.x -= 1;
+                self.pos.z -= 1;
+                if self.pos.z == 0 {
+                    self.leg = 2;
+                }
+            }
+            2 => {
+                self.pos.x += 1;
+                self.pos.z -= 1;
+                if self.pos.x == 0 {
+                    self.leg = 3;
+                }
+            }
+            3 => {
+                self.pos.x += 1;
+                self.pos.z += 1;
+                if self.pos.z == 0 {
+                    self.pos.x += 1;
+                    self.leg = 0;
+                    self.layer += 1;
+                    if self.layer == self.max_distance as i32 {
+                        return None;
+                    }
+                }
+            }
+            _ => unreachable!(),
+        }
+        Some(self.start + self.pos)
+    }
+}
diff --git a/azalea-world/src/world.rs b/azalea-world/src/world.rs
@@ -205,15 +205,13 @@ impl World {
 
         let nearest_to: BlockPos = nearest_to.into();
         let start_chunk: ChunkPos = (&nearest_to).into();
-        // todo (correctness): rename this to something like SquareChunkIterator and
-        // also have another one that iterates in a diagonal shape and use that
-        // here
         let 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;
 
             for (section_index, section) in chunk.read().sections.iter().enumerate() {
                 let maybe_has_block = match &section.states.palette {
@@ -242,15 +240,14 @@ impl World {
                             chunk_pos.z * 16 + (section_z as i32),
                         );
                         let this_block_pos = BlockPos { x, y, z };
+                        let this_block_distance = (nearest_to - this_block_pos).length_manhattan();
                         // only update if it's closer
-                        if let Some(nearest_found_pos) = nearest_found_pos {
-                            if this_block_pos.x + this_block_pos.y + this_block_pos.z
-                                >= nearest_found_pos.x + nearest_found_pos.y + nearest_found_pos.z
-                            {
-                                continue;
-                            }
+                        if !nearest_found_pos.is_some()
+                            || this_block_distance < nearest_found_distance
+                        {
+                            nearest_found_pos = Some(this_block_pos);
+                            nearest_found_distance = this_block_distance;
                         }
-                        nearest_found_pos = Some(this_block_pos);
                     }
                 }
             }