add Event::ReceiveChunk and find_blocks_in_chunk function

2 files changed, 90 insertions, 69 deletions

diff --git a/azalea-client/src/plugins/events.rs b/azalea-client/src/plugins/events.rs
index d9cbf912..36f48a05 100644
--- a/azalea-client/src/plugins/events.rs
+++ b/azalea-client/src/plugins/events.rs
@@ -4,7 +4,7 @@
 use std::sync::Arc;
 
 use azalea_chat::FormattedText;
-use azalea_core::tick::GameTick;
+use azalea_core::{position::ChunkPos, tick::GameTick};
 use azalea_entity::{Dead, InLoadedChunk};
 use azalea_protocol::packets::game::c_player_combat_kill::ClientboundPlayerCombatKill;
 use azalea_world::{InstanceName, MinecraftEntityId};
@@ -15,6 +15,7 @@ use tokio::sync::mpsc;
 
 use crate::{
     chat::{ChatPacket, ChatReceivedEvent},
+    chunks::ReceiveChunkEvent,
     disconnect::DisconnectEvent,
     packet::game::{
         AddPlayerEvent, DeathEvent, KeepAliveEvent, RemovePlayerEvent, UpdatePlayerEvent,
@@ -118,6 +119,7 @@ pub enum Event {
     KeepAlive(u64),
     /// The client disconnected from the server.
     Disconnect(Option<FormattedText>),
+    ReceiveChunk(ChunkPos),
 }
 
 /// A component that contains an event sender for events that are only
@@ -294,3 +296,17 @@ pub fn disconnect_listener(
         }
     }
 }
+
+pub fn receive_chunk_listener(
+    query: Query<&LocalPlayerEvents>,
+    mut events: EventReader<ReceiveChunkEvent>,
+) {
+    for event in events.read() {
+        if let Ok(local_player_events) = query.get(event.entity) {
+            let _ = local_player_events.send(Event::ReceiveChunk(ChunkPos::new(
+                event.packet.x,
+                event.packet.z,
+            )));
+        }
+    }
+}
diff --git a/azalea-world/src/find_blocks.rs b/azalea-world/src/find_blocks.rs
index b266d799..10068243 100644
--- a/azalea-world/src/find_blocks.rs
+++ b/azalea-world/src/find_blocks.rs
@@ -1,16 +1,7 @@
 use azalea_block::{BlockState, BlockStates};
 use azalea_core::position::{BlockPos, ChunkPos};
 
-use crate::{ChunkStorage, Instance, iterators::ChunkIterator, palette::Palette};
-
-fn palette_maybe_has_block(palette: &Palette<BlockState>, block_states: &BlockStates) -> bool {
-    match &palette {
-        Palette::SingleValue(id) => block_states.contains(id),
-        Palette::Linear(ids) => ids.iter().any(|id| block_states.contains(id)),
-        Palette::Hashmap(ids) => ids.iter().any(|id| block_states.contains(id)),
-        Palette::Global => true,
-    }
-}
+use crate::{Chunk, ChunkStorage, Instance, iterators::ChunkIterator, palette::Palette};
 
 impl Instance {
     /// Find the coordinates of a block in the world.
@@ -52,35 +43,20 @@ impl Instance {
                 continue;
             };
 
-            for (section_index, section) in chunk.read().sections.iter().enumerate() {
-                let maybe_has_block =
-                    palette_maybe_has_block(&section.states.palette, block_states);
-                if !maybe_has_block {
-                    continue;
-                }
-
-                for i in 0..4096 {
-                    let block_state = section.states.get_at_index(i);
-
-                    if block_states.contains(&block_state) {
-                        let section_pos = section.states.coords_from_index(i);
-                        let (x, y, z) = (
-                            chunk_pos.x * 16 + (section_pos.x as i32),
-                            self.chunks.min_y + (section_index * 16) as i32 + section_pos.y as i32,
-                            chunk_pos.z * 16 + (section_pos.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_none()
-                            || this_block_distance < nearest_found_distance
-                        {
-                            nearest_found_pos = Some(this_block_pos);
-                            nearest_found_distance = this_block_distance;
-                        }
+            find_blocks_in_chunk(
+                block_states,
+                chunk_pos,
+                &chunk.read(),
+                self.chunks.min_y,
+                |this_block_pos| {
+                    let this_block_distance = (nearest_to - this_block_pos).length_manhattan();
+                    // only update if it's closer
+                    if nearest_found_pos.is_none() || this_block_distance < nearest_found_distance {
+                        nearest_found_pos = Some(this_block_pos);
+                        nearest_found_distance = this_block_distance;
                     }
-                }
-            }
+                },
+            );
 
             if let Some(nearest_found_pos) = nearest_found_pos {
                 // this is required because find_block searches chunk-by-chunk, which can cause
@@ -179,38 +155,22 @@ impl Iterator for FindBlocks<'_> {
                 continue;
             };
 
-            for (section_index, section) in chunk.read().sections.iter().enumerate() {
-                let maybe_has_block =
-                    palette_maybe_has_block(&section.states.palette, self.block_states);
-                if !maybe_has_block {
-                    continue;
-                }
+            find_blocks_in_chunk(
+                self.block_states,
+                chunk_pos,
+                &chunk.read(),
+                self.chunks.min_y,
+                |this_block_pos| {
+                    let this_block_distance = (self.nearest_to - this_block_pos).length_manhattan();
 
-                for i in 0..4096 {
-                    let block_state = section.states.get_at_index(i);
-
-                    if self.block_states.contains(&block_state) {
-                        let section_pos = section.states.coords_from_index(i);
-                        let (x, y, z) = (
-                            chunk_pos.x * 16 + (section_pos.x as i32),
-                            self.chunks.min_y + (section_index * 16) as i32 + section_pos.y as i32,
-                            chunk_pos.z * 16 + (section_pos.z as i32),
-                        );
-                        let this_block_pos = BlockPos { x, y, z };
-                        let this_block_distance =
-                            (self.nearest_to - this_block_pos).length_manhattan();
-
-                        found.push((this_block_pos, this_block_distance));
-
-                        if nearest_found_pos.is_none()
-                            || this_block_distance < nearest_found_distance
-                        {
-                            nearest_found_pos = Some(this_block_pos);
-                            nearest_found_distance = this_block_distance;
-                        }
+                    found.push((this_block_pos, this_block_distance));
+
+                    if nearest_found_pos.is_none() || this_block_distance < nearest_found_distance {
+                        nearest_found_pos = Some(this_block_pos);
+                        nearest_found_distance = this_block_distance;
                     }
-                }
-            }
+                },
+            );
 
             if let Some(nearest_found_pos) = nearest_found_pos {
                 // this is required because find_block searches chunk-by-chunk, which can cause
@@ -242,6 +202,51 @@ impl Iterator for FindBlocks<'_> {
     }
 }
 
+/// An optimized function for finding the block positions in a chunk that match
+/// the given block states.
+///
+/// This is used internally by [`Instance::find_block`] and
+/// [`Instance::find_blocks`].
+pub fn find_blocks_in_chunk(
+    block_states: &BlockStates,
+    chunk_pos: ChunkPos,
+    chunk: &Chunk,
+    min_y: i32,
+    mut cb: impl FnMut(BlockPos),
+) {
+    for (section_index, section) in chunk.sections.iter().enumerate() {
+        let maybe_has_block = palette_maybe_has_block(&section.states.palette, block_states);
+        if !maybe_has_block {
+            continue;
+        }
+
+        for i in 0..4096 {
+            let block_state = section.states.get_at_index(i);
+
+            if block_states.contains(&block_state) {
+                let section_pos = section.states.coords_from_index(i);
+                let (x, y, z) = (
+                    chunk_pos.x * 16 + (section_pos.x as i32),
+                    min_y + (section_index * 16) as i32 + section_pos.y as i32,
+                    chunk_pos.z * 16 + (section_pos.z as i32),
+                );
+                let this_block_pos = BlockPos { x, y, z };
+
+                cb(this_block_pos);
+            }
+        }
+    }
+}
+
+fn palette_maybe_has_block(palette: &Palette<BlockState>, block_states: &BlockStates) -> bool {
+    match &palette {
+        Palette::SingleValue(id) => block_states.contains(id),
+        Palette::Linear(ids) => ids.iter().any(|id| block_states.contains(id)),
+        Palette::Hashmap(ids) => ids.iter().any(|id| block_states.contains(id)),
+        Palette::Global => true,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use azalea_registry::Block;