Skip to content

Commit

Permalink
>:(
Browse files Browse the repository at this point in the history
  • Loading branch information
BoxyUwU committed Mar 4, 2022
1 parent b6a647c commit a1cd173
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 41 deletions.
6 changes: 4 additions & 2 deletions crates/bevy_ecs/src/reflect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -103,15 +103,17 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
.entity_mut(destination_entity)
.insert(destination_component);
},
reflect_component: |world, entity| {
reflect_component: |world, entity| unsafe {
world
.get_entity(entity)?
.get::<C>()
// SAFE: lifetimes force correct usage of returned data
.get_unchecked::<C>()
.map(|c| c as &dyn Reflect)
},
reflect_component_mut: |world, entity| unsafe {
world
.get_entity(entity)?
// SAFE: lifetimes force correct usage of returned data
.get_unchecked_mut::<C>(world.last_change_tick(), world.read_change_tick())
.map(|c| ReflectMut {
value: c.value as &mut dyn Reflect,
Expand Down
9 changes: 6 additions & 3 deletions crates/bevy_ecs/src/system/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -728,9 +728,12 @@ where
.archetype_component_access
.has_read(archetype_component)
{
entity_ref
.get::<T>()
.ok_or(QueryComponentError::MissingComponent)
// SAFE: lifetimes enforce correct usage of returned borrow
unsafe {
entity_ref
.get_unchecked::<T>()
.ok_or(QueryComponentError::MissingComponent)
}
} else {
Err(QueryComponentError::MissingReadAccess)
}
Expand Down
70 changes: 36 additions & 34 deletions crates/bevy_ecs/src/world/entity_ref.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,17 +62,26 @@ impl<'w> EntityRef<'w> {
}

#[inline]
pub fn get<T: Component>(&self) -> Option<&'w T> {
pub fn get<T: Component>(&self) -> Option<&'_ T> {
// SAFE: lifetimes enforce correct usage of returned borrow
unsafe { self.get_unchecked::<T>() }
}

/// This allows aliased mutability use after free bugs.
/// # Safety
/// - Must not be a mutable reference to the component.
/// - Must not use the returned reference after the component is removed
#[inline]
pub unsafe fn get_unchecked<T: Component>(&self) -> Option<&'w T> {
// SAFE: entity location is valid and returned component is of type T
unsafe {
get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
.map(|value| &*value.cast::<T>())
}
get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
.map(|value| &*value.cast::<T>())
}

/// This allows aliased mutability and use after free bugs.
/// # Safety
/// This allows aliased mutability. You must make sure this call does not result in multiple
/// mutable references to the same component
/// - Must not be any other references to the component.
/// - Must not use the returned reference after the component is removed.
#[inline]
pub unsafe fn get_unchecked_mut<T: Component>(
&self,
Expand Down Expand Up @@ -145,39 +154,32 @@ impl<'w> EntityMut<'w> {
}

#[inline]
pub fn get<T: Component>(&self) -> Option<&'w T> {
// SAFE: entity location is valid and returned component is of type T
unsafe {
get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
.map(|value| &*value.cast::<T>())
}
pub fn get<T: Component>(&self) -> Option<&'_ T> {
// SAFE: lifetimes enforce correct usage of returned borrow
unsafe { self.get_unchecked::<T>() }
}

#[inline]
pub fn get_mut<T: Component>(&mut self) -> Option<Mut<'w, T>> {
// SAFE: world access is unique, entity location is valid, and returned component is of type
// T
unsafe {
get_component_and_ticks_with_type(
self.world,
TypeId::of::<T>(),
self.entity,
self.location,
)
.map(|(value, ticks)| Mut {
value: &mut *value.cast::<T>(),
ticks: Ticks {
component_ticks: &mut *ticks,
last_change_tick: self.world.last_change_tick(),
change_tick: self.world.change_tick(),
},
})
}
pub fn get_mut<T: Component>(&mut self) -> Option<Mut<'_, T>> {
// SAFE: world access is unique, and lifetimes enforce correct usage of returned borrow
unsafe { self.get_unchecked_mut::<T>() }
}

/// This allows aliased mutability use after free bugs.
/// # Safety
/// - Must not be a mutable reference to the component.
/// - Must not use the returned reference after the component is removed
#[inline]
pub unsafe fn get_unchecked<T: Component>(&self) -> Option<&'w T> {
// SAFE: entity location is valid and returned component is of type T
get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
.map(|value| &*value.cast::<T>())
}

/// This allows aliased mutability and use after free bugs.
/// # Safety
/// This allows aliased mutability. You must make sure this call does not result in multiple
/// mutable references to the same component
/// - Must not be any other references to the component.
/// - Must not use the returned reference after the component is removed.
#[inline]
pub unsafe fn get_unchecked_mut<T: Component>(&self) -> Option<Mut<'w, T>> {
get_component_and_ticks_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
Expand Down
6 changes: 4 additions & 2 deletions crates/bevy_ecs/src/world/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -416,7 +416,8 @@ impl World {
/// ```
#[inline]
pub fn get<T: Component>(&self, entity: Entity) -> Option<&T> {
self.get_entity(entity)?.get()
// SAFE: lifetimes enforce correct usage of returned borrow
unsafe { self.get_entity(entity)?.get_unchecked::<T>() }
}

/// Retrieves a mutable reference to the given `entity`'s [Component] of the given type.
Expand All @@ -439,7 +440,8 @@ impl World {
/// ```
#[inline]
pub fn get_mut<T: Component>(&mut self, entity: Entity) -> Option<Mut<T>> {
self.get_entity_mut(entity)?.get_mut()
// SAFE: lifetimes enforce correct usage of returned borrow
unsafe { self.get_entity_mut(entity)?.get_unchecked_mut::<T>() }
}

/// Despawns the given `entity`, if it exists. This will also remove all of the entity's
Expand Down

0 comments on commit a1cd173

Please sign in to comment.