Add World::find_block (#80)

* start adding World::find_block

* keep working on find_block

* BlockStates

* fix sorting

* update examples that use find_one_block

* azalea_block::properties

* fix tests

* add a gotoblock command to testbot

5 files changed, 403 insertions, 46 deletions

diff --git a/azalea-world/src/bit_storage.rs b/azalea-world/src/bit_storage.rs
index f6ca4cd6..09b68fae 100755
--- a/azalea-world/src/bit_storage.rs
+++ b/azalea-world/src/bit_storage.rs
@@ -158,13 +158,13 @@ impl BitStorage {
             .unwrap()
     }
 
+    /// Get the data at the given index.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the given index is greater than or equal to
+    /// the size of this storage.
     pub fn get(&self, index: usize) -> u64 {
-        // Validate.inclusiveBetween(0L, (long)(this.size - 1), (long)var1);
-        // int var2 = this.cellIndex(var1);
-        // long var3 = this.data[var2];
-        // int var5 = (var1 - var2 * this.valuesPerLong) * this.bits;
-        // return (int)(var3 >> var5 & this.mask);
-
         assert!(
             index < self.size,
             "Index {} out of bounds (must be less than {})",
diff --git a/azalea-world/src/iterators.rs b/azalea-world/src/iterators.rs
new file mode 100644
index 00000000..53a94898
--- /dev/null
+++ b/azalea-world/src/iterators.rs
@@ -0,0 +1,247 @@
+//! Iterators for iterating over Minecraft blocks and chunks, based on
+//! [prismarine-world's iterators](https://github.com/PrismarineJS/prismarine-world/blob/master/src/iterators.js).
+
+use azalea_core::{BlockPos, ChunkPos};
+
+/// An octahedron iterator, useful for iterating over blocks in a world.
+///
+/// ```
+/// # use azalea_core::BlockPos;
+/// # use azalea_world::iterators::BlockIterator;
+///
+/// let mut iter = BlockIterator::new(BlockPos::default(), 4);
+/// for block_pos in iter {
+///    println!("{:?}", block_pos);
+/// }
+/// ```
+pub struct BlockIterator {
+    start: BlockPos,
+    max_distance: u32,
+
+    pos: BlockPos,
+    apothem: u32,
+    left: i32,
+    right: i32,
+}
+impl BlockIterator {
+    pub fn new(start: BlockPos, max_distance: u32) -> Self {
+        Self {
+            start,
+            max_distance,
+
+            pos: BlockPos {
+                x: -1,
+                y: -1,
+                z: -1,
+            },
+            apothem: 1,
+            left: 1,
+            right: 2,
+        }
+    }
+}
+
+impl Iterator for BlockIterator {
+    type Item = BlockPos;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.apothem > self.max_distance {
+            return None;
+        }
+
+        self.right -= 1;
+        if self.right < 0 {
+            self.left -= 1;
+            if self.left < 0 {
+                self.pos.z += 2;
+                if self.pos.z > 1 {
+                    self.pos.y += 2;
+                    if self.pos.y > 1 {
+                        self.pos.x += 2;
+                        if self.pos.x > 1 {
+                            self.apothem += 1;
+                            self.pos.x = -1;
+                        }
+                        self.pos.y = -1;
+                    }
+                    self.pos.z = -1;
+                }
+                self.left = self.apothem as i32;
+            }
+            self.right = self.left;
+        }
+        let x = self.pos.x * self.right;
+        let y = self.pos.y * ((self.apothem as i32) - self.left);
+        let z = self.pos.z * ((self.apothem as i32) - (i32::abs(x) + i32::abs(y)));
+        Some(BlockPos { x: x, y, z } + self.start)
+    }
+}
+
+/// 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::SquareChunkIterator;
+///
+/// let mut iter = SquareChunkIterator::new(ChunkPos::default(), 4);
+/// for chunk_pos in iter {
+///   println!("{:?}", chunk_pos);
+/// }
+/// ```
+pub struct SquareChunkIterator {
+    start: ChunkPos,
+    number_of_points: u32,
+
+    dir: ChunkPos,
+
+    segment_len: u32,
+    pos: ChunkPos,
+    segment_passed: u32,
+    current_iter: u32,
+}
+impl SquareChunkIterator {
+    pub fn new(start: ChunkPos, max_distance: u32) -> Self {
+        Self {
+            start,
+            number_of_points: u32::pow(max_distance * 2 - 1, 2),
+
+            dir: ChunkPos { x: 1, z: 0 },
+
+            segment_len: 1,
+            pos: ChunkPos::default(),
+            segment_passed: 0,
+            current_iter: 0,
+        }
+    }
+
+    /// Change the distance that this iterator won't go past.
+    ///
+    /// ```
+    /// # use azalea_core::ChunkPos;
+    /// # use azalea_world::iterators::SquareChunkIterator;
+    ///
+    /// let mut iter = SquareChunkIterator::new(ChunkPos::default(), 2);
+    /// while let Some(chunk_pos) = iter.next() {
+    ///   println!("{:?}", chunk_pos);
+    /// }
+    /// iter.set_max_distance(4);
+    /// while let Some(chunk_pos) = iter.next() {
+    ///   println!("{:?}", chunk_pos);
+    /// }
+    /// ```
+    pub fn set_max_distance(&mut self, max_distance: u32) {
+        self.number_of_points = u32::pow(max_distance * 2 - 1, 2);
+    }
+}
+impl Iterator for SquareChunkIterator {
+    type Item = ChunkPos;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.current_iter > self.number_of_points {
+            return None;
+        }
+
+        let output = self.start + self.dir;
+
+        // make a step, add the direction to the current position
+        self.pos.x += self.dir.x;
+        self.pos.z += self.dir.z;
+        self.segment_passed += 1;
+
+        if self.segment_passed == self.segment_len {
+            // done with current segment
+            self.segment_passed = 0;
+
+            // rotate directions
+            (self.dir.x, self.dir.z) = (-self.dir.z, self.dir.x);
+
+            // increase segment length if necessary
+            if self.dir.z == 0 {
+                self.segment_len += 1;
+            }
+        }
+        self.current_iter += 1;
+        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/lib.rs b/azalea-world/src/lib.rs
index 1a419c3a..77498efd 100644
--- a/azalea-world/src/lib.rs
+++ b/azalea-world/src/lib.rs
@@ -7,6 +7,7 @@ mod bit_storage;
 mod chunk_storage;
 mod container;
 pub mod entity;
+pub mod iterators;
 pub mod palette;
 mod world;
 
diff --git a/azalea-world/src/palette.rs b/azalea-world/src/palette.rs
index d97f61a3..d10357ad 100755
--- a/azalea-world/src/palette.rs
+++ b/azalea-world/src/palette.rs
@@ -12,6 +12,11 @@ pub enum PalettedContainerType {
 #[derive(Clone, Debug)]
 pub struct PalettedContainer {
     pub bits_per_entry: u8,
+    /// This is usually a list of unique values that appear in the container so
+    /// they can be indexed by the bit storage.
+    ///
+    /// Sometimes it doesn't contain anything if there's too many unique items
+    /// in the bit storage, though.
     pub palette: Palette,
     /// Compacted list of indices pointing to entry IDs in the Palette.
     pub storage: BitStorage,
@@ -37,7 +42,7 @@ impl PalettedContainer {
         container_type: &'static PalettedContainerType,
     ) -> Result<Self, BufReadError> {
         let bits_per_entry = u8::read_from(buf)?;
-        let palette_type = PaletteType::from_bits_and_type(bits_per_entry, container_type);
+        let palette_type = PaletteKind::from_bits_and_type(bits_per_entry, container_type);
         let palette = palette_type.read(buf)?;
         let size = container_type.size();
 
@@ -57,15 +62,33 @@ impl PalettedContainer {
     }
 
     /// Calculates the index of the given coordinates.
-    pub fn get_index(&self, x: usize, y: usize, z: usize) -> usize {
+    pub fn index_from_coords(&self, x: usize, y: usize, z: usize) -> usize {
         let size_bits = self.container_type.size_bits();
 
         (((y << size_bits) | z) << size_bits) | x
     }
 
+    pub fn coords_from_index(&self, index: usize) -> (usize, usize, usize) {
+        let size_bits = self.container_type.size_bits();
+        let mask = (1 << size_bits) - 1;
+        (
+            index & mask,
+            (index >> size_bits >> size_bits) & mask,
+            (index >> size_bits) & mask,
+        )
+    }
+
     /// Returns the value at the given index.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if the index is greater than or equal to the number
+    /// of things in the storage. (So for block states, it must be less than
+    /// 4096).
     pub fn get_at_index(&self, index: usize) -> u32 {
+        // first get the pallete id
         let paletted_value = self.storage.get(index);
+        // and then get the value from that id
         self.palette.value_for(paletted_value as usize)
     }
 
@@ -73,14 +96,14 @@ impl PalettedContainer {
     pub fn get(&self, x: usize, y: usize, z: usize) -> u32 {
         // let paletted_value = self.storage.get(self.get_index(x, y, z));
         // self.palette.value_for(paletted_value as usize)
-        self.get_at_index(self.get_index(x, y, z))
+        self.get_at_index(self.index_from_coords(x, y, z))
     }
 
     /// Sets the id at the given coordinates and return the previous id
     pub fn get_and_set(&mut self, x: usize, y: usize, z: usize, value: u32) -> u32 {
         let paletted_value = self.id_for(value);
         self.storage
-            .get_and_set(self.get_index(x, y, z), paletted_value as u64) as u32
+            .get_and_set(self.index_from_coords(x, y, z), paletted_value as u64) as u32
     }
 
     /// Sets the id at the given index and return the previous id. You probably
@@ -92,12 +115,12 @@ impl PalettedContainer {
 
     /// Sets the id at the given coordinates and return the previous id
     pub fn set(&mut self, x: usize, y: usize, z: usize, value: u32) {
-        self.set_at_index(self.get_index(x, y, z), value);
+        self.set_at_index(self.index_from_coords(x, y, z), value);
     }
 
     fn create_or_reuse_data(&self, bits_per_entry: u8) -> PalettedContainer {
         let new_palette_type =
-            PaletteType::from_bits_and_type(bits_per_entry, &self.container_type);
+            PaletteKind::from_bits_and_type(bits_per_entry, &self.container_type);
         // note for whoever is trying to optimize this: vanilla has this
         // but it causes a stack overflow since it's not changing the bits per entry
         // i don't know how to fix this properly so glhf
@@ -188,13 +211,14 @@ impl McBufWritable for PalettedContainer {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub enum PaletteType {
+pub enum PaletteKind {
     SingleValue,
     Linear,
     Hashmap,
     Global,
 }
 
+/// A representation of the different types of chunk palettes Minecraft uses.
 #[derive(Clone, Debug)]
 pub enum Palette {
     /// ID of the corresponding entry in its global palette
@@ -211,13 +235,7 @@ impl Palette {
         match self {
             Palette::SingleValue(v) => *v,
             Palette::Linear(v) => v[id],
-            Palette::Hashmap(v) => {
-                if id >= v.len() {
-                    0
-                } else {
-                    v[id]
-                }
-            }
+            Palette::Hashmap(v) => v.get(id).copied().unwrap_or_default(),
             Palette::Global => id as u32,
         }
     }
@@ -241,49 +259,49 @@ impl McBufWritable for Palette {
     }
 }
 
-impl PaletteType {
+impl PaletteKind {
     pub fn from_bits_and_type(bits_per_entry: u8, container_type: &PalettedContainerType) -> Self {
         match container_type {
             PalettedContainerType::BlockStates => match bits_per_entry {
-                0 => PaletteType::SingleValue,
-                1..=4 => PaletteType::Linear,
-                5..=8 => PaletteType::Hashmap,
-                _ => PaletteType::Global,
+                0 => PaletteKind::SingleValue,
+                1..=4 => PaletteKind::Linear,
+                5..=8 => PaletteKind::Hashmap,
+                _ => PaletteKind::Global,
             },
             PalettedContainerType::Biomes => match bits_per_entry {
-                0 => PaletteType::SingleValue,
-                1..=3 => PaletteType::Linear,
-                _ => PaletteType::Global,
+                0 => PaletteKind::SingleValue,
+                1..=3 => PaletteKind::Linear,
+                _ => PaletteKind::Global,
             },
         }
     }
 
     pub fn read(&self, buf: &mut Cursor<&[u8]>) -> Result<Palette, BufReadError> {
         Ok(match self {
-            PaletteType::SingleValue => Palette::SingleValue(u32::var_read_from(buf)?),
-            PaletteType::Linear => Palette::Linear(Vec::<u32>::var_read_from(buf)?),
-            PaletteType::Hashmap => Palette::Hashmap(Vec::<u32>::var_read_from(buf)?),
-            PaletteType::Global => Palette::Global,
+            PaletteKind::SingleValue => Palette::SingleValue(u32::var_read_from(buf)?),
+            PaletteKind::Linear => Palette::Linear(Vec::<u32>::var_read_from(buf)?),
+            PaletteKind::Hashmap => Palette::Hashmap(Vec::<u32>::var_read_from(buf)?),
+            PaletteKind::Global => Palette::Global,
         })
     }
 
     pub fn as_empty_palette(&self) -> Palette {
         match self {
-            PaletteType::SingleValue => Palette::SingleValue(0),
-            PaletteType::Linear => Palette::Linear(Vec::new()),
-            PaletteType::Hashmap => Palette::Hashmap(Vec::new()),
-            PaletteType::Global => Palette::Global,
+            PaletteKind::SingleValue => Palette::SingleValue(0),
+            PaletteKind::Linear => Palette::Linear(Vec::new()),
+            PaletteKind::Hashmap => Palette::Hashmap(Vec::new()),
+            PaletteKind::Global => Palette::Global,
         }
     }
 }
 
-impl From<&Palette> for PaletteType {
+impl From<&Palette> for PaletteKind {
     fn from(palette: &Palette) -> Self {
         match palette {
-            Palette::SingleValue(_) => PaletteType::SingleValue,
-            Palette::Linear(_) => PaletteType::Linear,
-            Palette::Hashmap(_) => PaletteType::Hashmap,
-            Palette::Global => PaletteType::Global,
+            Palette::SingleValue(_) => PaletteKind::SingleValue,
+            Palette::Linear(_) => PaletteKind::Linear,
+            Palette::Hashmap(_) => PaletteKind::Hashmap,
+            Palette::Global => PaletteKind::Global,
         }
     }
 }
@@ -313,14 +331,14 @@ mod tests {
         assert_eq!(palette_container.bits_per_entry, 0);
         assert_eq!(palette_container.get_at_index(0), 0);
         assert_eq!(
-            PaletteType::from(&palette_container.palette),
-            PaletteType::SingleValue
+            PaletteKind::from(&palette_container.palette),
+            PaletteKind::SingleValue
         );
         palette_container.set_at_index(0, 1);
         assert_eq!(palette_container.get_at_index(0), 1);
         assert_eq!(
-            PaletteType::from(&palette_container.palette),
-            PaletteType::Linear
+            PaletteKind::from(&palette_container.palette),
+            PaletteKind::Linear
         );
     }
 
@@ -359,4 +377,22 @@ mod tests {
         palette_container.set_at_index(16, 16); // 5 bits
         assert_eq!(palette_container.bits_per_entry, 5);
     }
+
+    #[test]
+    fn test_coords_from_index() {
+        let palette_container =
+            PalettedContainer::new(&PalettedContainerType::BlockStates).unwrap();
+
+        for x in 0..15 {
+            for y in 0..15 {
+                for z in 0..15 {
+                    assert_eq!(
+                        palette_container
+                            .coords_from_index(palette_container.index_from_coords(x, y, z)),
+                        (x, y, z)
+                    );
+                }
+            }
+        }
+    }
 }
diff --git a/azalea-world/src/world.rs b/azalea-world/src/world.rs
index 41d83082..5bb9b0b7 100644
--- a/azalea-world/src/world.rs
+++ b/azalea-world/src/world.rs
@@ -2,9 +2,12 @@ use crate::{
     entity::{
         EntityInfos, EntityUuid, LoadedBy, Local, MinecraftEntityId, PartialEntityInfos, WorldName,
     },
+    iterators::ChunkIterator,
+    palette::Palette,
     ChunkStorage, PartialChunkStorage, WorldContainer,
 };
-use azalea_core::ChunkPos;
+use azalea_block::{BlockState, BlockStates};
+use azalea_core::{BlockPos, ChunkPos};
 use bevy_ecs::{
     entity::Entity,
     query::{Changed, With, Without},
@@ -187,6 +190,76 @@ impl Instance {
     pub fn entity_by_id(&self, entity_id: &MinecraftEntityId) -> Option<Entity> {
         self.entity_by_id.get(entity_id).copied()
     }
+
+    /// 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
+    /// optimization purposes.
+    pub fn find_block(
+        &self,
+        nearest_to: impl Into<BlockPos>,
+        block_states: &BlockStates,
+    ) -> Option<BlockPos> {
+        // iterate over every chunk in a 3d spiral pattern
+        // and then check the palette for the block state
+
+        let nearest_to: BlockPos = nearest_to.into();
+        let start_chunk: ChunkPos = (&nearest_to).into();
+        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 {
+                    Palette::SingleValue(id) => block_states.contains(&BlockState { id: *id }),
+                    Palette::Linear(ids) => ids
+                        .iter()
+                        .any(|&id| block_states.contains(&BlockState { id })),
+                    Palette::Hashmap(ids) => ids
+                        .iter()
+                        .any(|&id| block_states.contains(&BlockState { id })),
+                    Palette::Global => true,
+                };
+                if !maybe_has_block {
+                    continue;
+                }
+
+                for i in 0..4096 {
+                    let block_state = section.states.get_at_index(i);
+                    let block_state = BlockState { id: block_state };
+
+                    if block_states.contains(&block_state) {
+                        let (section_x, section_y, section_z) = section.states.coords_from_index(i);
+                        let (x, y, z) = (
+                            chunk_pos.x * 16 + (section_x as i32),
+                            self.chunks.min_y + (section_index * 16) as i32 + section_y as i32,
+                            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 !nearest_found_pos.is_some()
+                            || this_block_distance < nearest_found_distance
+                        {
+                            nearest_found_pos = Some(this_block_pos);
+                            nearest_found_distance = this_block_distance;
+                        }
+                    }
+                }
+            }
+
+            // if we found the position, return it
+            if nearest_found_pos.is_some() {
+                return nearest_found_pos;
+            }
+        }
+
+        None
+    }
 }
 
 impl Debug for PartialWorld {