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/schedule/executor_parallel.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use crate as bevy_ecs;
+use crate::world::unsafe_world_cell::UnsafeWorldCell;
 use crate::{
     archetype::ArchetypeComponentId,
     query::Access,
@@ -138,7 +139,12 @@ impl ParallelSystemExecutor for ParallelExecutor {
             }
         }
 
-        let thread_executor = world.get_resource::<MainThreadExecutor>().map(|e| &*e.0);
+        let thread_executor = world
+            .get_resource::<MainThreadExecutor>()
+            .map(|e| Arc::clone(&e.0));
+        let thread_executor = thread_executor.as_deref();
+
+        let world = world.as_unsafe_world_cell();
 
         ComputeTaskPool::init(TaskPool::default).scope_with_executor(
             false,
@@ -188,7 +194,7 @@ impl ParallelExecutor {
         &mut self,
         scope: &Scope<'_, 'scope, ()>,
         systems: &'scope mut [SystemContainer],
-        world: &'scope World,
+        world: UnsafeWorldCell<'scope>,
     ) {
         // These are used as a part of a unit test.
         #[cfg(test)]

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

@@ -14,7 +14,7 @@ use crate::{
         is_apply_system_buffers, BoxedCondition, ExecutorKind, SystemExecutor, SystemSchedule,
     },
     system::BoxedSystem,
-    world::World,
+    world::{unsafe_world_cell::UnsafeWorldCell, World},
 };
 
 /// A funky borrow split of [`SystemSchedule`] required by the [`MultiThreadedExecutor`].
@@ -252,6 +252,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)
@@ -280,6 +281,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);
@@ -400,7 +404,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() };
@@ -560,6 +564,7 @@ fn apply_system_buffers(
     unapplied_systems.clear();
 }
 
+// TODO: this is unsafe
 fn evaluate_and_fold_conditions(conditions: &mut [BoxedCondition], world: &World) -> bool {
     // not short-circuiting is intentional
     #[allow(clippy::unnecessary_fold)]
@@ -569,7 +574,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/commands/parallel_scope.rs

@@ -6,6 +6,7 @@ use crate::{
     entity::Entities,
     prelude::World,
     system::{SystemMeta, SystemParam},
+    world::unsafe_world_cell::UnsafeWorldCell,
 };
 
 use super::{CommandQueue, Commands};
@@ -70,7 +71,7 @@ unsafe impl SystemParam for ParallelCommands<'_, '_> {
     unsafe fn get_param<'w, 's>(
         state: &'s mut Self::State,
         _: &crate::system::SystemMeta,
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _: u32,
     ) -> Self::Item<'w, 's> {
         ParallelCommands {

crates/bevy_ecs/src/system/exclusive_function_system.rs

@@ -8,7 +8,7 @@ use crate::{
         check_system_change_tick, AsSystemLabel, ExclusiveSystemParam, ExclusiveSystemParamItem,
         In, InputMarker, IntoSystem, System, SystemMeta, SystemTypeIdLabel,
     },
-    world::{World, WorldId},
+    world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId},
 };
 use bevy_ecs_macros::all_tuples;
 use std::{any::TypeId, borrow::Cow, marker::PhantomData};
@@ -96,7 +96,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

@@ -6,7 +6,7 @@ use crate::{
     query::{Access, FilteredAccessSet},
     schedule::{SystemLabel, SystemLabelId},
     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, fmt::Debug, marker::PhantomData};
@@ -173,7 +173,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.
@@ -182,7 +182,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
@@ -242,7 +242,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
@@ -260,7 +260,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.
@@ -272,7 +272,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)
@@ -285,7 +285,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);
@@ -458,7 +458,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

@@ -2,8 +2,12 @@ use bevy_utils::tracing::warn;
 use core::fmt::Debug;
 
 use crate::{
-    archetype::ArchetypeComponentId, change_detection::MAX_CHANGE_AGE, component::ComponentId,
-    query::Access, schedule::SystemLabelId, world::World,
+    archetype::ArchetypeComponentId,
+    change_detection::MAX_CHANGE_AGE,
+    component::ComponentId,
+    query::Access,
+    schedule::SystemLabelId,
+    world::{unsafe_world_cell::UnsafeWorldCell, World},
 };
 
 use std::any::TypeId;
@@ -46,17 +50,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.