Skip to content

Commit

Permalink
fix mutable aliases for a very short time if WorldCell is already b…
Browse files Browse the repository at this point in the history
…orrowed (bevyengine#6639)

# Objective

Consider the test
```rust
let cell = world.cell();
let _value_a = cell.resource_mut::<A>();
let _value_b = cell.resource_mut::<A>();
```

Currently, this will roughly execute

```rust
// first call
let value = unsafe {
    self.world
    .get_non_send_unchecked_mut_with_id(component_id)?
};
return Some(WorldBorrowMut::new(value, archetype_component_id, self.access)))

// second call
let value = unsafe {
    self.world
    .get_non_send_unchecked_mut_with_id(component_id)?
};
return Some(WorldBorrowMut::new(value, archetype_component_id, self.access)))
```
where `WorldBorrowMut::new` will panic if the resource is already borrowed.

This means, that `_value_a` will be created, the access checked (OK), then `value_b` will be created, and the access checked (`panic`).
For a moment, both `_value_a` and `_value_b` existed as `&mut T` to the same location, which is insta-UB as far as I understand it.

## Solution
Flip the order so that `WorldBorrowMut::new` first checks the access, _then_ fetches creates the value. To do that, we pass a `impl FnOnce() -> Mut<T>` instead of the `Mut<T>` directly:

```rust
let get_value = || unsafe {
    self.world
    .get_non_send_unchecked_mut_with_id(component_id)?
};
return Some(WorldBorrowMut::new(get_value, archetype_component_id, self.access)))
```
  • Loading branch information
jakobhellermann authored and ItsDoot committed Feb 1, 2023
1 parent 4e7051e commit 3526b70
Showing 1 changed file with 24 additions and 28 deletions.
52 changes: 24 additions & 28 deletions crates/bevy_ecs/src/world/world_cell.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,21 +88,22 @@ pub struct WorldBorrow<'w, T> {
}

impl<'w, T> WorldBorrow<'w, T> {
fn new(
value: &'w T,
fn try_new(
value: impl FnOnce() -> Option<&'w T>,
archetype_component_id: ArchetypeComponentId,
access: Rc<RefCell<ArchetypeComponentAccess>>,
) -> Self {
) -> Option<Self> {
assert!(
access.borrow_mut().read(archetype_component_id),
"Attempted to immutably access {}, but it is already mutably borrowed",
std::any::type_name::<T>(),
);
Self {
let value = value()?;
Some(Self {
value,
archetype_component_id,
access,
}
})
}
}

Expand All @@ -129,21 +130,22 @@ pub struct WorldBorrowMut<'w, T> {
}

impl<'w, T> WorldBorrowMut<'w, T> {
fn new(
value: Mut<'w, T>,
fn try_new(
value: impl FnOnce() -> Option<Mut<'w, T>>,
archetype_component_id: ArchetypeComponentId,
access: Rc<RefCell<ArchetypeComponentAccess>>,
) -> Self {
) -> Option<Self> {
assert!(
access.borrow_mut().write(archetype_component_id),
"Attempted to mutably access {}, but it is already mutably borrowed",
std::any::type_name::<T>(),
);
Self {
let value = value()?;
Some(Self {
value,
archetype_component_id,
access,
}
})
}
}

Expand Down Expand Up @@ -189,12 +191,12 @@ impl<'w> WorldCell<'w> {
let archetype_component_id = self
.world
.get_resource_archetype_component_id(component_id)?;
Some(WorldBorrow::new(
WorldBorrow::try_new(
// SAFETY: ComponentId matches TypeId
unsafe { self.world.get_resource_with_id(component_id)? },
|| unsafe { self.world.get_resource_with_id(component_id) },
archetype_component_id,
self.access.clone(),
))
)
}

/// Gets a reference to the resource of the given type
Expand Down Expand Up @@ -222,15 +224,12 @@ impl<'w> WorldCell<'w> {
let archetype_component_id = self
.world
.get_resource_archetype_component_id(component_id)?;
Some(WorldBorrowMut::new(
WorldBorrowMut::try_new(
// SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
unsafe {
self.world
.get_resource_unchecked_mut_with_id(component_id)?
},
|| unsafe { self.world.get_resource_unchecked_mut_with_id(component_id) },
archetype_component_id,
self.access.clone(),
))
)
}

/// Gets a mutable reference to the resource of the given type
Expand Down Expand Up @@ -258,12 +257,12 @@ impl<'w> WorldCell<'w> {
let archetype_component_id = self
.world
.get_resource_archetype_component_id(component_id)?;
Some(WorldBorrow::new(
WorldBorrow::try_new(
// SAFETY: ComponentId matches TypeId
unsafe { self.world.get_non_send_with_id(component_id)? },
|| unsafe { self.world.get_non_send_with_id(component_id) },
archetype_component_id,
self.access.clone(),
))
)
}

/// Gets an immutable reference to the non-send resource of the given type, if it exists.
Expand Down Expand Up @@ -291,15 +290,12 @@ impl<'w> WorldCell<'w> {
let archetype_component_id = self
.world
.get_resource_archetype_component_id(component_id)?;
Some(WorldBorrowMut::new(
WorldBorrowMut::try_new(
// SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
unsafe {
self.world
.get_non_send_unchecked_mut_with_id(component_id)?
},
|| unsafe { self.world.get_non_send_unchecked_mut_with_id(component_id) },
archetype_component_id,
self.access.clone(),
))
)
}

/// Gets a mutable reference to the non-send resource of the given type, if it exists.
Expand Down

0 comments on commit 3526b70

Please sign in to comment.