use UnsafeWorldCell for System and ParamState

Author: jakobhellermann

crates/bevy_ecs/macros/src/lib.rs

@@ -294,7 +294,7 @@ pub fn impl_param_set(_input: TokenStream) -> TokenStream {
                 unsafe fn get_param<'w, 's>(
                     state: &'s mut Self::State,
                     system_meta: &SystemMeta,
-                    world: &'w World,
+                    world: UnsafeWorldCell<'w>,
                     change_tick: u32,
                 ) -> Self::Item<'w, 's> {
                     ParamSet {
@@ -498,7 +498,7 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
                 unsafe fn get_param<'w2, 's2>(
                     state: &'s2 mut Self::State,
                     system_meta: &#path::system::SystemMeta,
-                    world: &'w2 #path::world::World,
+                    world: #path::world::unsafe_world_cell::UnsafeWorldCell<'w2>,
                     change_tick: u32,
                 ) -> Self::Item<'w2, 's2> {
                     let (#(#tuple_patterns,)*) = <

crates/bevy_ecs/src/removal_detection.rs

@@ -8,7 +8,7 @@ use crate::{
     prelude::Local,
     storage::SparseSet,
     system::{ReadOnlySystemParam, SystemMeta, SystemParam},
-    world::World,
+    world::{unsafe_world_cell::UnsafeWorldCell, World},
 };
 
 use std::{
@@ -178,9 +178,10 @@ unsafe impl<'a> SystemParam for &'a RemovedComponentEvents {
     unsafe fn get_param<'w, 's>(
         _state: &'s mut Self::State,
         _system_meta: &SystemMeta,
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _change_tick: u32,
     ) -> Self::Item<'w, 's> {
-        world.removed_components()
+        // SAFETY: only metadata is accessed
+        unsafe { world.world_metadata() }.removed_components()
     }
 }

crates/bevy_ecs/src/schedule/executor/multi_threaded.rs

@@ -18,7 +18,7 @@ use crate::{
         is_apply_system_buffers, BoxedCondition, ExecutorKind, SystemExecutor, SystemSchedule,
     },
     system::BoxedSystem,
-    world::World,
+    world::{unsafe_world_cell::UnsafeWorldCell, World},
 };
 
 use crate as bevy_ecs;
@@ -281,6 +281,7 @@ impl MultiThreadedExecutor {
             // SAFETY: No exclusive system is running.
             // Therefore, there is no existing mutable reference to the world.
             let world = unsafe { &*cell.get() };
+
             if !self.can_run(system_index, system, conditions, world) {
                 // NOTE: exclusive systems with ambiguities are susceptible to
                 // being significantly displaced here (compared to single-threaded order)
@@ -309,6 +310,9 @@ impl MultiThreadedExecutor {
                 break;
             }
 
+            // TODO: do this earlier?
+            let world = world.as_unsafe_world_cell_migration_internal();
+
             // SAFETY: No other reference to this system exists.
             unsafe {
                 self.spawn_system_task(scope, system_index, systems, world);
@@ -425,7 +429,7 @@ impl MultiThreadedExecutor {
         scope: &Scope<'_, 'scope, ()>,
         system_index: usize,
         systems: &'scope [SyncUnsafeCell<BoxedSystem>],
-        world: &'scope World,
+        world: UnsafeWorldCell<'scope>,
     ) {
         // SAFETY: this system is not running, no other reference exists
         let system = unsafe { &mut *systems[system_index].get() };
@@ -604,6 +608,7 @@ fn apply_system_buffers(
     }
 }
 
+// TODO: this is unsafe
 fn evaluate_and_fold_conditions(conditions: &mut [BoxedCondition], world: &World) -> bool {
     // not short-circuiting is intentional
     #[allow(clippy::unnecessary_fold)]
@@ -613,7 +618,7 @@ fn evaluate_and_fold_conditions(conditions: &mut [BoxedCondition], world: &World
             #[cfg(feature = "trace")]
             let _condition_span = info_span!("condition", name = &*condition.name()).entered();
             // SAFETY: caller ensures system access is compatible
-            unsafe { condition.run_unsafe((), world) }
+            unsafe { condition.run_unsafe((), world.as_unsafe_world_cell_readonly()) }
         })
         .fold(true, |acc, res| acc && res)
 }

crates/bevy_ecs/src/system/exclusive_function_system.rs

@@ -7,7 +7,7 @@ use crate::{
         check_system_change_tick, ExclusiveSystemParam, ExclusiveSystemParamItem, In, InputMarker,
         IntoSystem, System, SystemMeta,
     },
-    world::{World, WorldId},
+    world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId},
 };
 use bevy_ecs_macros::all_tuples;
 use std::{any::TypeId, borrow::Cow, marker::PhantomData};
@@ -95,7 +95,7 @@ where
     }
 
     #[inline]
-    unsafe fn run_unsafe(&mut self, _input: Self::In, _world: &World) -> Self::Out {
+    unsafe fn run_unsafe(&mut self, _input: Self::In, _world: UnsafeWorldCell<'_>) -> Self::Out {
         panic!("Cannot run exclusive systems with a shared World reference");
     }
 

crates/bevy_ecs/src/system/function_system.rs

@@ -5,7 +5,7 @@ use crate::{
     prelude::FromWorld,
     query::{Access, FilteredAccessSet},
     system::{check_system_change_tick, ReadOnlySystemParam, System, SystemParam, SystemParamItem},
-    world::{World, WorldId},
+    world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId},
 };
 use bevy_ecs_macros::all_tuples;
 use std::{any::TypeId, borrow::Cow, marker::PhantomData};
@@ -174,7 +174,7 @@ impl<Param: SystemParam> SystemState<Param> {
         self.validate_world(world);
         self.update_archetypes(world);
         // SAFETY: Param is read-only and doesn't allow mutable access to World. It also matches the World this SystemState was created with.
-        unsafe { self.get_unchecked_manual(world) }
+        unsafe { self.get_unchecked_manual(world.as_unsafe_world_cell_readonly()) }
     }
 
     /// Retrieve the mutable [`SystemParam`] values.
@@ -183,7 +183,7 @@ impl<Param: SystemParam> SystemState<Param> {
         self.validate_world(world);
         self.update_archetypes(world);
         // SAFETY: World is uniquely borrowed and matches the World this SystemState was created with.
-        unsafe { self.get_unchecked_manual(world) }
+        unsafe { self.get_unchecked_manual(world.as_unsafe_world_cell()) }
     }
 
     /// Applies all state queued up for [`SystemParam`] values. For example, this will apply commands queued up
@@ -243,7 +243,7 @@ impl<Param: SystemParam> SystemState<Param> {
         self.validate_world(world);
         let change_tick = world.read_change_tick();
         // SAFETY: Param is read-only and doesn't allow mutable access to World. It also matches the World this SystemState was created with.
-        unsafe { self.fetch(world, change_tick) }
+        unsafe { self.fetch(world.as_unsafe_world_cell_readonly(), change_tick) }
     }
 
     /// Retrieve the mutable [`SystemParam`] values.  This will not update the state's view of the world's archetypes
@@ -261,7 +261,7 @@ impl<Param: SystemParam> SystemState<Param> {
         self.validate_world(world);
         let change_tick = world.change_tick();
         // SAFETY: World is uniquely borrowed and matches the World this SystemState was created with.
-        unsafe { self.fetch(world, change_tick) }
+        unsafe { self.fetch(world.as_unsafe_world_cell(), change_tick) }
     }
 
     /// Retrieve the [`SystemParam`] values. This will not update archetypes automatically.
@@ -273,7 +273,7 @@ impl<Param: SystemParam> SystemState<Param> {
     #[inline]
     pub unsafe fn get_unchecked_manual<'w, 's>(
         &'s mut self,
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
     ) -> SystemParamItem<'w, 's, Param> {
         let change_tick = world.increment_change_tick();
         self.fetch(world, change_tick)
@@ -286,7 +286,7 @@ impl<Param: SystemParam> SystemState<Param> {
     #[inline]
     unsafe fn fetch<'w, 's>(
         &'s mut self,
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         change_tick: u32,
     ) -> SystemParamItem<'w, 's, Param> {
         let param = Param::get_param(&mut self.param_state, &self.meta, world, change_tick);
@@ -459,7 +459,7 @@ where
     }
 
     #[inline]
-    unsafe fn run_unsafe(&mut self, input: Self::In, world: &World) -> Self::Out {
+    unsafe fn run_unsafe(&mut self, input: Self::In, world: UnsafeWorldCell<'_>) -> Self::Out {
         let change_tick = world.increment_change_tick();
 
         // Safety:

crates/bevy_ecs/src/system/system.rs

@@ -3,7 +3,7 @@ use core::fmt::Debug;
 
 use crate::{
     archetype::ArchetypeComponentId, change_detection::MAX_CHANGE_AGE, component::ComponentId,
-    query::Access, world::World,
+    query::Access, world::unsafe_world_cell::UnsafeWorldCell, world::World,
 };
 
 use std::any::TypeId;
@@ -46,17 +46,14 @@ pub trait System: Send + Sync + 'static {
     ///
     /// # Safety
     ///
-    /// This might access world and resources in an unsafe manner. This should only be called in one
-    /// of the following contexts:
-    ///     1. This system is the only system running on the given world across all threads.
-    ///     2. This system only runs in parallel with other systems that do not conflict with the
-    ///        [`System::archetype_component_access()`].
-    unsafe fn run_unsafe(&mut self, input: Self::In, world: &World) -> Self::Out;
+    /// This function may only be called with a [`UnsafeWorldCell`] that can be used for
+    /// everything listed in [`System::archetype_component_access`].
+    unsafe fn run_unsafe(&mut self, input: Self::In, world: UnsafeWorldCell<'_>) -> Self::Out;
     /// Runs the system with the given input in the world.
     fn run(&mut self, input: Self::In, world: &mut World) -> Self::Out {
         self.update_archetype_component_access(world);
         // SAFETY: world and resources are exclusively borrowed
-        unsafe { self.run_unsafe(input, world) }
+        unsafe { self.run_unsafe(input, world.as_unsafe_world_cell()) }
     }
     fn apply_buffers(&mut self, world: &mut World);
     /// Initialize the system.