From 608ccb8e54454460a1b41a456f99b16d70fb5913 Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Sun, 12 Jan 2025 05:07:53 +0000
Subject: [PATCH] fix panic on bot disconnect

---
 azalea-entity/src/plugin/mod.rs |  8 ++++--
 azalea-physics/src/fluids.rs    |  2 +-
 azalea-physics/src/lib.rs       |  6 ++---
 azalea-physics/tests/physics.rs | 46 +++++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/azalea-entity/src/plugin/mod.rs b/azalea-entity/src/plugin/mod.rs
index 52178c7bc..e8889637e 100644
--- a/azalea-entity/src/plugin/mod.rs
+++ b/azalea-entity/src/plugin/mod.rs
@@ -4,7 +4,10 @@ mod relative_updates;
 use std::collections::HashSet;
 
 use azalea_block::{fluid_state::FluidKind, BlockState};
-use azalea_core::position::{BlockPos, ChunkPos, Vec3};
+use azalea_core::{
+    position::{BlockPos, ChunkPos, Vec3},
+    tick::GameTick,
+};
 use azalea_world::{InstanceContainer, InstanceName, MinecraftEntityId};
 use bevy_app::{App, Plugin, PreUpdate, Update};
 use bevy_ecs::prelude::*;
@@ -59,7 +62,7 @@ impl Plugin for EntityPlugin {
             ),
         )
         .add_systems(Update, update_bounding_box)
-        .add_systems(PreUpdate, update_in_loaded_chunk)
+        .add_systems(GameTick, update_in_loaded_chunk)
         .init_resource::<EntityUuidIndex>();
     }
 }
@@ -215,6 +218,7 @@ pub fn update_in_loaded_chunk(
     for (entity, instance_name, position) in &query {
         let player_chunk_pos = ChunkPos::from(position);
         let Some(instance_lock) = instance_container.get(instance_name) else {
+            commands.entity(entity).remove::<InLoadedChunk>();
             continue;
         };
 
diff --git a/azalea-physics/src/fluids.rs b/azalea-physics/src/fluids.rs
index eca4266e4..6e89cdea6 100644
--- a/azalea-physics/src/fluids.rs
+++ b/azalea-physics/src/fluids.rs
@@ -23,7 +23,7 @@ pub fn update_in_water_state_and_do_fluid_pushing(
     for (mut physics, position, instance_name) in &mut query {
         let world_lock = instance_container
             .get(instance_name)
-            .expect("All entities should be in a valid world");
+            .expect("All entities with InLoadedChunk should be in a valid world");
         let world = world_lock.read();
 
         physics.water_fluid_height = 0.;
diff --git a/azalea-physics/src/lib.rs b/azalea-physics/src/lib.rs
index 6ea4e9461..c626dcdf7 100644
--- a/azalea-physics/src/lib.rs
+++ b/azalea-physics/src/lib.rs
@@ -109,10 +109,8 @@ pub fn ai_step(
                     if !physics.is_in_lava()
                         || physics.on_ground() && fluid_height <= fluid_jump_threshold
                     {
-                        if physics.on_ground()
-                            || in_water
-                                && fluid_height <= fluid_jump_threshold
-                                && physics.no_jump_delay == 0
+                        if (physics.on_ground() || in_water && fluid_height <= fluid_jump_threshold)
+                            && physics.no_jump_delay == 0
                         {
                             jump_from_ground(
                                 &mut physics,
diff --git a/azalea-physics/tests/physics.rs b/azalea-physics/tests/physics.rs
index c7e850063..d2986ff09 100644
--- a/azalea-physics/tests/physics.rs
+++ b/azalea-physics/tests/physics.rs
@@ -363,3 +363,49 @@ fn test_negative_coordinates_weird_wall_collision() {
     let entity_pos = app.world_mut().get::<Position>(entity).unwrap();
     assert_eq!(entity_pos.y, 70.5);
 }
+
+#[test]
+fn spawn_and_unload_world() {
+    let mut app = make_test_app();
+    let world_lock = app.world_mut().resource_mut::<InstanceContainer>().insert(
+        ResourceLocation::new("minecraft:overworld"),
+        384,
+        -64,
+    );
+    let mut partial_world = PartialInstance::default();
+
+    partial_world.chunks.set(
+        &ChunkPos { x: -1, z: -1 },
+        Some(Chunk::default()),
+        &mut world_lock.write().chunks,
+    );
+    let _entity = app
+        .world_mut()
+        .spawn((
+            EntityBundle::new(
+                Uuid::nil(),
+                Vec3 {
+                    x: -7.5,
+                    y: 73.,
+                    z: -7.5,
+                },
+                azalea_registry::EntityKind::Player,
+                ResourceLocation::new("minecraft:overworld"),
+            ),
+            MinecraftEntityId(0),
+            LocalEntity,
+        ))
+        .id();
+
+    // do a tick
+    app.world_mut().run_schedule(GameTick);
+    app.update();
+
+    // now unload the partial_world and world_lock
+    drop(partial_world);
+    drop(world_lock);
+
+    // do another tick
+    app.world_mut().run_schedule(GameTick);
+    app.update();
+}