make entities have a reference to WeakWorlds instead

... and other exciting bug fixes

4 files changed, 223 insertions, 154 deletions

diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
index 7d4486e4..6d3907f0 100755
--- a/azalea-world/src/chunk_storage.rs
+++ b/azalea-world/src/chunk_storage.rs
@@ -1,20 +1,16 @@
 use crate::palette::PalettedContainer;
 use crate::palette::PalettedContainerType;
-use crate::World;
 use azalea_block::BlockState;
-use azalea_buf::BufReadError;
-use azalea_buf::{McBufReadable, McBufWritable};
+use azalea_buf::{BufReadError, McBufReadable, McBufWritable};
 use azalea_core::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos};
-use log::debug;
-use log::trace;
-use log::warn;
-use parking_lot::Mutex;
+use log::{debug, trace, warn};
 use parking_lot::RwLock;
-use std::collections::HashMap;
-use std::fmt::Debug;
-use std::io::Cursor;
-use std::sync::Weak;
-use std::{io::Write, sync::Arc};
+use std::{
+    collections::HashMap,
+    fmt::Debug,
+    io::{Cursor, Write},
+    sync::{Arc, Weak},
+};
 
 const SECTION_HEIGHT: u32 = 16;
 
@@ -30,7 +26,7 @@ pub struct PartialChunkStorage {
     chunk_radius: u32,
     view_range: u32,
     // chunks is a list of size chunk_radius * chunk_radius
-    chunks: Vec<Option<Arc<Mutex<Chunk>>>>,
+    chunks: Vec<Option<Arc<RwLock<Chunk>>>>,
 }
 
 /// A storage for chunks where they're only stored weakly, so if they're not
@@ -39,7 +35,7 @@ pub struct PartialChunkStorage {
 pub struct WeakChunkStorage {
     pub height: u32,
     pub min_y: i32,
-    pub chunks: HashMap<ChunkPos, Weak<Mutex<Chunk>>>,
+    pub chunks: HashMap<ChunkPos, Weak<RwLock<Chunk>>>,
 }
 
 /// A storage of potentially infinite chunks in a world. Chunks are stored as
@@ -47,7 +43,7 @@ pub struct WeakChunkStorage {
 pub struct ChunkStorage {
     pub height: u32,
     pub min_y: i32,
-    pub chunks: HashMap<ChunkPos, Arc<Mutex<Chunk>>>,
+    pub chunks: HashMap<ChunkPos, Arc<RwLock<Chunk>>>,
 }
 
 /// A single chunk in a world (16*?*16 blocks). This only contains the blocks
@@ -113,20 +109,13 @@ impl PartialChunkStorage {
             && (chunk_pos.z - self.view_center.z).unsigned_abs() <= self.chunk_radius
     }
 
-    pub fn get_block_state(&self, pos: &BlockPos) -> Option<BlockState> {
-        let chunk_pos = ChunkPos::from(pos);
-        let chunk = self.get(&chunk_pos)?;
-        let chunk = chunk.lock();
-        chunk.get(&ChunkBlockPos::from(pos), self.min_y())
-    }
-
     pub fn set_block_state(&self, pos: &BlockPos, state: BlockState) -> Option<BlockState> {
         if pos.y < self.min_y() || pos.y >= (self.min_y() + self.height() as i32) {
             return None;
         }
         let chunk_pos = ChunkPos::from(pos);
         let chunk = self.get(&chunk_pos)?;
-        let mut chunk = chunk.lock();
+        let mut chunk = chunk.write();
         Some(chunk.get_and_set(&ChunkBlockPos::from(pos), state, self.min_y()))
     }
 
@@ -145,7 +134,7 @@ impl PartialChunkStorage {
             return Ok(());
         }
 
-        let chunk = Arc::new(Mutex::new(Chunk::read_with_dimension_height(
+        let chunk = Arc::new(RwLock::new(Chunk::read_with_dimension_height(
             data,
             self.height(),
         )?));
@@ -158,7 +147,7 @@ impl PartialChunkStorage {
 
     /// Get a [`Chunk`] within render distance, or `None` if it's not loaded.
     /// Use [`PartialChunkStorage::get`] to get a chunk from the shared storage.
-    pub fn limited_get(&self, pos: &ChunkPos) -> Option<&Arc<Mutex<Chunk>>> {
+    pub fn limited_get(&self, pos: &ChunkPos) -> Option<&Arc<RwLock<Chunk>>> {
         if !self.in_range(pos) {
             warn!(
                 "Chunk at {:?} is not in the render distance (center: {:?}, {} chunks)",
@@ -173,7 +162,7 @@ impl PartialChunkStorage {
     /// Get a mutable reference to a [`Chunk`] within render distance, or
     /// `None` if it's not loaded. Use [`PartialChunkStorage::get`] to get
     /// a chunk from the shared storage.
-    pub fn limited_get_mut(&mut self, pos: &ChunkPos) -> Option<&mut Option<Arc<Mutex<Chunk>>>> {
+    pub fn limited_get_mut(&mut self, pos: &ChunkPos) -> Option<&mut Option<Arc<RwLock<Chunk>>>> {
         if !self.in_range(pos) {
             return None;
         }
@@ -183,7 +172,7 @@ impl PartialChunkStorage {
     }
 
     /// Get a chunk,
-    pub fn get(&self, pos: &ChunkPos) -> Option<Arc<Mutex<Chunk>>> {
+    pub fn get(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
         self.shared
             .read()
             .chunks
@@ -196,7 +185,7 @@ impl PartialChunkStorage {
     ///
     /// # Panics
     /// If the chunk is not in the render distance.
-    pub fn set(&mut self, pos: &ChunkPos, chunk: Option<Arc<Mutex<Chunk>>>) {
+    pub fn set(&mut self, pos: &ChunkPos, chunk: Option<Arc<RwLock<Chunk>>>) {
         if let Some(chunk) = &chunk {
             self.shared
                 .write()
@@ -219,16 +208,30 @@ impl WeakChunkStorage {
             chunks: HashMap::new(),
         }
     }
-}
 
-impl Chunk {
-    pub fn read_with_dimension(
-        buf: &mut Cursor<&[u8]>,
-        data: &World,
-    ) -> Result<Self, BufReadError> {
-        Self::read_with_dimension_height(buf, data.height())
+    pub fn get(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
+        self.chunks.get(pos).and_then(|chunk| chunk.upgrade())
+    }
+
+    pub fn get_block_state(&self, pos: &BlockPos) -> Option<BlockState> {
+        let chunk_pos = ChunkPos::from(pos);
+        let chunk = self.get(&chunk_pos)?;
+        let chunk = chunk.read();
+        chunk.get(&ChunkBlockPos::from(pos), self.min_y)
+    }
+
+    pub fn set_block_state(&self, pos: &BlockPos, state: BlockState) -> Option<BlockState> {
+        if pos.y < self.min_y || pos.y >= (self.min_y + self.height as i32) {
+            return None;
+        }
+        let chunk_pos = ChunkPos::from(pos);
+        let chunk = self.get(&chunk_pos)?;
+        let mut chunk = chunk.write();
+        Some(chunk.get_and_set(&ChunkBlockPos::from(pos), state, self.min_y))
     }
+}
 
+impl Chunk {
     pub fn read_with_dimension_height(
         buf: &mut Cursor<&[u8]>,
         dimension_height: u32,
@@ -242,18 +245,12 @@ impl Chunk {
         Ok(Chunk { sections })
     }
 
-    pub fn section_index(&self, y: i32, min_y: i32) -> u32 {
-        assert!(y >= min_y, "y ({y}) must be at least {min_y}");
-        let min_section_index = min_y.div_floor(16);
-        (y.div_floor(16) - min_section_index) as u32
-    }
-
     pub fn get(&self, pos: &ChunkBlockPos, min_y: i32) -> Option<BlockState> {
         if pos.y < min_y {
             // y position is out of bounds
             return None;
         }
-        let section_index = self.section_index(pos.y, min_y) as usize;
+        let section_index = section_index(pos.y, min_y) as usize;
         if section_index >= self.sections.len() {
             // y position is out of bounds
             return None;
@@ -270,7 +267,7 @@ impl Chunk {
         state: BlockState,
         min_y: i32,
     ) -> BlockState {
-        let section_index = self.section_index(pos.y, min_y);
+        let section_index = section_index(pos.y, min_y);
         // TODO: make sure the section exists
         let section = &mut self.sections[section_index as usize];
         let chunk_section_pos = ChunkSectionBlockPos::from(pos);
@@ -278,7 +275,7 @@ impl Chunk {
     }
 
     pub fn set(&mut self, pos: &ChunkBlockPos, state: BlockState, min_y: i32) {
-        let section_index = self.section_index(pos.y, min_y);
+        let section_index = section_index(pos.y, min_y);
         // TODO: make sure the section exists
         let section = &mut self.sections[section_index as usize];
         let chunk_section_pos = ChunkSectionBlockPos::from(pos);
@@ -384,21 +381,28 @@ impl Default for WeakChunkStorage {
     }
 }
 
+/// Get the index of where a section is in a chunk based on its y coordinate
+/// and the minimum y coordinate of the world.
+pub fn section_index(y: i32, min_y: i32) -> u32 {
+    assert!(y >= min_y, "y ({y}) must be at least {min_y}");
+    let min_section_index = min_y.div_floor(16);
+    (y.div_floor(16) - min_section_index) as u32
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn test_section_index() {
-        let chunk = Chunk::default();
-        assert_eq!(chunk.section_index(0, 0), 0);
-        assert_eq!(chunk.section_index(128, 0), 8);
-        assert_eq!(chunk.section_index(127, 0), 7);
-        assert_eq!(chunk.section_index(0, -64), 4);
-        assert_eq!(chunk.section_index(-64, -64), 0);
-        assert_eq!(chunk.section_index(-49, -64), 0);
-        assert_eq!(chunk.section_index(-48, -64), 1);
-        assert_eq!(chunk.section_index(128, -64), 12);
+        assert_eq!(section_index(0, 0), 0);
+        assert_eq!(section_index(128, 0), 8);
+        assert_eq!(section_index(127, 0), 7);
+        assert_eq!(section_index(0, -64), 4);
+        assert_eq!(section_index(-64, -64), 0);
+        assert_eq!(section_index(-49, -64), 0);
+        assert_eq!(section_index(-48, -64), 1);
+        assert_eq!(section_index(128, -64), 12);
     }
 
     #[test]
@@ -406,21 +410,31 @@ mod tests {
         let mut chunk_storage = PartialChunkStorage::default();
         chunk_storage.set(
             &ChunkPos { x: 0, z: 0 },
-            Some(Arc::new(Mutex::new(Chunk::default()))),
+            Some(Arc::new(RwLock::new(Chunk::default()))),
         );
         assert!(chunk_storage
+            .shared
+            .read()
             .get_block_state(&BlockPos { x: 0, y: 319, z: 0 })
             .is_some());
         assert!(chunk_storage
+            .shared
+            .read()
             .get_block_state(&BlockPos { x: 0, y: 320, z: 0 })
             .is_none());
         assert!(chunk_storage
+            .shared
+            .read()
             .get_block_state(&BlockPos { x: 0, y: 338, z: 0 })
             .is_none());
         assert!(chunk_storage
+            .shared
+            .read()
             .get_block_state(&BlockPos { x: 0, y: -64, z: 0 })
             .is_some());
         assert!(chunk_storage
+            .shared
+            .read()
             .get_block_state(&BlockPos { x: 0, y: -65, z: 0 })
             .is_none());
     }
diff --git a/azalea-world/src/entity/mod.rs b/azalea-world/src/entity/mod.rs
index 814ef6d2..94362f2f 100644
--- a/azalea-world/src/entity/mod.rs
+++ b/azalea-world/src/entity/mod.rs
@@ -5,7 +5,7 @@ pub mod metadata;
 
 use self::attributes::{AttributeInstance, AttributeModifiers};
 pub use self::metadata::EntityMetadata;
-use crate::World;
+use crate::WeakWorld;
 use azalea_block::BlockState;
 use azalea_core::{BlockPos, Vec3, AABB};
 pub use data::*;
@@ -17,7 +17,7 @@ use uuid::Uuid;
 
 /// A reference to an entity in a world.
 #[derive(Debug)]
-pub struct Entity<'d, D = &'d mut World> {
+pub struct Entity<'d, D = &'d WeakWorld> {
     /// The world this entity is in.
     pub world: D,
     /// The incrementing numerical id of the entity.
@@ -26,7 +26,7 @@ pub struct Entity<'d, D = &'d mut World> {
     _marker: PhantomData<&'d ()>,
 }
 
-impl<'d, D: Deref<Target = World>> Entity<'d, D> {
+impl<'d, D: Deref<Target = WeakWorld>> Entity<'d, D> {
     pub fn new(world: D, id: u32, data: NonNull<EntityData>) -> Self {
         // TODO: have this be based on the entity type
         Self {
@@ -38,7 +38,7 @@ impl<'d, D: Deref<Target = World>> Entity<'d, D> {
     }
 }
 
-impl<'d, D: DerefMut<Target = World>> Entity<'d, D> {
+impl<'d, D: Deref<Target = WeakWorld>> Entity<'d, D> {
     /// Sets the position of the entity. This doesn't update the cache in
     /// azalea-world, and should only be used within azalea-world!
     ///
@@ -95,7 +95,7 @@ impl<'d, D: DerefMut<Target = World>> Entity<'d, D> {
     }
 }
 
-impl<'d, D: Deref<Target = World>> Entity<'d, D> {
+impl<'d, D: Deref<Target = WeakWorld>> Entity<'d, D> {
     #[inline]
     pub fn pos(&self) -> &Vec3 {
         &self.pos
@@ -150,8 +150,8 @@ impl<'d, D: Deref<Target = World>> Entity<'d, D> {
 
 // impl<
 //         'd,
-//         D: DerefMut<Target = World> + Deref<Target = World>,
-//         D2: Deref<Target = World>,
+//         D: Deref<Target = WeakWorld> + Deref<Target = WeakWorld>,
+//         D2: Deref<Target = WeakWorld>,
 //     > From<Entity<'d, D>> for Entity<'d, D2>
 // {
 //     fn from(entity: Entity<'d, D>) -> Entity<'d, D> {
@@ -164,13 +164,13 @@ impl<'d, D: Deref<Target = World>> Entity<'d, D> {
 //     }
 // }
 
-impl<D: DerefMut<Target = World>> DerefMut for Entity<'_, D> {
+impl<D: Deref<Target = WeakWorld>> DerefMut for Entity<'_, D> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         unsafe { self.data.as_mut() }
     }
 }
 
-impl<D: Deref<Target = World>> Deref for Entity<'_, D> {
+impl<D: Deref<Target = WeakWorld>> Deref for Entity<'_, D> {
     type Target = EntityData;
 
     fn deref(&self) -> &Self::Target {
@@ -287,10 +287,11 @@ impl EntityData {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::PartialWorld;
 
     #[test]
     fn from_mut_entity_to_ref_entity() {
-        let mut world = World::default();
+        let mut world = PartialWorld::default();
         let uuid = Uuid::from_u128(100);
         world.add_entity(
             0,
@@ -301,7 +302,6 @@ mod tests {
             ),
         );
         let entity: Entity = world.entity_mut(0).unwrap();
-        let entity_ref = Entity::from(entity);
-        assert_eq!(entity_ref.uuid, uuid);
+        assert_eq!(entity.uuid, uuid);
     }
 }
diff --git a/azalea-world/src/entity_storage.rs b/azalea-world/src/entity_storage.rs
index 0e8fc0b5..71c5413c 100755
--- a/azalea-world/src/entity_storage.rs
+++ b/azalea-world/src/entity_storage.rs
@@ -34,7 +34,11 @@ use uuid::Uuid;
 pub struct PartialEntityStorage {
     pub shared: Arc<RwLock<WeakEntityStorage>>,
 
-    /// The entity id of the player that owns this struct.
+    /// The entity id of the player that owns this partial world. This will
+    /// make [`PartialWorld::entity_mut`] pretend the entity doesn't exist so
+    /// it doesn't get modified from outside sources.
+    ///
+    /// [`PartialWorld::entity_mut`]: crate::PartialWorld::entity_mut
     pub owner_entity_id: u32,
     pub updates_received: IntMap<u32, u32>,
     /// Strong references to the entities we have loaded.
@@ -172,11 +176,7 @@ impl PartialEntityStorage {
     /// Get an entity in the shared storage by its id, if it exists.
     #[inline]
     pub fn get_by_id(&self, id: u32) -> Option<Arc<EntityData>> {
-        self.shared
-            .read()
-            .data_by_id
-            .get(&id)
-            .and_then(|e| e.upgrade())
+        self.shared.read().get_by_id(id)
     }
 
     /// Get a reference to an entity by its UUID, if it's being loaded by this
@@ -204,13 +204,7 @@ impl PartialEntityStorage {
     /// Get an entity in the shared storage by its UUID, if it exists.
     #[inline]
     pub fn get_by_uuid(&self, uuid: &Uuid) -> Option<Arc<EntityData>> {
-        self.shared.read().id_by_uuid.get(uuid).and_then(|id| {
-            self.shared
-                .read()
-                .data_by_id
-                .get(id)
-                .and_then(|e| e.upgrade())
-        })
+        self.shared.read().get_by_uuid(uuid)
     }
 
     /// Clear all entities in a chunk. This will not clear them from the
@@ -241,46 +235,23 @@ impl PartialEntityStorage {
         old_chunk: &ChunkPos,
         new_chunk: &ChunkPos,
     ) {
-        if let Some(entities) = self.shared.write().ids_by_chunk.get_mut(old_chunk) {
-            entities.remove(&entity_id);
-        }
         self.shared
             .write()
-            .ids_by_chunk
-            .entry(*new_chunk)
-            .or_default()
-            .insert(entity_id);
+            .update_entity_chunk(entity_id, old_chunk, new_chunk);
     }
 
-    pub fn find_one_entity<F>(&self, mut f: F) -> Option<Arc<EntityData>>
+    pub fn find_entity<F>(&self, mut f: F) -> Option<Arc<EntityData>>
     where
         F: FnMut(&Arc<EntityData>) -> bool,
     {
-        for entity in self.shared.read().entities() {
-            if let Some(entity) = entity.upgrade() {
-                if f(&entity) {
-                    return Some(entity);
-                }
-            }
-        }
-        None
+        self.shared.read().find_entity(|e| f(e))
     }
 
-    pub fn find_one_entity_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<Arc<EntityData>>
+    pub fn find_entity_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<Arc<EntityData>>
     where
         F: FnMut(&EntityData) -> bool,
     {
-        let shared = self.shared.read();
-        if let Some(entities) = shared.ids_by_chunk.get(chunk) {
-            for entity_id in entities {
-                if let Some(entity) = shared.data_by_id.get(entity_id).and_then(|e| e.upgrade()) {
-                    if f(&entity) {
-                        return Some(entity);
-                    }
-                }
-            }
-        }
-        None
+        self.shared.read().find_entity_in_chunk(chunk, |e| f(e))
     }
 }
 
@@ -366,6 +337,67 @@ impl WeakEntityStorage {
     pub fn contains_id(&self, id: &u32) -> bool {
         self.data_by_id.contains_key(id)
     }
+
+    /// Get an entity by its id, if it exists.
+    #[inline]
+    pub fn get_by_id(&self, id: u32) -> Option<Arc<EntityData>> {
+        self.data_by_id.get(&id).and_then(|e| e.upgrade())
+    }
+
+    /// Get an entity in the shared storage by its UUID, if it exists.
+    #[inline]
+    pub fn get_by_uuid(&self, uuid: &Uuid) -> Option<Arc<EntityData>> {
+        self.id_by_uuid
+            .get(uuid)
+            .and_then(|id| self.data_by_id.get(id).and_then(|e| e.upgrade()))
+    }
+
+    pub fn find_entity<F>(&self, mut f: F) -> Option<Arc<EntityData>>
+    where
+        F: FnMut(&Arc<EntityData>) -> bool,
+    {
+        for entity in self.entities() {
+            if let Some(entity) = entity.upgrade() {
+                if f(&entity) {
+                    return Some(entity);
+                }
+            }
+        }
+        None
+    }
+
+    pub fn find_entity_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<Arc<EntityData>>
+    where
+        F: FnMut(&EntityData) -> bool,
+    {
+        if let Some(entities) = self.ids_by_chunk.get(chunk) {
+            for entity_id in entities {
+                if let Some(entity) = self.data_by_id.get(entity_id).and_then(|e| e.upgrade()) {
+                    if f(&entity) {
+                        return Some(entity);
+                    }
+                }
+            }
+        }
+        None
+    }
+
+    /// Move an entity from its old chunk to a new chunk.
+    #[inline]
+    pub fn update_entity_chunk(
+        &mut self,
+        entity_id: u32,
+        old_chunk: &ChunkPos,
+        new_chunk: &ChunkPos,
+    ) {
+        if let Some(entities) = self.ids_by_chunk.get_mut(old_chunk) {
+            entities.remove(&entity_id);
+        }
+        self.ids_by_chunk
+            .entry(*new_chunk)
+            .or_default()
+            .insert(entity_id);
+    }
 }
 
 #[cfg(test)]
diff --git a/azalea-world/src/world.rs b/azalea-world/src/world.rs
index 65d001c3..648613b9 100644
--- a/azalea-world/src/world.rs
+++ b/azalea-world/src/world.rs
@@ -6,18 +6,24 @@ use crate::{
 use azalea_block::BlockState;
 use azalea_buf::BufReadError;
 use azalea_core::{BlockPos, ChunkPos, PositionDelta8, Vec3};
-use parking_lot::{Mutex, RwLock};
+use parking_lot::RwLock;
 use std::{backtrace::Backtrace, fmt::Debug};
 use std::{fmt::Formatter, io::Cursor, sync::Arc};
 use uuid::Uuid;
 
-/// A world is a collection of chunks and entities. They're called "levels" in
-/// Minecraft's source code.
+/// PartialWorlds are usually owned by clients, and hold strong references to
+/// chunks and entities in [`WeakWorld`]s.
+///
+/// Basically, they hold the chunks and entities that are within render
+/// distance but can still access chunks and entities owned by other
+/// `PartialWorld`s that have the same `WeakWorld`.
+///
+/// This is primarily useful for having multiple clients in the same world.
 #[derive(Default)]
-pub struct World {
+pub struct PartialWorld {
     // we just need to keep a strong reference to `shared` so it doesn't get
     // dropped, we don't need to do anything with it
-    _shared: Arc<WeakWorld>,
+    pub shared: Arc<WeakWorld>,
 
     pub chunk_storage: PartialChunkStorage,
     pub entity_storage: PartialEntityStorage,
@@ -31,10 +37,10 @@ pub struct WeakWorld {
     pub entity_storage: Arc<RwLock<WeakEntityStorage>>,
 }
 
-impl World {
+impl PartialWorld {
     pub fn new(chunk_radius: u32, shared: Arc<WeakWorld>, owner_entity_id: u32) -> Self {
-        World {
-            _shared: shared.clone(),
+        PartialWorld {
+            shared: shared.clone(),
             chunk_storage: PartialChunkStorage::new(chunk_radius, shared.chunk_storage.clone()),
             entity_storage: PartialEntityStorage::new(
                 shared.entity_storage.clone(),
@@ -51,13 +57,13 @@ impl World {
         self.chunk_storage.replace_with_packet_data(pos, data)
     }
 
-    pub fn get_chunk(&self, pos: &ChunkPos) -> Option<Arc<Mutex<Chunk>>> {
+    pub fn get_chunk(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
         self.chunk_storage.get(pos)
     }
 
     pub fn set_chunk(&mut self, pos: &ChunkPos, chunk: Option<Chunk>) -> Result<(), BufReadError> {
         self.chunk_storage
-            .set(pos, chunk.map(|c| Arc::new(Mutex::new(c))));
+            .set(pos, chunk.map(|c| Arc::new(RwLock::new(c))));
         Ok(())
     }
 
@@ -65,14 +71,24 @@ impl World {
         self.chunk_storage.view_center = *pos;
     }
 
-    pub fn get_block_state(&self, pos: &BlockPos) -> Option<BlockState> {
-        self.chunk_storage.get_block_state(pos)
-    }
-
     pub fn set_block_state(&mut self, pos: &BlockPos, state: BlockState) -> Option<BlockState> {
         self.chunk_storage.set_block_state(pos, state)
     }
 
+    /// Returns a mutable reference to the entity with the given ID.
+    pub fn entity_mut(&mut self, id: u32) -> Option<Entity<'_, &WeakWorld>> {
+        // no entity for you (we're processing this entity somewhere else)
+        if id != self.entity_storage.owner_entity_id && !self.entity_storage.maybe_update(id) {
+            return None;
+        }
+
+        self.shared.entity(id)
+    }
+
+    pub fn add_entity(&mut self, id: u32, entity: EntityData) {
+        self.entity_storage.insert(id, entity);
+    }
+
     pub fn set_entity_pos(&mut self, entity_id: u32, new_pos: Vec3) -> Result<(), MoveEntityError> {
         let mut entity = self
             .entity_mut(entity_id)
@@ -109,71 +125,78 @@ impl World {
         }
         Ok(())
     }
+}
 
-    pub fn add_entity(&mut self, id: u32, entity: EntityData) {
-        self.entity_storage.insert(id, entity);
+impl WeakWorld {
+    pub fn new(height: u32, min_y: i32) -> Self {
+        WeakWorld {
+            chunk_storage: Arc::new(RwLock::new(WeakChunkStorage::new(height, min_y))),
+            entity_storage: Arc::new(RwLock::new(WeakEntityStorage::new())),
+        }
     }
 
+    /// Read the total height of the world. You can add this to [`Self::min_y`]
+    /// to get the highest possible y coordinate a block can be placed at.
     pub fn height(&self) -> u32 {
-        self.chunk_storage.height()
+        self.chunk_storage.read().height
     }
 
+    /// Get the lowest possible y coordinate a block can be placed at.
     pub fn min_y(&self) -> i32 {
-        self.chunk_storage.min_y()
+        self.chunk_storage.read().min_y
     }
 
     pub fn entity_data_by_id(&self, id: u32) -> Option<Arc<EntityData>> {
-        self.entity_storage.get_by_id(id)
+        self.entity_storage.read().get_by_id(id)
     }
 
-    pub fn entity(&self, id: u32) -> Option<Entity<&World>> {
-        let entity_data = self.entity_storage.get_by_id(id)?;
-        let entity_ptr = unsafe { entity_data.as_ptr() };
-        Some(Entity::new(self, id, entity_ptr))
-    }
-
-    /// Returns a mutable reference to the entity with the given ID.
-    pub fn entity_mut(&mut self, id: u32) -> Option<Entity<'_, &mut World>> {
-        // no entity for you (we're processing this entity somewhere else)
-        if id != self.entity_storage.owner_entity_id && !self.entity_storage.maybe_update(id) {
-            return None;
-        }
-
-        let entity_data = self.entity_storage.get_by_id(id)?;
+    /// Returns a entity with the given ID.
+    ///
+    /// The returned Entity can technically be mutated, but you should avoid
+    /// doing any relative mutations.
+    pub fn entity(&self, id: u32) -> Option<Entity<&WeakWorld>> {
+        let entity_data = self.entity_storage.read().get_by_id(id)?;
         let entity_ptr = unsafe { entity_data.as_ptr() };
         Some(Entity::new(self, id, entity_ptr))
     }
 
     pub fn entity_by_uuid(&self, uuid: &Uuid) -> Option<Arc<EntityData>> {
-        self.entity_storage.get_by_uuid(uuid)
+        self.entity_storage.read().get_by_uuid(uuid)
     }
 
-    pub fn find_one_entity<F>(&self, mut f: F) -> Option<Arc<EntityData>>
+    pub fn find_entity<F>(&self, mut f: F) -> Option<Arc<EntityData>>
     where
         F: FnMut(&EntityData) -> bool,
     {
-        self.entity_storage.find_one_entity(|entity| f(entity))
+        self.entity_storage.read().find_entity(|entity| f(entity))
     }
-}
 
-impl WeakWorld {
-    pub fn new(height: u32, min_y: i32) -> Self {
-        WeakWorld {
-            chunk_storage: Arc::new(RwLock::new(WeakChunkStorage::new(height, min_y))),
-            entity_storage: Arc::new(RwLock::new(WeakEntityStorage::new())),
+    pub fn set_entity_pos(&self, entity_id: u32, new_pos: Vec3) -> Result<(), MoveEntityError> {
+        let mut entity = self
+            .entity(entity_id)
+            .ok_or_else(|| MoveEntityError::EntityDoesNotExist(Backtrace::capture()))?;
+        let old_chunk = ChunkPos::from(entity.pos());
+        let new_chunk = ChunkPos::from(&new_pos);
+        // this is fine because we update the chunk below
+        unsafe { entity.move_unchecked(new_pos) };
+        if old_chunk != new_chunk {
+            self.entity_storage
+                .write()
+                .update_entity_chunk(entity_id, &old_chunk, &new_chunk);
         }
+        Ok(())
     }
 
-    pub fn height(&self) -> u32 {
-        self.chunk_storage.read().height
+    pub fn get_block_state(&self, pos: &BlockPos) -> Option<BlockState> {
+        self.chunk_storage.read().get_block_state(pos)
     }
 
-    pub fn min_y(&self) -> i32 {
-        self.chunk_storage.read().min_y
+    pub fn get_chunk(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
+        self.chunk_storage.read().get(pos)
     }
 }
 
-impl Debug for World {
+impl Debug for PartialWorld {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("World")
             .field("chunk_storage", &self.chunk_storage)