fix chunk errors when joining a world with a same name but different height

6 files changed, 163 insertions, 3 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 68702b49..74f2d23a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,3 +34,4 @@ write down most non-trivial breaking changes.
 - Several protocol fixes, including for ClientboundSetPlayerTeam and a few data components.
 - Update the `InstanceName` component correctly when we receive a respawn or second login packet.
 - Block shapes and some properties were using data from `1.20.3-pre4` due to using an old data generator (Pixlyzer), which has now been replaced with the data generator from [Pumpkin](https://github.com/Pumpkin-MC/Extractor).
+- No more chunk errors when the client joins another world with the same name but different height.
diff --git a/azalea-client/src/local_player.rs b/azalea-client/src/local_player.rs
index 455cc470..bf1609ab 100644
--- a/azalea-client/src/local_player.rs
+++ b/azalea-client/src/local_player.rs
@@ -144,6 +144,20 @@ impl InstanceHolder {
             ))),
         }
     }
+
+    /// Reset the `Instance` to a new reference to an empty instance, but with
+    /// the same registries as the current one.
+    ///
+    /// This is used by Azalea when entering the config state.
+    pub fn reset(&mut self) {
+        let registries = self.instance.read().registries.clone();
+
+        let mut new_instance = Instance::default();
+        new_instance.registries = registries;
+        self.instance = Arc::new(RwLock::new(new_instance));
+
+        self.partial_instance.write().reset();
+    }
 }
 
 #[derive(Error, Debug)]
diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index 71766f8b..b301973f 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -1496,10 +1496,11 @@ impl GamePacketHandler<'_> {
     pub fn start_configuration(&mut self, _p: &ClientboundStartConfiguration) {
         debug!("Got start configuration packet");
 
-        as_system::<(Commands, Query<&mut RawConnection>)>(
+        as_system::<(Commands, Query<(&mut RawConnection, &mut InstanceHolder)>)>(
             self.ecs,
             |(mut commands, mut query)| {
-                let Some(mut raw_conn) = query.get_mut(self.player).ok() else {
+                let Some((mut raw_conn, mut instance_holder)) = query.get_mut(self.player).ok()
+                else {
                     warn!("Got start configuration packet but player doesn't have a RawConnection");
                     return;
                 };
@@ -1515,6 +1516,8 @@ impl GamePacketHandler<'_> {
                     .insert(crate::client::InConfigState)
                     .remove::<crate::JoinedClientBundle>()
                     .remove::<EntityBundle>();
+
+                instance_holder.reset();
             },
         );
     }
diff --git a/azalea-client/tests/login_to_dimension_with_same_name.rs b/azalea-client/tests/login_to_dimension_with_same_name.rs
new file mode 100644
index 00000000..be362bb7
--- /dev/null
+++ b/azalea-client/tests/login_to_dimension_with_same_name.rs
@@ -0,0 +1,136 @@
+use azalea_client::{InConfigState, InGameState, InstanceHolder, test_simulation::*};
+use azalea_core::{position::ChunkPos, resource_location::ResourceLocation};
+use azalea_entity::LocalEntity;
+use azalea_protocol::packets::{
+    ConnectionProtocol, Packet,
+    config::{ClientboundFinishConfiguration, ClientboundRegistryData},
+    game::ClientboundStartConfiguration,
+};
+use azalea_registry::{DataRegistry, DimensionType};
+use azalea_world::InstanceName;
+use bevy_log::tracing_subscriber;
+use simdnbt::owned::{NbtCompound, NbtTag};
+
+#[test]
+fn test_login_to_dimension_with_same_name() {
+    let _ = tracing_subscriber::fmt().try_init();
+
+    generic_test_login_to_dimension_with_same_name(true);
+    generic_test_login_to_dimension_with_same_name(false);
+}
+
+fn generic_test_login_to_dimension_with_same_name(using_respawn: bool) {
+    let make_basic_login_or_respawn_packet = if using_respawn {
+        |dimension: DimensionType, instance_name: ResourceLocation| {
+            make_basic_respawn_packet(dimension, instance_name).into_variant()
+        }
+    } else {
+        |dimension: DimensionType, instance_name: ResourceLocation| {
+            make_basic_login_packet(dimension, instance_name).into_variant()
+        }
+    };
+
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut simulation = Simulation::new(ConnectionProtocol::Configuration);
+    assert!(simulation.has_component::<InConfigState>());
+    assert!(!simulation.has_component::<InGameState>());
+
+    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)),
+            ])),
+        )]
+        .into_iter()
+        .collect(),
+    });
+    simulation.tick();
+    simulation.receive_packet(ClientboundFinishConfiguration);
+    simulation.tick();
+
+    assert!(!simulation.has_component::<InConfigState>());
+    assert!(simulation.has_component::<InGameState>());
+    assert!(simulation.has_component::<LocalEntity>());
+
+    //
+    // OVERWORLD 1
+    //
+
+    simulation.receive_packet(make_basic_login_packet(
+        DimensionType::new_raw(0), // overworld
+        ResourceLocation::new("azalea:overworld"),
+    ));
+    simulation.tick();
+
+    assert_eq!(
+        *simulation.component::<InstanceName>(),
+        ResourceLocation::new("azalea:overworld"),
+        "InstanceName should be azalea:overworld after setting dimension to that"
+    );
+
+    simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), (384 + 64) / 16));
+    simulation.tick();
+    // make sure the chunk exists
+    simulation
+        .chunk(ChunkPos::new(0, 0))
+        .expect("chunk should exist");
+
+    //
+    // OVERWORLD 2
+    //
+
+    simulation.receive_packet(ClientboundStartConfiguration);
+    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(256)),
+                ("min_y".into(), NbtTag::Int(0)),
+            ])),
+        )]
+        .into_iter()
+        .collect(),
+    });
+    simulation.receive_packet(ClientboundFinishConfiguration);
+    simulation.receive_packet(make_basic_login_or_respawn_packet(
+        DimensionType::new_raw(0),
+        ResourceLocation::new("azalea:overworld"),
+    ));
+    simulation.tick();
+
+    assert!(
+        simulation.chunk(ChunkPos::new(0, 0)).is_none(),
+        "chunk should not exist immediately after changing dimensions"
+    );
+    assert_eq!(
+        *simulation.component::<InstanceName>(),
+        ResourceLocation::new("azalea:overworld"),
+        "InstanceName should still be azalea:overworld after changing dimensions to that"
+    );
+    assert_eq!(
+        simulation
+            .component::<InstanceHolder>()
+            .instance
+            .read()
+            .chunks
+            .height,
+        256
+    );
+
+    simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), 256 / 16));
+    simulation.tick();
+    // make sure the chunk exists
+    simulation
+        .chunk(ChunkPos::new(0, 0))
+        .expect("chunk should exist");
+    simulation.receive_packet(make_basic_login_or_respawn_packet(
+        DimensionType::new_raw(2), // nether
+        ResourceLocation::new("minecraft:nether"),
+    ));
+    simulation.tick();
+}
diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
index d8357a95..9e97a7a7 100644
--- a/azalea-world/src/chunk_storage.rs
+++ b/azalea-world/src/chunk_storage.rs
@@ -26,7 +26,7 @@ const SECTION_HEIGHT: u32 = 16;
 pub struct PartialChunkStorage {
     /// The center of the view, i.e. the chunk the player is currently in.
     view_center: ChunkPos,
-    chunk_radius: u32,
+    pub(crate) chunk_radius: u32,
     view_range: u32,
     // chunks is a list of size chunk_radius * chunk_radius
     chunks: Box<[Option<Arc<RwLock<Chunk>>>]>,
diff --git a/azalea-world/src/world.rs b/azalea-world/src/world.rs
index a650f152..4b2b2a19 100644
--- a/azalea-world/src/world.rs
+++ b/azalea-world/src/world.rs
@@ -41,6 +41,12 @@ impl PartialInstance {
             entity_infos: PartialEntityInfos::new(owner_entity),
         }
     }
+
+    /// Clears the internal references to chunks in the PartialInstance and
+    /// resets the view center.
+    pub fn reset(&mut self) {
+        self.chunks = PartialChunkStorage::new(self.chunks.chunk_radius);
+    }
 }
 
 /// An entity ID used by Minecraft.