Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify usage of EntityWorldMut::move_entity_from_remove by making it a proper method #17360

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 71 additions & 80 deletions crates/bevy_ecs/src/world/entity_ref.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
use crate::{
archetype::{Archetype, ArchetypeId, Archetypes},
archetype::{Archetype, ArchetypeId},
bundle::{Bundle, BundleId, BundleInfo, BundleInserter, DynamicBundle, InsertMode},
change_detection::MutUntyped,
component::{Component, ComponentId, ComponentTicks, Components, Mutable, StorageType},
entity::{
Entities, Entity, EntityBorrow, EntityCloneBuilder, EntityLocation, TrustedEntityBorrow,
},
entity::{Entity, EntityBorrow, EntityCloneBuilder, EntityLocation, TrustedEntityBorrow},
event::Event,
observer::Observer,
query::{Access, ReadOnlyQueryData},
Expand Down Expand Up @@ -1728,10 +1726,8 @@ impl<'w> EntityWorldMut<'w> {
);
}

let archetypes = &mut world.archetypes;
let storages = &mut world.storages;
let components = &mut world.components;
let entities = &mut world.entities;
let removed_components = &mut world.removed_components;

let entity = self.entity;
Expand All @@ -1756,46 +1752,38 @@ impl<'w> EntityWorldMut<'w> {
})
};

#[expect(
clippy::undocumented_unsafe_blocks,
reason = "Needs to be documented; see #17345."
)]
// SAFETY: `new_archetype_id` has a subset of the components in the entity's current archetype
// because it was created by removing a bundle from that archetype.
unsafe {
Self::move_entity_from_remove::<false>(
entity,
&mut self.location,
old_location.archetype_id,
old_location,
entities,
archetypes,
storages,
new_archetype_id,
);
self.move_entity_from_remove::<false>(new_archetype_id);
}
self.world.flush();
self.update_location();
Some(result)
}

/// Moves the entity to a new archetype, where the new archetype
/// was a result of removing components from the entity's current archetype.
///
/// When `DROP` is `true`, removed components will be dropped.
///
/// When `DROP` is `false`, removed components will be forgotten, making someone else
/// responsible for dropping them.
///
/// # Safety
///
/// `new_archetype_id` must have the same or a subset of the components
/// in `old_archetype_id`. Probably more safety stuff too, audit a call to
/// this fn as if the code here was written inline
/// in the entity's current archetype.
///
/// when DROP is true removed components will be dropped otherwise they will be forgotten
// We use a const generic here so that we are less reliant on
// inlining for rustc to optimize out the `match DROP`
unsafe fn move_entity_from_remove<const DROP: bool>(
entity: Entity,
self_location: &mut EntityLocation,
old_archetype_id: ArchetypeId,
old_location: EntityLocation,
entities: &mut Entities,
archetypes: &mut Archetypes,
storages: &mut Storages,
new_archetype_id: ArchetypeId,
) {
// inlining for rustc to optimize out the `if DROP`
unsafe fn move_entity_from_remove<const DROP: bool>(&mut self, new_archetype_id: ArchetypeId) {
let old_archetype_id = self.location.archetype_id;
let old_location = self.location;
let entities = &mut self.world.entities;
let archetypes = &mut self.world.archetypes;
let storages = &mut self.world.storages;

let old_archetype = &mut archetypes[old_archetype_id];
let remove_result = old_archetype.swap_remove(old_location.archetype_row);
// if an entity was moved into this entity's archetype row, update its archetype row
Expand All @@ -1817,7 +1805,7 @@ impl<'w> EntityWorldMut<'w> {
let new_archetype = &mut archetypes[new_archetype_id];

let new_location = if old_table_id == new_archetype.table_id() {
new_archetype.allocate(entity, old_table_row)
new_archetype.allocate(self.entity, old_table_row)
} else {
let (old_table, new_table) = storages
.tables
Expand All @@ -1832,7 +1820,7 @@ impl<'w> EntityWorldMut<'w> {
};

// SAFETY: move_result.new_row is a valid position in new_archetype's table
let new_location = unsafe { new_archetype.allocate(entity, move_result.new_row) };
let new_location = unsafe { new_archetype.allocate(self.entity, move_result.new_row) };

// if an entity was moved into this entity's table row, update its table row
if let Some(swapped_entity) = move_result.swapped_entity {
Expand All @@ -1854,40 +1842,42 @@ impl<'w> EntityWorldMut<'w> {
new_location
};

*self_location = new_location;
self.location = new_location;
// SAFETY: The entity is valid and has been moved to the new location already.
unsafe {
entities.set(entity.index(), new_location);
entities.set(self.entity.index(), new_location);
}
}

/// Remove the components of `bundle` from `entity`.
///
/// # Safety
/// - A `BundleInfo` with the corresponding `BundleId` must have been initialized.
unsafe fn remove_bundle(&mut self, bundle: BundleId) -> EntityLocation {
unsafe fn remove_bundle(&mut self, bundle: BundleId) {
let entity = self.entity;
let world = &mut self.world;
let location = self.location;
// SAFETY: the caller guarantees that the BundleInfo for this id has been initialized.
let bundle_info = world.bundles.get_unchecked(bundle);
let bundle_info = unsafe { world.bundles.get_unchecked(bundle) };

// SAFETY: `archetype_id` exists because it is referenced in `location` which is valid
// and components in `bundle_info` must exist due to this function's safety invariants.
let new_archetype_id = bundle_info
.remove_bundle_from_archetype(
&mut world.archetypes,
&mut world.storages,
&world.components,
&world.observers,
location.archetype_id,
// components from the bundle that are not present on the entity are ignored
true,
)
.expect("intersections should always return a result");
let new_archetype_id = unsafe {
bundle_info
.remove_bundle_from_archetype(
&mut world.archetypes,
&mut world.storages,
&world.components,
&world.observers,
location.archetype_id,
// components from the bundle that are not present on the entity are ignored
true,
)
.expect("intersections should always return a result")
};

if new_archetype_id == location.archetype_id {
return location;
return;
}

// SAFETY: Archetypes and Bundles cannot be mutably aliased through DeferredWorld
Expand Down Expand Up @@ -1929,21 +1919,11 @@ impl<'w> EntityWorldMut<'w> {
}
}

// SAFETY: `new_archetype_id` is a subset of the components in `old_location.archetype_id`
// because it is created by removing a bundle from these components.
let mut new_location = location;
Self::move_entity_from_remove::<true>(
entity,
&mut new_location,
location.archetype_id,
location,
&mut world.entities,
&mut world.archetypes,
&mut world.storages,
new_archetype_id,
);

new_location
// SAFETY: `new_archetype_id` has a subset of the components in the entity's current archetype
// because it was created by removing a bundle from that archetype.
unsafe {
self.move_entity_from_remove::<true>(new_archetype_id);
}
}

/// Removes any components in the [`Bundle`] from the entity.
Expand All @@ -1958,10 +1938,12 @@ impl<'w> EntityWorldMut<'w> {
self.assert_not_despawned();
let storages = &mut self.world.storages;
let components = &mut self.world.components;
let bundle_info = self.world.bundles.register_info::<T>(components, storages);
let bundle_id = self.world.bundles.register_info::<T>(components, storages);

// SAFETY: the `BundleInfo` is initialized above
self.location = unsafe { self.remove_bundle(bundle_info) };
// SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
unsafe {
self.remove_bundle(bundle_id);
}
self.world.flush();
self.update_location();
self
Expand All @@ -1980,8 +1962,10 @@ impl<'w> EntityWorldMut<'w> {

let bundle_id = bundles.register_contributed_bundle_info::<T>(components, storages);

// SAFETY: the dynamic `BundleInfo` is initialized above
self.location = unsafe { self.remove_bundle(bundle_id) };
// SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
unsafe {
self.remove_bundle(bundle_id);
}
self.world.flush();
self.update_location();
self
Expand Down Expand Up @@ -2013,8 +1997,10 @@ impl<'w> EntityWorldMut<'w> {
.collect::<Vec<_>>();
let remove_bundle = self.world.bundles.init_dynamic_info(components, to_remove);

// SAFETY: the `BundleInfo` for the components to remove is initialized above
self.location = unsafe { self.remove_bundle(remove_bundle) };
// SAFETY: the `BundleInfo` for `remove_bundle` is initialized above
unsafe {
self.remove_bundle(remove_bundle);
}
self.world.flush();
self.update_location();
self
Expand All @@ -2037,8 +2023,10 @@ impl<'w> EntityWorldMut<'w> {
.bundles
.init_component_info(components, component_id);

// SAFETY: the `BundleInfo` for this `component_id` is initialized above
self.location = unsafe { self.remove_bundle(bundle_id) };
// SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
unsafe {
self.remove_bundle(bundle_id);
}
self.world.flush();
self.update_location();
self
Expand All @@ -2062,8 +2050,9 @@ impl<'w> EntityWorldMut<'w> {
.init_dynamic_info(components, component_ids);

// SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
unsafe { self.remove_bundle(bundle_id) };

unsafe {
self.remove_bundle(bundle_id);
}
self.world.flush();
self.update_location();
self
Expand All @@ -2084,8 +2073,10 @@ impl<'w> EntityWorldMut<'w> {
.bundles
.init_dynamic_info(components, component_ids.as_slice());

// SAFETY: the `BundleInfo` for this `component_id` is initialized above
self.location = unsafe { self.remove_bundle(bundle_id) };
// SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
unsafe {
self.remove_bundle(bundle_id);
}
self.world.flush();
self.update_location();
self
Expand Down
Loading