Improvements to azalea-world for entities

Cargo.lock

@@ -198,6 +198,7 @@ dependencies = [
  "azalea-entity",
  "azalea-nbt",
  "azalea-protocol",
+ "log",
  "nohash-hasher",
 ]
 

azalea-client/src/client.rs

@@ -21,9 +21,11 @@ use azalea_protocol::{
     resolver, ServerAddress,
 };
 use azalea_world::{ChunkStorage, EntityStorage, World};
-use std::{fmt::Debug, sync::Arc};
+use std::{
+    fmt::Debug,
+    sync::{Arc, Mutex},
+};
 use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
-use tokio::sync::Mutex;
 
 #[derive(Default)]
 pub struct ClientState {
@@ -55,7 +57,7 @@ pub enum ChatPacket {
 /// A player that you can control that is currently in a Minecraft server.
 pub struct Client {
     event_receiver: UnboundedReceiver<Event>,
-    pub conn: Arc<Mutex<GameConnection>>,
+    pub conn: Arc<tokio::sync::Mutex<GameConnection>>,
     pub state: Arc<Mutex<ClientState>>,
     // game_loop
 }
@@ -63,6 +65,8 @@ pub struct Client {
 /// Whether we should ignore errors when decoding packets.
 const IGNORE_ERRORS: bool = !cfg!(debug_assertions);
 
+struct HandleError(String);
+
 impl Client {
     /// Connect to a Minecraft server with an account.
     pub async fn join(account: &Account, address: &ServerAddress) -> Result<Self, String> {
@@ -137,7 +141,7 @@ impl Client {
             }
         };
 
-        let conn = Arc::new(Mutex::new(conn));
+        let conn = Arc::new(tokio::sync::Mutex::new(conn));
 
         let (tx, rx) = mpsc::unbounded_channel();
 
@@ -161,14 +165,16 @@ impl Client {
     }
 
     async fn game_loop(
-        conn: Arc<Mutex<GameConnection>>,
+        conn: Arc<tokio::sync::Mutex<GameConnection>>,
         tx: UnboundedSender<Event>,
         state: Arc<Mutex<ClientState>>,
     ) {
         loop {
             let r = conn.lock().await.read().await;
             match r {
-                Ok(packet) => Self::handle(&packet, &tx, &state, &conn).await,
+                Ok(packet) => {
+                    Self::handle(&packet, &tx, &state, &conn).await;
+                }
                 Err(e) => {
                     if IGNORE_ERRORS {
                         println!("Error: {:?}", e);
@@ -187,13 +193,14 @@ impl Client {
         packet: &GamePacket,
         tx: &UnboundedSender<Event>,
         state: &Arc<Mutex<ClientState>>,
-        conn: &Arc<Mutex<GameConnection>>,
+        conn: &Arc<tokio::sync::Mutex<GameConnection>>,
-    ) {
+    ) -> Result<(), HandleError> {
         match packet {
             GamePacket::ClientboundLoginPacket(p) => {
                 println!("Got login packet {:?}", p);
 
-                let mut state = state.lock().await;
+                {
+                    let mut state = state.lock()?;
 
                     // // write p into login.txt
                     // std::io::Write::write_all(
@@ -257,12 +264,8 @@ impl Client {
                     .try_into()
                     .expect("min_y is not an i32");
 
-                state.world = Some(World {
+                    state.world = Some(World::new(16, height, min_y));
-                    height,
+                }
-                    min_y,
-                    storage: ChunkStorage::new(16),
-                    entities: EntityStorage::new(),
-                });
 
                 conn.lock()
                     .await
@@ -321,8 +324,7 @@ impl Client {
             GamePacket::ClientboundSetChunkCacheCenterPacket(p) => {
                 println!("Got chunk cache center packet {:?}", p);
                 state
-                    .lock()
+                    .lock()?
-                    .await
                     .world
                     .as_mut()
                     .unwrap()
@@ -334,8 +336,7 @@ impl Client {
                 // let chunk = Chunk::read_with_world_height(&mut p.chunk_data);
                 // println("chunk {:?}")
                 state
-                    .lock()
+                    .lock()?
-                    .await
                     .world
                     .as_mut()
                     .expect("World doesn't exist! We should've gotten a login packet by now.")
@@ -349,13 +350,11 @@ impl Client {
                 println!("Got add entity packet {:?}", p);
                 let entity = Entity::from(p);
                 state
-                    .lock()
+                    .lock()?
-                    .await
                     .world
                     .as_mut()
                     .expect("World doesn't exist! We should've gotten a login packet by now.")
-                    .entities
+                    .add_entity(entity);
-                    .insert(entity);
             }
             GamePacket::ClientboundSetEntityDataPacket(p) => {
                 // println!("Got set entity data packet {:?}", p);
@@ -392,6 +391,18 @@ impl Client {
             }
             GamePacket::ClientboundTeleportEntityPacket(p) => {
                 // println!("Got teleport entity packet {:?}", p);
+                let state_lock = state.lock()?;
+
+                // let entity = state_lock
+                //     .world
+                //     .unwrap()
+                //     .entity_by_id(p.id)
+                //     .ok_or("Teleporting entity that doesn't exist.".to_string())?;
+                // state_lock
+                //     .world
+                //     .as_mut()
+                //     .expect("World doesn't exist! We should've gotten a login packet by now.")
+                //     .move_entity(&mut entity, new_pos)
             }
             GamePacket::ClientboundUpdateAdvancementsPacket(p) => {
                 println!("Got update advancements packet {:?}", p);
@@ -457,9 +468,23 @@ impl Client {
             }
             _ => panic!("Unexpected packet {:?}", packet),
         }
+
+        Ok(())
     }
 
     pub async fn next(&mut self) -> Option<Event> {
         self.event_receiver.recv().await
     }
 }
+
+impl<T> From<std::sync::PoisonError<T>> for HandleError {
+    fn from(e: std::sync::PoisonError<T>) -> Self {
+        HandleError(e.to_string())
+    }
+}
+
+impl From<String> for HandleError {
+    fn from(e: String) -> Self {
+        HandleError(e)
+    }
+}

azalea-core/src/position.rs

@@ -27,7 +27,7 @@ impl Rem<i32> for BlockPos {
     }
 }
 
-#[derive(Clone, Copy, Debug, Default, PartialEq)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
 pub struct ChunkPos {
     pub x: i32,
     pub z: i32,
@@ -164,6 +164,12 @@ impl From<&EntityPos> for BlockPos {
     }
 }
 
+impl From<&EntityPos> for ChunkPos {
+    fn from(pos: &EntityPos) -> Self {
+        ChunkPos::from(&BlockPos::from(pos))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -181,4 +187,15 @@ mod tests {
         let chunk_block_pos = ChunkBlockPos::from(&block_pos);
         assert_eq!(chunk_block_pos, ChunkBlockPos::new(5, 78, 14));
     }
+
+    #[test]
+    fn test_from_entity_pos_to_chunk_pos() {
+        let entity_pos = EntityPos {
+            x: 33.5,
+            y: 77.0,
+            z: -19.6,
+        };
+        let chunk_pos = ChunkPos::from(&entity_pos);
+        assert_eq!(chunk_pos, ChunkPos::new(2, -2));
+    }
 }

azalea-entity/src/lib.rs

@@ -1,4 +1,4 @@
-use azalea_core::EntityPos;
+use azalea_core::{ChunkPos, EntityPos};
 #[cfg(feature = "protocol")]
 use azalea_protocol::packets::game::clientbound_add_entity_packet::ClientboundAddEntityPacket;
 use uuid::Uuid;
@@ -16,19 +16,16 @@ impl Entity {
         &self.pos
     }
 
-    pub fn set_pos(&mut self, pos: EntityPos) {
+    /// Sets the position of the entity. This doesn't update the cache in
-        // TODO: check if it moved to another chunk
+    /// azalea-world, and should only be used within azalea-world!
-        self.pos = pos;
+    pub fn unsafe_move(&mut self, new_pos: EntityPos) {
+        self.pos = new_pos;
     }
 }
 
 #[cfg(feature = "protocol")]
-impl From<&azalea_protocol::packets::game::clientbound_add_entity_packet::ClientboundAddEntityPacket>
+impl From<&ClientboundAddEntityPacket> for Entity {
-    for Entity
+    fn from(p: &ClientboundAddEntityPacket) -> Self {
-{
-    fn from(
-        p: &azalea_protocol::packets::game::clientbound_add_entity_packet::ClientboundAddEntityPacket,
-    ) -> Self {
         Self {
             id: p.id,
             uuid: p.uuid,

azalea-world/Cargo.toml

@@ -11,6 +11,7 @@ azalea-core = {path = "../azalea-core"}
 azalea-entity = {path = "../azalea-entity"}
 azalea-nbt = {path = "../azalea-nbt"}
 azalea-protocol = {path = "../azalea-protocol"}
+log = "0.4.17"
 nohash-hasher = "0.2.0"
 
 [profile.release]

azalea-world/src/chunk.rs

@@ -1,4 +1,3 @@
-use crate::bit_storage::BitStorage;
 use crate::palette::PalettedContainer;
 use crate::palette::PalettedContainerType;
 use crate::World;
@@ -13,10 +12,13 @@ use std::{
 
 const SECTION_HEIGHT: u32 = 16;
 
+#[derive(Debug)]
 pub struct ChunkStorage {
     pub view_center: ChunkPos,
     chunk_radius: u32,
     view_range: u32,
+    pub height: u32,
+    pub min_y: i32,
     // chunks is a list of size chunk_radius * chunk_radius
     chunks: Vec<Option<Arc<Mutex<Chunk>>>>,
 }
@@ -32,12 +34,14 @@ fn floor_mod(x: i32, y: u32) -> u32 {
 }
 
 impl ChunkStorage {
-    pub fn new(chunk_radius: u32) -> Self {
+    pub fn new(chunk_radius: u32, height: u32, min_y: i32) -> Self {
         let view_range = chunk_radius * 2 + 1;
         ChunkStorage {
             view_center: ChunkPos::new(0, 0),
             chunk_radius,
             view_range,
+            height,
+            min_y,
             chunks: vec![None; (view_range * view_range) as usize],
         }
     }
@@ -61,6 +65,29 @@ impl ChunkStorage {
             None => None,
         }
     }
+
+    pub fn replace_with_packet_data(
+        &mut self,
+        pos: &ChunkPos,
+        data: &mut impl Read,
+    ) -> Result<(), String> {
+        if !self.in_range(pos) {
+            println!(
+                "Ignoring chunk since it's not in the view range: {}, {}",
+                pos.x, pos.z
+            );
+            return Ok(());
+        }
+
+        let chunk = Arc::new(Mutex::new(Chunk::read_with_world_height(
+            data,
+            self.height,
+        )?));
+        println!("Loaded chunk {:?}", pos);
+        self[pos] = Some(chunk);
+
+        Ok(())
+    }
 }
 
 impl Index<&ChunkPos> for ChunkStorage {
@@ -84,7 +111,7 @@ pub struct Chunk {
 
 impl Chunk {
     pub fn read_with_world(buf: &mut impl Read, data: &World) -> Result<Self, String> {
-        Self::read_with_world_height(buf, data.height)
+        Self::read_with_world_height(buf, data.height())
     }
 
     pub fn read_with_world_height(buf: &mut impl Read, world_height: u32) -> Result<Self, String> {

azalea-world/src/entity.rs

@@ -1,14 +1,14 @@
-use std::collections::HashMap;
-
 use azalea_core::ChunkPos;
 use azalea_entity::Entity;
-use nohash_hasher::IntMap;
+use log::warn;
+use nohash_hasher::{IntMap, IntSet};
+use std::collections::HashMap;
 
 #[derive(Debug)]
 pub struct EntityStorage {
     by_id: IntMap<u32, Entity>,
     // TODO: this doesn't work yet (should be updated in the set_pos method in azalea-entity)
-    by_chunk: HashMap<ChunkPos, u32>,
+    by_chunk: HashMap<ChunkPos, IntSet<u32>>,
 }
 
 impl EntityStorage {
@@ -22,13 +22,24 @@ impl EntityStorage {
     /// Add an entity to the storage.
     #[inline]
     pub fn insert(&mut self, entity: Entity) {
+        self.by_chunk
+            .entry(ChunkPos::from(entity.pos()))
+            .or_default()
+            .insert(entity.id);
         self.by_id.insert(entity.id, entity);
     }
 
     /// Remove an entity from the storage by its id.
     #[inline]
     pub fn remove_by_id(&mut self, id: u32) {
-        self.by_id.remove(&id);
+        if let Some(entity) = self.by_id.remove(&id) {
+            let entity_chunk = ChunkPos::from(entity.pos());
+            if let None = self.by_chunk.remove(&entity_chunk) {
+                warn!("Tried to remove entity with id {id} from chunk {entity_chunk:?} but it was not found.");
+            }
+        } else {
+            warn!("Tried to remove entity with id {id} but it was not found.")
+        }
     }
 
     /// Get a reference to an entity by its id.
@@ -42,4 +53,30 @@ impl EntityStorage {
     pub fn get_mut_by_id(&mut self, id: u32) -> Option<&mut Entity> {
         self.by_id.get_mut(&id)
     }
+
+    /// Clear all entities in a chunk.
+    pub fn clear_chunk(&mut self, chunk: &ChunkPos) {
+        if let Some(entities) = self.by_chunk.remove(chunk) {
+            for entity_id in entities {
+                self.by_id.remove(&entity_id);
+            }
+        }
+    }
+
+    /// Updates an entity from its old chunk.
+    #[inline]
+    pub fn update_entity_chunk(
+        &mut self,
+        entity_id: u32,
+        old_chunk: &ChunkPos,
+        new_chunk: &ChunkPos,
+    ) {
+        if let Some(entities) = self.by_chunk.get_mut(old_chunk) {
+            entities.remove(&entity_id);
+        }
+        self.by_chunk
+            .entry(*new_chunk)
+            .or_default()
+            .insert(entity_id);
+    }
 }

azalea-world/src/lib.rs

@@ -6,7 +6,8 @@ mod entity;
 mod palette;
 
 use azalea_block::BlockState;
-use azalea_core::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos};
+use azalea_core::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos, EntityPos};
+use azalea_entity::Entity;
 use azalea_protocol::mc_buf::{McBufReadable, McBufWritable};
 pub use bit_storage::BitStorage;
 pub use chunk::{Chunk, ChunkStorage};
@@ -26,61 +27,78 @@ mod tests {
     }
 }
 
+#[derive(Debug)]
 pub struct World {
-    pub storage: ChunkStorage,
+    chunk_storage: ChunkStorage,
-    pub entities: EntityStorage,
+    entity_storage: EntityStorage,
-    pub height: u32,
-    pub min_y: i32,
 }
 
 impl World {
+    pub fn new(chunk_radius: u32, height: u32, min_y: i32) -> Self {
+        World {
+            chunk_storage: ChunkStorage::new(chunk_radius, height, min_y),
+            entity_storage: EntityStorage::new(),
+        }
+    }
+
     pub fn replace_with_packet_data(
         &mut self,
         pos: &ChunkPos,
         data: &mut impl Read,
     ) -> Result<(), String> {
-        if !self.storage.in_range(pos) {
+        self.chunk_storage.replace_with_packet_data(pos, data)
-            println!(
-                "Ignoring chunk since it's not in the view range: {}, {}",
-                pos.x, pos.z
-            );
-            return Ok(());
-        }
-        // let existing_chunk = &self.storage[pos];
-
-        let chunk = Arc::new(Mutex::new(Chunk::read_with_world(data, self)?));
-        println!("Loaded chunk {:?}", pos);
-        self.storage[pos] = Some(chunk);
-
-        Ok(())
     }
 
     pub fn update_view_center(&mut self, pos: &ChunkPos) {
-        self.storage.view_center = *pos;
+        self.chunk_storage.view_center = *pos;
     }
 
     pub fn get_block_state(&self, pos: &BlockPos) -> Option<BlockState> {
-        self.storage.get_block_state(pos, self.min_y)
+        self.chunk_storage.get_block_state(pos, self.min_y())
+    }
+
+    pub fn move_entity(&mut self, entity_id: u32, new_pos: EntityPos) -> Result<(), String> {
+        let entity = self
+            .entity_storage
+            .get_mut_by_id(entity_id)
+            .ok_or("Moving entity that doesn't exist".to_string())?;
+        let old_chunk = ChunkPos::from(entity.pos());
+        let new_chunk = ChunkPos::from(&new_pos);
+        // this is fine because we update the chunk below
+        entity.unsafe_move(new_pos);
+        if old_chunk != new_chunk {
+            self.entity_storage
+                .update_entity_chunk(entity_id, &old_chunk, &new_chunk);
+        }
+        Ok(())
+    }
+
+    pub fn add_entity(&mut self, entity: Entity) {
+        self.entity_storage.insert(entity);
+    }
+
+    pub fn height(&self) -> u32 {
+        self.chunk_storage.height
+    }
+
+    pub fn min_y(&self) -> i32 {
+        self.chunk_storage.min_y
+    }
+
+    pub fn entity_by_id(&self, id: u32) -> Option<&Entity> {
+        self.entity_storage.get_by_id(id)
     }
 }
+
 impl Index<&ChunkPos> for World {
     type Output = Option<Arc<Mutex<Chunk>>>;
 
     fn index(&self, pos: &ChunkPos) -> &Self::Output {
-        &self.storage[pos]
+        &self.chunk_storage[pos]
     }
 }
 impl IndexMut<&ChunkPos> for World {
     fn index_mut<'a>(&'a mut self, pos: &ChunkPos) -> &'a mut Self::Output {
-        &mut self.storage[pos]
+        &mut self.chunk_storage[pos]
     }
 }
-// impl Index<&BlockPos> for World {
-//     type Output = Option<Arc<Mutex<Chunk>>>;
-
-//     fn index(&self, pos: &BlockPos) -> &Self::Output {
-//         let chunk = &self[ChunkPos::from(pos)];
-//         // chunk.
-
-//     }
-// }

bot/src/main.rs

@@ -1,7 +1,7 @@
 use azalea_client::{Account, Event};
 
 #[tokio::main]
-async fn main() {
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
     println!("Hello, world!");
 
     // let address = "95.111.249.143:10000";
@@ -15,18 +15,18 @@ async fn main() {
     let mut client = account.join(&address.try_into().unwrap()).await.unwrap();
     println!("connected");
 
-    while let Some(e) = client.next().await {
+    while let Some(e) = &client.next().await {
         match e {
             // TODO: have a "loaded" or "ready" event that fires when all chunks are loaded
             Event::Login => {}
             Event::Chat(_p) => {
-                let state = client.state.lock().await;
+                let state = &client.state.lock()?;
                 let world = state.world.as_ref().unwrap();
-                println!("{:?}", world.entities);
+                println!("{:?}", world);
                 // world.get_block_state(state.player.entity.pos);
                 // println!("{}", p.message.to_ansi(None));
                 // if p.message.to_ansi(None) == "<py5> ok" {
-                //     let state = client.state.lock().await;
+                //     let state = client.state.lock();
                 //     let world = state.world.as_ref().unwrap();
                 //     let c = world.get_block_state(&BlockPos::new(5, 78, -2)).unwrap();
                 //     println!("block state: {:?}", c);
@@ -36,4 +36,6 @@ async fn main() {
     }
 
     println!("done");
+
+    Ok(())
 }

examples/pvp.rs

@@ -21,7 +21,7 @@ async fn main() {
                     if bot.entity.can_reach(target.bounding_box) {
                         bot.swing();
                     }
-                    if !h.using_held_item() && bot.state.lock().await.hunger <= 17 {
+                    if !h.using_held_item() && bot.state.lock().hunger <= 17 {
                         bot.hold(azalea::ItemGroup::Food);
                         tokio::task::spawn(bot.use_held_item());
                     }