fix despawning entities on dimension change

5 files changed, 141 insertions, 17 deletions

diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index c3381e59..a13f1490 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -207,18 +207,22 @@ impl GamePacketHandler<'_> {
 
         as_system::<(
             Commands,
-            Query<(
-                &GameProfileComponent,
-                &ClientInformation,
-                Option<&mut InstanceName>,
-                Option<&mut LoadedBy>,
-                &mut EntityIdIndex,
-                &mut InstanceHolder,
-            )>,
+            Query<
+                (
+                    &GameProfileComponent,
+                    &ClientInformation,
+                    Option<&mut InstanceName>,
+                    Option<&mut LoadedBy>,
+                    &mut EntityIdIndex,
+                    &mut InstanceHolder,
+                ),
+                With<LocalEntity>,
+            >,
             EventWriter<InstanceLoadedEvent>,
             ResMut<InstanceContainer>,
             ResMut<EntityUuidIndex>,
             EventWriter<SendPacketEvent>,
+            Query<&mut LoadedBy, Without<LocalEntity>>,
         )>(
             self.ecs,
             |(
@@ -228,6 +232,7 @@ impl GamePacketHandler<'_> {
                 mut instance_container,
                 mut entity_uuid_index,
                 mut send_packet_events,
+                mut loaded_by_query,
             )| {
                 let (
                     game_profile,
@@ -317,6 +322,11 @@ impl GamePacketHandler<'_> {
                     &mut instance_holder.instance.write(),
                 );
 
+                // every entity is now unloaded by this player
+                for mut loaded_by in &mut loaded_by_query.iter_mut() {
+                    loaded_by.remove(&self.player);
+                }
+
                 // update or insert loaded_by
                 if let Some(mut loaded_by) = loaded_by {
                     loaded_by.insert(self.player);
@@ -1413,16 +1423,20 @@ impl GamePacketHandler<'_> {
 
         as_system::<(
             Commands,
-            Query<(
-                &mut InstanceHolder,
-                &GameProfileComponent,
-                &ClientInformation,
-            )>,
+            Query<
+                (
+                    &mut InstanceHolder,
+                    &GameProfileComponent,
+                    &ClientInformation,
+                ),
+                With<LocalEntity>,
+            >,
             EventWriter<_>,
             ResMut<InstanceContainer>,
+            Query<&mut LoadedBy, Without<LocalEntity>>,
         )>(
             self.ecs,
-            |(mut commands, mut query, mut events, mut instance_container)| {
+            |(mut commands, mut query, mut events, mut instance_container, mut loaded_by_query)| {
                 let (mut instance_holder, game_profile, client_information) =
                     query.get_mut(self.player).unwrap();
 
@@ -1461,6 +1475,11 @@ impl GamePacketHandler<'_> {
                 );
                 instance_holder.instance = weak_instance;
 
+                // every entity is now unloaded by this player
+                for mut loaded_by in &mut loaded_by_query.iter_mut() {
+                    loaded_by.remove(&self.player);
+                }
+
                 // this resets a bunch of our components like physics and stuff
                 let entity_bundle = EntityBundle::new(
                     game_profile.uuid,
diff --git a/azalea-client/tests/despawn_entities_when_changing_dimension.rs b/azalea-client/tests/despawn_entities_when_changing_dimension.rs
new file mode 100644
index 00000000..39cea091
--- /dev/null
+++ b/azalea-client/tests/despawn_entities_when_changing_dimension.rs
@@ -0,0 +1,97 @@
+use azalea_client::{InConfigState, InGameState, test_simulation::*};
+use azalea_core::{
+    delta::PositionDelta8,
+    position::{ChunkPos, Vec3},
+    resource_location::ResourceLocation,
+};
+use azalea_entity::{LocalEntity, metadata::Cow};
+use azalea_protocol::packets::{
+    ConnectionProtocol,
+    config::{ClientboundFinishConfiguration, ClientboundRegistryData},
+    game::ClientboundAddEntity,
+};
+use azalea_registry::DimensionType;
+use azalea_world::InstanceName;
+use bevy_ecs::{entity::Entity, query::With};
+use bevy_log::tracing_subscriber;
+use simdnbt::owned::{NbtCompound, NbtTag};
+use uuid::Uuid;
+
+#[test]
+fn test_despawn_entities_when_changing_dimension() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut simulation = Simulation::new(ConnectionProtocol::Configuration);
+    simulation.receive_packet(ClientboundRegistryData {
+        registry_id: ResourceLocation::new("minecraft:dimension_type"),
+        entries: vec![
+            (
+                ResourceLocation::new("minecraft:overworld"),
+                Some(NbtCompound::from_values(vec![
+                    ("height".into(), NbtTag::Int(384)),
+                    ("min_y".into(), NbtTag::Int(-64)),
+                ])),
+            ),
+            (
+                ResourceLocation::new("minecraft:nether"),
+                Some(NbtCompound::from_values(vec![
+                    ("height".into(), NbtTag::Int(256)),
+                    ("min_y".into(), NbtTag::Int(0)),
+                ])),
+            ),
+        ]
+        .into_iter()
+        .collect(),
+    });
+    simulation.tick();
+    simulation.receive_packet(ClientboundFinishConfiguration);
+    simulation.tick();
+
+    //
+    // OVERWORLD
+    //
+
+    simulation.receive_packet(make_basic_login_packet(
+        DimensionType::new_raw(0), // overworld
+        ResourceLocation::new("azalea:a"),
+    ));
+    simulation.tick();
+
+    simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), (384 + 64) / 16));
+    simulation.tick();
+    // spawn a cow
+    simulation.receive_packet(ClientboundAddEntity {
+        id: 123.into(),
+        uuid: Uuid::from_u128(1234),
+        entity_type: azalea_registry::EntityKind::Cow,
+        position: Vec3::new(0., 64., 0.),
+        x_rot: 0,
+        y_rot: 0,
+        y_head_rot: 0,
+        data: 0,
+        velocity: PositionDelta8::default(),
+    });
+    simulation.tick();
+    // make sure it's spawned
+    let mut cow_query = simulation.app.world_mut().query_filtered::<(), With<Cow>>();
+    let cow_iter = cow_query.iter(simulation.app.world());
+    assert_eq!(cow_iter.count(), 1, "cow should be spawned");
+
+    //
+    // NETHER
+    //
+
+    simulation.receive_packet(make_basic_respawn_packet(
+        DimensionType::new_raw(1), // nether
+        ResourceLocation::new("azalea:b"),
+    ));
+    simulation.tick();
+
+    // cow should be completely deleted from the ecs
+    let cow_iter = cow_query.iter(simulation.app.world());
+    assert_eq!(
+        cow_iter.count(),
+        0,
+        "cow should be despawned after switching dimensions"
+    );
+}
diff --git a/azalea-entity/src/plugin/indexing.rs b/azalea-entity/src/plugin/indexing.rs
index 21dd273a..fefecb06 100644
--- a/azalea-entity/src/plugin/indexing.rs
+++ b/azalea-entity/src/plugin/indexing.rs
@@ -7,7 +7,7 @@ use azalea_world::{Instance, InstanceContainer, InstanceName, MinecraftEntityId}
 use bevy_ecs::{
     component::Component,
     entity::Entity,
-    query::{Added, Changed},
+    query::{Added, Changed, Without},
     system::{Commands, Query, Res, ResMut, Resource},
 };
 use derive_more::{Deref, DerefMut};
@@ -16,7 +16,7 @@ use tracing::{debug, warn};
 use uuid::Uuid;
 
 use super::LoadedBy;
-use crate::{EntityUuid, Position};
+use crate::{EntityUuid, LocalEntity, Position};
 
 #[derive(Resource, Default)]
 pub struct EntityUuidIndex {
@@ -152,7 +152,7 @@ pub fn remove_despawned_entities_from_indexes(
             &InstanceName,
             &LoadedBy,
         ),
-        Changed<LoadedBy>,
+        (Changed<LoadedBy>, Without<LocalEntity>),
     >,
 ) {
     for (entity, uuid, minecraft_id, position, instance_name, loaded_by) in &query {
diff --git a/azalea-entity/src/plugin/mod.rs b/azalea-entity/src/plugin/mod.rs
index bddb4b66..b6ae369f 100644
--- a/azalea-entity/src/plugin/mod.rs
+++ b/azalea-entity/src/plugin/mod.rs
@@ -206,6 +206,9 @@ pub fn update_bounding_box(mut query: Query<(&Position, &mut Physics), Changed<P
 
 /// Marks an entity that's in a loaded chunk. This is updated at the beginning
 /// of every tick.
+///
+/// Internally, this is only used for player physics. Not to be confused with
+/// the somewhat similarly named [`LoadedBy`].
 #[derive(Component, Clone, Debug, Copy)]
 pub struct InLoadedChunk;
 
diff --git a/azalea-world/src/world.rs b/azalea-world/src/world.rs
index 3428ab5e..6b72c2c5 100644
--- a/azalea-world/src/world.rs
+++ b/azalea-world/src/world.rs
@@ -97,6 +97,11 @@ impl Display for MinecraftEntityId {
         write!(f, "eid({})", self.0)
     }
 }
+impl From<i32> for MinecraftEntityId {
+    fn from(id: i32) -> Self {
+        Self(id)
+    }
+}
 
 /// Keep track of certain metadatas that are only relevant for this partial
 /// world.