fix PartialWorld::default() for tests

4 files changed, 37 insertions, 48 deletions

diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 44c0f8cc..c88cfc31 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -523,7 +523,7 @@ impl Client {
                     *world_lock = PartialWorld::new(
                         client.client_information.read().view_distance.into(),
                         weak_world,
-                        p.player_id,
+                        Some(p.player_id),
                     );
 
                     let entity = EntityData::new(
diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
index 6d3907f0..c0c5f43a 100755
--- a/azalea-world/src/chunk_storage.rs
+++ b/azalea-world/src/chunk_storage.rs
@@ -20,7 +20,7 @@ const SECTION_HEIGHT: u32 = 16;
 /// [`ChunkStorage`] instead.
 pub struct PartialChunkStorage {
     /// Chunk storage that can be shared by clients.
-    shared: Arc<RwLock<WeakChunkStorage>>,
+    pub shared: Arc<RwLock<WeakChunkStorage>>,
 
     pub view_center: ChunkPos,
     chunk_radius: u32,
@@ -32,6 +32,7 @@ pub struct PartialChunkStorage {
 /// A storage for chunks where they're only stored weakly, so if they're not
 /// actively being used somewhere else they'll be forgotten. This is used for
 /// shared worlds.
+#[derive(Debug)]
 pub struct WeakChunkStorage {
     pub height: u32,
     pub min_y: i32,
diff --git a/azalea-world/src/entity_storage.rs b/azalea-world/src/entity_storage.rs
index 2a4c83ff..1bc1adc5 100755
--- a/azalea-world/src/entity_storage.rs
+++ b/azalea-world/src/entity_storage.rs
@@ -30,6 +30,8 @@ use uuid::Uuid;
 /// `storage.shared.read().entities` [`WeakEntityStorage::entities`].
 ///
 /// This is meant to be used with shared worlds.
+///
+/// You can access the shared storage with `world.shared.read()`.
 #[derive(Debug, Default)]
 pub struct PartialEntityStorage {
     pub shared: Arc<RwLock<WeakEntityStorage>>,
@@ -39,7 +41,7 @@ pub struct PartialEntityStorage {
     /// it doesn't get modified from outside sources.
     ///
     /// [`PartialWorld::entity_mut`]: crate::PartialWorld::entity_mut
-    pub owner_entity_id: u32,
+    pub owner_entity_id: Option<u32>,
     pub updates_received: IntMap<u32, u32>,
     /// Strong references to the entities we have loaded.
     data_by_id: IntMap<u32, Arc<EntityData>>,
@@ -59,8 +61,10 @@ pub struct WeakEntityStorage {
 }
 
 impl PartialEntityStorage {
-    pub fn new(shared: Arc<RwLock<WeakEntityStorage>>, owner_entity_id: u32) -> Self {
-        shared.write().updates_received.insert(owner_entity_id, 0);
+    pub fn new(shared: Arc<RwLock<WeakEntityStorage>>, owner_entity_id: Option<u32>) -> Self {
+        if let Some(owner_entity_id) = owner_entity_id {
+            shared.write().updates_received.insert(owner_entity_id, 0);
+        }
         Self {
             shared,
             owner_entity_id,
@@ -94,7 +98,7 @@ impl PartialEntityStorage {
         // not there in which case set both to 1
         if let Some(&shared_updates_received) = shared.updates_received.get(&id) {
             // 0 means we're never tracking updates for this entity
-            if shared_updates_received != 0 || id == self.owner_entity_id {
+            if shared_updates_received != 0 || Some(id) == self.owner_entity_id {
                 self.updates_received.insert(id, 1);
             }
         } else {
@@ -127,15 +131,6 @@ impl PartialEntityStorage {
         self.data_by_id.contains_key(id)
     }
 
-    /// Whether the entity with the given id is in the shared storage (i.e.
-    /// it's possible we don't see the entity but something else in the shared
-    /// storage does). To check whether the entity is being loaded by this
-    /// storage, use [`PartialEntityStorage::limited_contains_id`].
-    #[inline]
-    pub fn contains_id(&self, id: &u32) -> bool {
-        self.shared.read().data_by_id.contains_key(id)
-    }
-
     /// Get a reference to an entity by its id, if it's being loaded by this
     /// storage.
     #[inline]
@@ -173,12 +168,6 @@ 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().get_by_id(id)
-    }
-
     /// Get a reference to an entity by its UUID, if it's being loaded by this
     /// storage.
     #[inline]
@@ -201,12 +190,6 @@ impl PartialEntityStorage {
             .and_then(|id| self.data_by_id.get_mut(id))
     }
 
-    /// 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().get_by_uuid(uuid)
-    }
-
     /// Clear all entities in a chunk. This will not clear them from the
     /// shared storage, unless there are no other references to them.
     pub fn clear_chunk(&mut self, chunk: &ChunkPos) {
@@ -239,20 +222,6 @@ impl PartialEntityStorage {
             .write()
             .update_entity_chunk(entity_id, old_chunk, new_chunk);
     }
-
-    pub fn entity_by<F>(&self, mut f: F) -> Option<Arc<EntityData>>
-    where
-        F: FnMut(&Arc<EntityData>) -> bool,
-    {
-        self.shared.read().entity_by(|e| f(e))
-    }
-
-    pub fn entity_by_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<Arc<EntityData>>
-    where
-        F: FnMut(&EntityData) -> bool,
-    {
-        self.shared.read().entity_by_in_chunk(chunk, |e| f(e))
-    }
 }
 
 impl WeakEntityStorage {
@@ -425,7 +394,8 @@ mod tests {
     #[test]
     fn test_store_entity() {
         let mut storage = PartialEntityStorage::default();
-        assert!(storage.get_by_id(0).is_none());
+        assert!(storage.limited_get_by_id(0).is_none());
+        assert!(storage.shared.read().get_by_id(0).is_none());
 
         let uuid = Uuid::from_u128(100);
         storage.insert(
@@ -436,9 +406,11 @@ mod tests {
                 EntityMetadata::Player(metadata::Player::default()),
             ),
         );
-        assert_eq!(storage.get_by_id(0).unwrap().uuid, uuid);
+        assert_eq!(storage.limited_get_by_id(0).unwrap().uuid, uuid);
+        assert_eq!(storage.shared.read().get_by_id(0).unwrap().uuid, uuid);
 
         storage.remove_by_id(0);
-        assert!(storage.get_by_id(0).is_none());
+        assert!(storage.limited_get_by_id(0).is_none());
+        assert!(storage.shared.read().get_by_id(0).is_none());
     }
 }
diff --git a/azalea-world/src/world.rs b/azalea-world/src/world.rs
index ad6037a5..6a3652f8 100644
--- a/azalea-world/src/world.rs
+++ b/azalea-world/src/world.rs
@@ -19,7 +19,6 @@ use uuid::Uuid;
 /// `PartialWorld`s that have the same `WeakWorld`.
 ///
 /// This is primarily useful for having multiple clients in the same world.
-#[derive(Default)]
 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
@@ -31,14 +30,14 @@ pub struct PartialWorld {
 
 /// A world where the chunks are stored as weak pointers. This is used for
 /// shared worlds.
-#[derive(Default)]
+#[derive(Default, Debug)]
 pub struct WeakWorld {
     pub chunk_storage: Arc<RwLock<WeakChunkStorage>>,
     pub entity_storage: Arc<RwLock<WeakEntityStorage>>,
 }
 
 impl PartialWorld {
-    pub fn new(chunk_radius: u32, shared: Arc<WeakWorld>, owner_entity_id: u32) -> Self {
+    pub fn new(chunk_radius: u32, shared: Arc<WeakWorld>, owner_entity_id: Option<u32>) -> Self {
         PartialWorld {
             shared: shared.clone(),
             chunk_storage: PartialChunkStorage::new(chunk_radius, shared.chunk_storage.clone()),
@@ -78,7 +77,8 @@ impl PartialWorld {
     /// 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) {
+        if Some(id) != self.entity_storage.owner_entity_id && !self.entity_storage.maybe_update(id)
+        {
             return None;
         }
 
@@ -208,6 +208,22 @@ impl Debug for PartialWorld {
         f.debug_struct("World")
             .field("chunk_storage", &self.chunk_storage)
             .field("entity_storage", &self.entity_storage)
+            .field("shared", &self.shared)
             .finish()
     }
 }
+
+impl Default for PartialWorld {
+    fn default() -> Self {
+        let chunk_storage = PartialChunkStorage::default();
+        let entity_storage = PartialEntityStorage::default();
+        Self {
+            shared: Arc::new(WeakWorld {
+                chunk_storage: chunk_storage.shared.clone(),
+                entity_storage: entity_storage.shared.clone(),
+            }),
+            chunk_storage,
+            entity_storage,
+        }
+    }
+}