Skip to content

Commit

Permalink
add UnsafeWorldCell abstraction (bevyengine#6404)
Browse files Browse the repository at this point in the history
alternative to bevyengine#5922, implements bevyengine#5956 
builds on top of bevyengine#6402

# Objective

bevyengine#5956 goes into more detail, but the TLDR is:
- bevy systems ensure disjoint accesses to resources and components, and for that to work there are methods `World::get_resource_unchecked_mut(&self)`, ..., `EntityRef::get_mut_unchecked(&self)` etc.
- we don't have these unchecked methods for `by_id` variants, so third-party crate authors cannot build their own safe disjoint-access abstractions with these
- having `_unchecked_mut` methods is not great, because in their presence safe code can accidentally violate subtle invariants. Having to go through `world.as_unsafe_world_cell().unsafe_method()` forces you to stop and think about what you want to write in your `// SAFETY` comment.

The alternative is to keep exposing `_unchecked_mut` variants for every operation that we want third-party crates to build upon, but we'd prefer to avoid using these methods alltogether: bevyengine#5922 (comment)

Also, this is something that **cannot be implemented outside of bevy**, so having either this PR or bevyengine#5922 as an escape hatch with lots of discouraging comments would be great.

## Solution

- add `UnsafeWorldCell` with `unsafe fn get_resource(&self)`, `unsafe fn get_resource_mut(&self)`
- add `fn World::as_unsafe_world_cell(&mut self) -> UnsafeWorldCell<'_>` (and `as_unsafe_world_cell_readonly(&self)`)
- add `UnsafeWorldCellEntityRef` with `unsafe fn get`, `unsafe fn get_mut` and the other utilities on `EntityRef` (no methods for spawning, despawning, insertion)
- use the `UnsafeWorldCell` abstraction in `ReflectComponent`, `ReflectResource` and `ReflectAsset`, so these APIs are easier to reason about
- remove `World::get_resource_mut_unchecked`, `EntityRef::get_mut_unchecked` and use `unsafe { world.as_unsafe_world_cell().get_mut() }` and `unsafe { world.as_unsafe_world_cell().get_entity(entity)?.get_mut() }` instead

This PR does **not** make use of `UnsafeWorldCell` for anywhere else in `bevy_ecs` such as `SystemParam` or `Query`. That is a much larger change, and I am convinced that having `UnsafeWorldCell` is already useful for third-party crates.

Implemented API:

```rust
struct World { .. }
impl World {
  fn as_unsafe_world_cell(&self) -> UnsafeWorldCell<'_>;
}

struct UnsafeWorldCell<'w>(&'w World);
impl<'w> UnsafeWorldCell {
  unsafe fn world(&self) -> &World;

  fn get_entity(&self) -> UnsafeWorldCellEntityRef<'w>; // returns 'w which is `'self` of the `World::as_unsafe_world_cell(&'w self)`

  unsafe fn get_resource<T>(&self) -> Option<&'w T>;
  unsafe fn get_resource_by_id(&self, ComponentId) -> Option<&'w T>;
  unsafe fn get_resource_mut<T>(&self) -> Option<Mut<'w, T>>;
  unsafe fn get_resource_mut_by_id(&self) -> Option<MutUntyped<'w>>;
  unsafe fn get_non_send_resource<T>(&self) -> Option<&'w T>;
  unsafe fn get_non_send_resource_mut<T>(&self) -> Option<Mut<'w, T>>>;

  // not included: remove, remove_resource, despawn, anything that might change archetypes
}

struct UnsafeWorldCellEntityRef<'w> { .. }
impl UnsafeWorldCellEntityRef<'w> {
  unsafe fn get<T>(&self, Entity) -> Option<&'w T>;
  unsafe fn get_by_id(&self, Entity, ComponentId) -> Option<Ptr<'w>>;
  unsafe fn get_mut<T>(&self, Entity) -> Option<Mut<'w, T>>;
  unsafe fn get_mut_by_id(&self, Entity, ComponentId) -> Option<MutUntyped<'w>>;
  unsafe fn get_change_ticks<T>(&self, Entity) -> Option<Mut<'w, T>>;
  // fn id, archetype, contains, contains_id, containts_type_id
}
```

<details>
<summary>UnsafeWorldCell docs</summary>

Variant of the [`World`] where resource and component accesses takes a `&World`, and the responsibility to avoid
aliasing violations are given to the caller instead of being checked at compile-time by rust's unique XOR shared rule.

### Rationale
In rust, having a `&mut World` means that there are absolutely no other references to the safe world alive at the same time,
without exceptions. Not even unsafe code can change this.

But there are situations where careful shared mutable access through a type is possible and safe. For this, rust provides the [`UnsafeCell`](std::cell::UnsafeCell)
escape hatch, which allows you to get a `*mut T` from a `&UnsafeCell<T>` and around which safe abstractions can be built.

Access to resources and components can be done uniquely using [`World::resource_mut`] and [`World::entity_mut`], and shared using [`World::resource`] and [`World::entity`].
These methods use lifetimes to check at compile time that no aliasing rules are being broken.

This alone is not enough to implement bevy systems where multiple systems can access *disjoint* parts of the world concurrently. For this, bevy stores all values of
resources and components (and [`ComponentTicks`](crate::component::ComponentTicks)) in [`UnsafeCell`](std::cell::UnsafeCell)s, and carefully validates disjoint access patterns using
APIs like [`System::component_access`](crate::system::System::component_access).

A system then can be executed using [`System::run_unsafe`](crate::system::System::run_unsafe) with a `&World` and use methods with interior mutability to access resource values.
access resource values.

### Example Usage

[`UnsafeWorldCell`] can be used as a building block for writing APIs that safely allow disjoint access into the world.
In the following example, the world is split into a resource access half and a component access half, where each one can
safely hand out mutable references.

```rust
use bevy_ecs::world::World;
use bevy_ecs::change_detection::Mut;
use bevy_ecs::system::Resource;
use bevy_ecs::world::unsafe_world_cell_world::UnsafeWorldCell;

// INVARIANT: existance of this struct means that users of it are the only ones being able to access resources in the world
struct OnlyResourceAccessWorld<'w>(UnsafeWorldCell<'w>);
// INVARIANT: existance of this struct means that users of it are the only ones being able to access components in the world
struct OnlyComponentAccessWorld<'w>(UnsafeWorldCell<'w>);

impl<'w> OnlyResourceAccessWorld<'w> {
    fn get_resource_mut<T: Resource>(&mut self) -> Option<Mut<'w, T>> {
        // SAFETY: resource access is allowed through this UnsafeWorldCell
        unsafe { self.0.get_resource_mut::<T>() }
    }
}
// impl<'w> OnlyComponentAccessWorld<'w> {
//     ...
// }

// the two interior mutable worlds borrow from the `&mut World`, so it cannot be accessed while they are live
fn split_world_access(world: &mut World) -> (OnlyResourceAccessWorld<'_>, OnlyComponentAccessWorld<'_>) {
    let resource_access = OnlyResourceAccessWorld(unsafe { world.as_unsafe_world_cell() });
    let component_access = OnlyComponentAccessWorld(unsafe { world.as_unsafe_world_cell() });
    (resource_access, component_access)
}
```


</details>
  • Loading branch information
jakobhellermann authored and ItsDoot committed Feb 1, 2023
1 parent ca710ab commit 66d222d
Show file tree
Hide file tree
Showing 8 changed files with 774 additions and 335 deletions.
39 changes: 17 additions & 22 deletions crates/bevy_asset/src/reflect.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::any::{Any, TypeId};

use bevy_ecs::world::World;
use bevy_ecs::world::{unsafe_world_cell::UnsafeWorldCell, World};
use bevy_reflect::{FromReflect, FromType, Reflect, Uuid};

use crate::{Asset, Assets, Handle, HandleId, HandleUntyped};
Expand All @@ -18,8 +18,10 @@ pub struct ReflectAsset {
assets_resource_type_id: TypeId,

get: fn(&World, HandleUntyped) -> Option<&dyn Reflect>,
get_mut: fn(&mut World, HandleUntyped) -> Option<&mut dyn Reflect>,
get_unchecked_mut: unsafe fn(&World, HandleUntyped) -> Option<&mut dyn Reflect>,
// SAFETY:
// - may only be called with a [`IteriorMutableWorld`] which can be used to access the corresponding `Assets<T>` resource mutably
// - may only be used to access **at most one** access at once
get_unchecked_mut: unsafe fn(UnsafeWorldCell<'_>, HandleUntyped) -> Option<&mut dyn Reflect>,
add: fn(&mut World, &dyn Reflect) -> HandleUntyped,
set: fn(&mut World, HandleUntyped, &dyn Reflect) -> HandleUntyped,
len: fn(&World) -> usize,
Expand Down Expand Up @@ -54,10 +56,11 @@ impl ReflectAsset {
world: &'w mut World,
handle: HandleUntyped,
) -> Option<&'w mut dyn Reflect> {
(self.get_mut)(world, handle)
// SAFETY: unique world access
unsafe { (self.get_unchecked_mut)(world.as_unsafe_world_cell(), handle) }
}

/// Equivalent of [`Assets::get_mut`], but does not require a mutable reference to the world.
/// Equivalent of [`Assets::get_mut`], but works with an [`UnsafeWorldCell`].
///
/// Only use this method when you have ensured that you are the *only* one with access to the [`Assets`] resource of the asset type.
/// Furthermore, this does *not* allow you to have look up two distinct handles,
Expand All @@ -67,11 +70,12 @@ impl ReflectAsset {
/// # use bevy_asset::{ReflectAsset, HandleUntyped};
/// # use bevy_ecs::prelude::World;
/// # let reflect_asset: ReflectAsset = unimplemented!();
/// # let world: World = unimplemented!();
/// # let mut world: World = unimplemented!();
/// # let handle_1: HandleUntyped = unimplemented!();
/// # let handle_2: HandleUntyped = unimplemented!();
/// let a = unsafe { reflect_asset.get_unchecked_mut(&world, handle_1).unwrap() };
/// let b = unsafe { reflect_asset.get_unchecked_mut(&world, handle_2).unwrap() };
/// let unsafe_world_cell = world.as_unsafe_world_cell();
/// let a = unsafe { reflect_asset.get_unchecked_mut(unsafe_world_cell, handle_1).unwrap() };
/// let b = unsafe { reflect_asset.get_unchecked_mut(unsafe_world_cell, handle_2).unwrap() };
/// // ^ not allowed, two mutable references through the same asset resource, even though the
/// // handles are distinct
///
Expand All @@ -81,12 +85,11 @@ impl ReflectAsset {
/// # Safety
/// This method does not prevent you from having two mutable pointers to the same data,
/// violating Rust's aliasing rules. To avoid this:
/// * Only call this method when you have exclusive access to the world
/// (or use a scheduler that enforces unique access to the `Assets` resource).
/// * Only call this method if you know that the [`UnsafeWorldCell`] may be used to access the corresponding `Assets<T>`
/// * Don't call this method more than once in the same scope.
pub unsafe fn get_unchecked_mut<'w>(
&self,
world: &'w World,
world: UnsafeWorldCell<'w>,
handle: HandleUntyped,
) -> Option<&'w mut dyn Reflect> {
// SAFETY: requirements are deferred to the caller
Expand Down Expand Up @@ -140,18 +143,10 @@ impl<A: Asset + FromReflect> FromType<A> for ReflectAsset {
let asset = assets.get(&handle.typed());
asset.map(|asset| asset as &dyn Reflect)
},
get_mut: |world, handle| {
let assets = world.resource_mut::<Assets<A>>().into_inner();
let asset = assets.get_mut(&handle.typed());
asset.map(|asset| asset as &mut dyn Reflect)
},
get_unchecked_mut: |world, handle| {
let assets = unsafe {
world
.get_resource_unchecked_mut::<Assets<A>>()
.unwrap()
.into_inner()
};
// SAFETY: `get_unchecked_mut` must be callied with `UnsafeWorldCell` having access to `Assets<A>`,
// and must ensure to only have at most one reference to it live at all times.
let assets = unsafe { world.get_resource_mut::<Assets<A>>().unwrap().into_inner() };
let asset = assets.get_mut(&handle.typed());
asset.map(|asset| asset as &mut dyn Reflect)
},
Expand Down
49 changes: 26 additions & 23 deletions crates/bevy_ecs/src/reflect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use crate::{
component::Component,
entity::{Entity, EntityMap, MapEntities, MapEntitiesError},
system::Resource,
world::{FromWorld, World},
world::{unsafe_world_cell::UnsafeWorldCell, FromWorld, World},
};
use bevy_reflect::{
impl_from_reflect_value, impl_reflect_value, FromType, Reflect, ReflectDeserialize,
Expand Down Expand Up @@ -52,7 +52,10 @@ pub struct ReflectComponentFns {
/// Function pointer implementing [`ReflectComponent::reflect()`].
pub reflect: fn(&World, Entity) -> Option<&dyn Reflect>,
/// Function pointer implementing [`ReflectComponent::reflect_mut()`].
pub reflect_mut: unsafe fn(&World, Entity) -> Option<Mut<dyn Reflect>>,
///
/// # Safety
/// The function may only be called with an [`UnsafeWorldCell`] that can be used to mutably access the relevant component on the given entity.
pub reflect_mut: unsafe fn(UnsafeWorldCell<'_>, Entity) -> Option<Mut<'_, dyn Reflect>>,
/// Function pointer implementing [`ReflectComponent::copy()`].
pub copy: fn(&World, &mut World, Entity, Entity),
}
Expand Down Expand Up @@ -117,20 +120,20 @@ impl ReflectComponent {
entity: Entity,
) -> Option<Mut<'a, dyn Reflect>> {
// SAFETY: unique world access
unsafe { (self.0.reflect_mut)(world, entity) }
unsafe { (self.0.reflect_mut)(world.as_unsafe_world_cell(), entity) }
}

/// # Safety
/// This method does not prevent you from having two mutable pointers to the same data,
/// violating Rust's aliasing rules. To avoid this:
/// * Only call this method in an exclusive system to avoid sharing across threads (or use a
/// scheduler that enforces safe memory access).
/// * Only call this method with a [`UnsafeWorldCell`] that may be used to mutably access the component on the entity `entity`
/// * Don't call this method more than once in the same scope for a given [`Component`].
pub unsafe fn reflect_unchecked_mut<'a>(
&self,
world: &'a World,
world: UnsafeWorldCell<'a>,
entity: Entity,
) -> Option<Mut<'a, dyn Reflect>> {
// SAFETY: safety requirements deferred to caller
(self.0.reflect_mut)(world, entity)
}

Expand Down Expand Up @@ -209,16 +212,14 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
.map(|c| c as &dyn Reflect)
},
reflect_mut: |world, entity| {
// SAFETY: reflect_mut is an unsafe function pointer used by `reflect_unchecked_mut` which promises to never
// produce aliasing mutable references, and reflect_mut, which has mutable world access
// SAFETY: reflect_mut is an unsafe function pointer used by
// 1. `reflect_unchecked_mut` which must be called with an UnsafeWorldCell with access to the the component `C` on the `entity`, and
// 2. `reflect_mut`, which has mutable world access
unsafe {
world
.get_entity(entity)?
.get_unchecked_mut::<C>(world.last_change_tick(), world.read_change_tick())
.map(|c| Mut {
value: c.value as &mut dyn Reflect,
ticks: c.ticks,
})
world.get_entity(entity)?.get_mut::<C>().map(|c| Mut {
value: c.value as &mut dyn Reflect,
ticks: c.ticks,
})
}
},
})
Expand Down Expand Up @@ -265,7 +266,10 @@ pub struct ReflectResourceFns {
/// Function pointer implementing [`ReflectResource::reflect()`].
pub reflect: fn(&World) -> Option<&dyn Reflect>,
/// Function pointer implementing [`ReflectResource::reflect_unchecked_mut()`].
pub reflect_unchecked_mut: unsafe fn(&World) -> Option<Mut<dyn Reflect>>,
///
/// # Safety
/// The function may only be called with an [`UnsafeWorldCell`] that can be used to mutably access the relevant resource.
pub reflect_unchecked_mut: unsafe fn(UnsafeWorldCell<'_>) -> Option<Mut<'_, dyn Reflect>>,
/// Function pointer implementing [`ReflectResource::copy()`].
pub copy: fn(&World, &mut World),
}
Expand Down Expand Up @@ -314,19 +318,18 @@ impl ReflectResource {
/// Gets the value of this [`Resource`] type from the world as a mutable reflected reference.
pub fn reflect_mut<'a>(&self, world: &'a mut World) -> Option<Mut<'a, dyn Reflect>> {
// SAFETY: unique world access
unsafe { (self.0.reflect_unchecked_mut)(world) }
unsafe { (self.0.reflect_unchecked_mut)(world.as_unsafe_world_cell()) }
}

/// # Safety
/// This method does not prevent you from having two mutable pointers to the same data,
/// violating Rust's aliasing rules. To avoid this:
/// * Only call this method in an exclusive system to avoid sharing across threads (or use a
/// scheduler that enforces safe memory access).
/// * Only call this method with an [`UnsafeWorldCell`] which can be used to mutably access the resource.
/// * Don't call this method more than once in the same scope for a given [`Resource`].
pub unsafe fn reflect_unchecked_mut<'a>(
pub unsafe fn reflect_unchecked_mut<'w>(
&self,
world: &'a World,
) -> Option<Mut<'a, dyn Reflect>> {
world: UnsafeWorldCell<'w>,
) -> Option<Mut<'w, dyn Reflect>> {
// SAFETY: caller promises to uphold uniqueness guarantees
(self.0.reflect_unchecked_mut)(world)
}
Expand Down Expand Up @@ -385,7 +388,7 @@ impl<C: Resource + Reflect + FromWorld> FromType<C> for ReflectResource {
// SAFETY: all usages of `reflect_unchecked_mut` guarantee that there is either a single mutable
// reference or multiple immutable ones alive at any given point
unsafe {
world.get_resource_unchecked_mut::<C>().map(|res| Mut {
world.get_resource_mut::<C>().map(|res| Mut {
value: res.value as &mut dyn Reflect,
ticks: res.ticks,
})
Expand Down
3 changes: 2 additions & 1 deletion crates/bevy_ecs/src/system/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1134,6 +1134,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Query<'w, 's, Q, F> {
}
let world = self.world;
let entity_ref = world
.as_unsafe_world_cell_migration_internal()
.get_entity(entity)
.ok_or(QueryComponentError::NoSuchEntity)?;
let component_id = world
Expand All @@ -1150,7 +1151,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Query<'w, 's, Q, F> {
.has_write(archetype_component)
{
entity_ref
.get_unchecked_mut::<T>(self.last_change_tick, self.change_tick)
.get_mut_using_ticks::<T>(self.last_change_tick, self.change_tick)
.ok_or(QueryComponentError::MissingComponent)
} else {
Err(QueryComponentError::MissingWriteAccess)
Expand Down
6 changes: 4 additions & 2 deletions crates/bevy_ecs/src/system/system_param.rs
Original file line number Diff line number Diff line change
Expand Up @@ -541,7 +541,8 @@ unsafe impl<'a, T: Resource> SystemParam for ResMut<'a, T> {
change_tick: u32,
) -> Self::Item<'w, 's> {
let value = world
.get_resource_unchecked_mut_with_id(component_id)
.as_unsafe_world_cell_migration_internal()
.get_resource_mut_with_id(component_id)
.unwrap_or_else(|| {
panic!(
"Resource requested by {} does not exist: {}",
Expand Down Expand Up @@ -578,7 +579,8 @@ unsafe impl<'a, T: Resource> SystemParam for Option<ResMut<'a, T>> {
change_tick: u32,
) -> Self::Item<'w, 's> {
world
.get_resource_unchecked_mut_with_id(component_id)
.as_unsafe_world_cell_migration_internal()
.get_resource_mut_with_id(component_id)
.map(|value| ResMut {
value: value.value,
ticks: TicksMut {
Expand Down
Loading

0 comments on commit 66d222d

Please sign in to comment.