Skip to content

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

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

Closed
wants to merge 9 commits into from
150 changes: 71 additions & 79 deletions crates/bevy_ecs/src/world/entity_ref.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use crate::{
archetype::{Archetype, ArchetypeId, Archetypes},
archetype::{Archetype, ArchetypeId},
bundle::{
Bundle, BundleEffect, BundleFromComponents, BundleId, BundleInfo, BundleInserter,
DynamicBundle, InsertMode,
Expand All @@ -10,8 +10,7 @@ use crate::{
StorageType,
},
entity::{
ContainsEntity, Entities, Entity, EntityCloner, EntityClonerBuilder, EntityEquivalent,
EntityLocation,
ContainsEntity, Entity, EntityCloner, EntityClonerBuilder, EntityEquivalent, EntityLocation,
},
event::Event,
observer::Observer,
Expand Down Expand Up @@ -2032,10 +2031,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 @@ -2061,46 +2058,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.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The notes about DROP may have been in the "Safety" section because it would be unsound to drop the values if they've already been moved out.

On the other hand, having the component data be valid is the normal state. The fact that this doesn't drop when false is just an extra behavior used to justify the safety requirements for the take_component call.

So, yup, moving the notes on DROP out of the "safety" section is an improvement! (Sorry for the noise; the safety logic here is a little subtle and I wanted to write down my train of thought when reviewing.)

///
/// 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 @@ -2122,7 +2111,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 @@ -2137,7 +2126,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 @@ -2159,40 +2148,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, caller: MaybeLocation) -> EntityLocation {
unsafe fn remove_bundle(&mut self, bundle: BundleId, caller: MaybeLocation) {
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 @@ -2236,21 +2227,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 @@ -2274,13 +2255,15 @@ impl<'w> EntityWorldMut<'w> {
let mut registrator = unsafe {
ComponentsRegistrator::new(&mut self.world.components, &mut self.world.component_ids)
};
let bundle_info = self
let bundle_id = self
.world
.bundles
.register_info::<T>(&mut registrator, storages);

// SAFETY: the `BundleInfo` is initialized above
self.location = unsafe { self.remove_bundle(bundle_info, caller) };
// SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
unsafe {
self.remove_bundle(bundle_id, caller);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not a big deal, but if you put the ; outside the unsafe block then rustfmt might put it all on one line. (Here and below.)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Slightly different semantics though; semicolon-outside is used to return something from the block to a variable.

I don't know how much we actually care about that, but we seem to follow it in this file at least.

}
self.world.flush();
self.update_location();
self
Expand Down Expand Up @@ -2310,8 +2293,10 @@ impl<'w> EntityWorldMut<'w> {

let bundle_id = bundles.register_contributed_bundle_info::<T>(&mut registrator, storages);

// SAFETY: the dynamic `BundleInfo` is initialized above
self.location = unsafe { self.remove_bundle(bundle_id, caller) };
// SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
unsafe {
self.remove_bundle(bundle_id, caller);
}
self.world.flush();
self.update_location();
self
Expand Down Expand Up @@ -2358,8 +2343,10 @@ impl<'w> EntityWorldMut<'w> {
.bundles
.init_dynamic_info(&mut self.world.storages, &registrator, to_remove);

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

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

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

unsafe {
self.remove_bundle(bundle_id, MaybeLocation::caller());
}
self.world.flush();
self.update_location();
self
Expand Down Expand Up @@ -2449,8 +2439,10 @@ impl<'w> EntityWorldMut<'w> {
component_ids.as_slice(),
);

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