From e60249e59d0e8f62eb3fff2222e5cf6c3bc1ba33 Mon Sep 17 00:00:00 2001
From: Joseph <21144246+JoJoJet@users.noreply.github.com>
Date: Thu, 21 Sep 2023 13:57:06 -0700
Subject: [PATCH] Improve codegen for world validation (#9464)

# Objective

Improve code-gen for `QueryState::validate_world` and
`SystemState::validate_world`.

## Solution

* Move panics into separate, non-inlined functions, to reduce the code
size of the outer methods.
* Mark the panicking functions with `#[cold]` to help the compiler
optimize for the happy path.
* Mark the functions with `#[track_caller]` to make debugging easier.

---------

Co-authored-by: James Liu <contact@jamessliu.com>
---
 crates/bevy_ecs/src/query/state.rs            | 16 +++++++++++-----
 crates/bevy_ecs/src/system/function_system.rs | 12 +++++++++++-
 crates/bevy_ecs/src/system/mod.rs             |  2 +-
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs
index 45115c2d540f5..e64c257ec04b5 100644
--- a/crates/bevy_ecs/src/query/state.rs
+++ b/crates/bevy_ecs/src/query/state.rs
@@ -224,12 +224,18 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
     /// Many unsafe query methods require the world to match for soundness. This function is the easiest
     /// way of ensuring that it matches.
     #[inline]
+    #[track_caller]
     pub fn validate_world(&self, world_id: WorldId) {
-        assert!(
-            world_id == self.world_id,
-            "Attempted to use {} with a mismatched World. QueryStates can only be used with the World they were created from.",
-                std::any::type_name::<Self>(),
-        );
+        #[inline(never)]
+        #[track_caller]
+        #[cold]
+        fn panic_mismatched(this: WorldId, other: WorldId) -> ! {
+            panic!("Encountered a mismatched World. This QueryState was created from {this:?}, but a method was called using {other:?}.");
+        }
+
+        if self.world_id != world_id {
+            panic_mismatched(self.world_id, world_id);
+        }
     }
 
     /// Update the current [`QueryState`] with information from the provided [`Archetype`]
diff --git a/crates/bevy_ecs/src/system/function_system.rs b/crates/bevy_ecs/src/system/function_system.rs
index aa11b1a03f3cb..9e5fd21513dc9 100644
--- a/crates/bevy_ecs/src/system/function_system.rs
+++ b/crates/bevy_ecs/src/system/function_system.rs
@@ -232,8 +232,18 @@ impl<Param: SystemParam> SystemState<Param> {
 
     /// Asserts that the [`SystemState`] matches the provided world.
     #[inline]
+    #[track_caller]
     fn validate_world(&self, world_id: WorldId) {
-        assert!(self.matches_world(world_id), "Encountered a mismatched World. A SystemState cannot be used with Worlds other than the one it was created with.");
+        #[inline(never)]
+        #[track_caller]
+        #[cold]
+        fn panic_mismatched(this: WorldId, other: WorldId) -> ! {
+            panic!("Encountered a mismatched World. This SystemState was created from {this:?}, but a method was called using {other:?}.");
+        }
+
+        if !self.matches_world(world_id) {
+            panic_mismatched(self.world_id, world_id);
+        }
     }
 
     /// Updates the state's internal view of the [`World`]'s archetypes. If this is not called before fetching the parameters,
diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs
index ac73d676c9d6e..9079d3a9d4745 100644
--- a/crates/bevy_ecs/src/system/mod.rs
+++ b/crates/bevy_ecs/src/system/mod.rs
@@ -1775,7 +1775,7 @@ mod tests {
     }
 
     #[test]
-    #[should_panic = "Attempted to use bevy_ecs::query::state::QueryState<()> with a mismatched World."]
+    #[should_panic = "Encountered a mismatched World."]
     fn query_validates_world_id() {
         let mut world1 = World::new();
         let world2 = World::new();