Add BundleRemover (#18521)

# Objective

It has long been a todo item in the ecs to create a `BundleRemover`
alongside the inserter, spawner, etc.

This is an uncontroversial first step of #18514.

## Solution

Move existing code from complex helper functions to one generalized
`BundleRemover`.

## Testing

Existing tests.

crates/bevy_ecs/src/bundle.rs

@@ -20,7 +20,10 @@ use crate::{
     query::DebugCheckedUnwrap,
     relationship::RelationshipHookMode,
     storage::{SparseSetIndex, SparseSets, Storages, Table, TableRow},
-    world::{unsafe_world_cell::UnsafeWorldCell, EntityWorldMut, ON_ADD, ON_INSERT, ON_REPLACE},
+    world::{
+        unsafe_world_cell::UnsafeWorldCell, EntityWorldMut, ON_ADD, ON_INSERT, ON_REMOVE,
+        ON_REPLACE,
+    },
 };
 use alloc::{boxed::Box, vec, vec::Vec};
 use bevy_platform::collections::{HashMap, HashSet};
@@ -1371,6 +1374,274 @@ impl<'w> BundleInserter<'w> {
     }
 }
 
+// SAFETY: We have exclusive world access so our pointers can't be invalidated externally
+pub(crate) struct BundleRemover<'w> {
+    world: UnsafeWorldCell<'w>,
+    bundle_info: ConstNonNull<BundleInfo>,
+    old_and_new_table: Option<(NonNull<Table>, NonNull<Table>)>,
+    old_archetype: NonNull<Archetype>,
+    new_archetype: NonNull<Archetype>,
+}
+
+impl<'w> BundleRemover<'w> {
+    /// Creates a new [`BundleRemover`], if such a remover would do anything.
+    ///
+    /// If `require_all` is true, the [`BundleRemover`] is only created if the entire bundle is present on the archetype.
+    ///
+    /// # Safety
+    /// Caller must ensure that `archetype_id` is valid
+    #[inline]
+    pub(crate) unsafe fn new<T: Bundle>(
+        world: &'w mut World,
+        archetype_id: ArchetypeId,
+        require_all: bool,
+    ) -> Option<Self> {
+        // SAFETY: These come from the same world. `world.components_registrator` can't be used since we borrow other fields too.
+        let mut registrator =
+            unsafe { ComponentsRegistrator::new(&mut world.components, &mut world.component_ids) };
+        let bundle_id = world
+            .bundles
+            .register_info::<T>(&mut registrator, &mut world.storages);
+        // SAFETY: we initialized this bundle_id in `init_info`, and caller ensures archetype is valid.
+        unsafe { Self::new_with_id(world, archetype_id, bundle_id, require_all) }
+    }
+
+    /// Creates a new [`BundleRemover`], if such a remover would do anything.
+    ///
+    /// If `require_all` is true, the [`BundleRemover`] is only created if the entire bundle is present on the archetype.
+    ///
+    /// # Safety
+    /// Caller must ensure that `bundle_id` exists in `world.bundles` and `archetype_id` is valid.
+    #[inline]
+    pub(crate) unsafe fn new_with_id(
+        world: &'w mut World,
+        archetype_id: ArchetypeId,
+        bundle_id: BundleId,
+        require_all: bool,
+    ) -> Option<Self> {
+        let bundle_info = world.bundles.get_unchecked(bundle_id);
+        // SAFETY: Caller ensures archetype and bundle ids are correct.
+        let new_archetype_id = unsafe {
+            bundle_info.remove_bundle_from_archetype(
+                &mut world.archetypes,
+                &mut world.storages,
+                &world.components,
+                &world.observers,
+                archetype_id,
+                !require_all,
+            )?
+        };
+        if new_archetype_id == archetype_id {
+            return None;
+        }
+        let (old_archetype, new_archetype) =
+            world.archetypes.get_2_mut(archetype_id, new_archetype_id);
+
+        let tables = if old_archetype.table_id() == new_archetype.table_id() {
+            None
+        } else {
+            let (old, new) = world
+                .storages
+                .tables
+                .get_2_mut(old_archetype.table_id(), new_archetype.table_id());
+            Some((old.into(), new.into()))
+        };
+
+        Some(Self {
+            bundle_info: bundle_info.into(),
+            new_archetype: new_archetype.into(),
+            old_archetype: old_archetype.into(),
+            old_and_new_table: tables,
+            world: world.as_unsafe_world_cell(),
+        })
+    }
+
+    /// This can be passed to [`remove`](Self::remove) as the `pre_remove` function if you don't want to do anything before removing.
+    pub fn empty_pre_remove(
+        _: &mut SparseSets,
+        _: Option<&mut Table>,
+        _: &Components,
+        _: &[ComponentId],
+    ) -> (bool, ()) {
+        (true, ())
+    }
+
+    /// Performs the removal.
+    ///
+    /// `pre_remove` should return a bool for if the components still need to be dropped.
+    ///
+    /// # Safety
+    /// The `location` must have the same archetype as the remover.
+    #[inline]
+    pub(crate) unsafe fn remove<T: 'static>(
+        &mut self,
+        entity: Entity,
+        location: EntityLocation,
+        caller: MaybeLocation,
+        pre_remove: impl FnOnce(
+            &mut SparseSets,
+            Option<&mut Table>,
+            &Components,
+            &[ComponentId],
+        ) -> (bool, T),
+    ) -> (EntityLocation, T) {
+        // Hooks
+        // SAFETY: all bundle components exist in World
+        unsafe {
+            // SAFETY: We only keep access to archetype/bundle data.
+            let mut deferred_world = self.world.into_deferred();
+            let bundle_components_in_archetype = || {
+                self.bundle_info
+                    .as_ref()
+                    .iter_explicit_components()
+                    .filter(|component_id| self.old_archetype.as_ref().contains(*component_id))
+            };
+            if self.old_archetype.as_ref().has_replace_observer() {
+                deferred_world.trigger_observers(
+                    ON_REPLACE,
+                    entity,
+                    bundle_components_in_archetype(),
+                    caller,
+                );
+            }
+            deferred_world.trigger_on_replace(
+                self.old_archetype.as_ref(),
+                entity,
+                bundle_components_in_archetype(),
+                caller,
+                RelationshipHookMode::Run,
+            );
+            if self.old_archetype.as_ref().has_remove_observer() {
+                deferred_world.trigger_observers(
+                    ON_REMOVE,
+                    entity,
+                    bundle_components_in_archetype(),
+                    caller,
+                );
+            }
+            deferred_world.trigger_on_remove(
+                self.old_archetype.as_ref(),
+                entity,
+                bundle_components_in_archetype(),
+                caller,
+            );
+        }
+
+        // SAFETY: We still have the cell, so this is unique, it doesn't conflict with other references, and we drop it shortly.
+        let world = unsafe { self.world.world_mut() };
+
+        let (needs_drop, pre_remove_result) = pre_remove(
+            &mut world.storages.sparse_sets,
+            self.old_and_new_table
+                .as_ref()
+                // SAFETY: There is no conflicting access for this scope.
+                .map(|(old, _)| unsafe { &mut *old.as_ptr() }),
+            &world.components,
+            self.bundle_info.as_ref().explicit_components(),
+        );
+
+        // Handle sparse set removes
+        for component_id in self.bundle_info.as_ref().iter_explicit_components() {
+            if self.old_archetype.as_ref().contains(component_id) {
+                world.removed_components.send(component_id, entity);
+
+                // Make sure to drop components stored in sparse sets.
+                // Dense components are dropped later in `move_to_and_drop_missing_unchecked`.
+                if let Some(StorageType::SparseSet) =
+                    self.old_archetype.as_ref().get_storage_type(component_id)
+                {
+                    world
+                        .storages
+                        .sparse_sets
+                        .get_mut(component_id)
+                        // Set exists because the component existed on the entity
+                        .unwrap()
+                        // If it was already forgotten, it would not be in the set.
+                        .remove(entity);
+                }
+            }
+        }
+
+        // Handle archetype change
+        let remove_result = self
+            .old_archetype
+            .as_mut()
+            .swap_remove(location.archetype_row);
+        // if an entity was moved into this entity's archetype row, update its archetype row
+        if let Some(swapped_entity) = remove_result.swapped_entity {
+            let swapped_location = world.entities.get(swapped_entity).unwrap();
+
+            world.entities.set(
+                swapped_entity.index(),
+                EntityLocation {
+                    archetype_id: swapped_location.archetype_id,
+                    archetype_row: location.archetype_row,
+                    table_id: swapped_location.table_id,
+                    table_row: swapped_location.table_row,
+                },
+            );
+        }
+
+        // Handle table change
+        let new_location = if let Some((mut old_table, mut new_table)) = self.old_and_new_table {
+            let move_result = if needs_drop {
+                // SAFETY: old_table_row exists
+                unsafe {
+                    old_table
+                        .as_mut()
+                        .move_to_and_drop_missing_unchecked(location.table_row, new_table.as_mut())
+                }
+            } else {
+                // SAFETY: old_table_row exists
+                unsafe {
+                    old_table.as_mut().move_to_and_forget_missing_unchecked(
+                        location.table_row,
+                        new_table.as_mut(),
+                    )
+                }
+            };
+
+            // SAFETY: move_result.new_row is a valid position in new_archetype's table
+            let new_location = unsafe {
+                self.new_archetype
+                    .as_mut()
+                    .allocate(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 {
+                let swapped_location = world.entities.get(swapped_entity).unwrap();
+
+                world.entities.set(
+                    swapped_entity.index(),
+                    EntityLocation {
+                        archetype_id: swapped_location.archetype_id,
+                        archetype_row: swapped_location.archetype_row,
+                        table_id: swapped_location.table_id,
+                        table_row: location.table_row,
+                    },
+                );
+                world.archetypes[swapped_location.archetype_id]
+                    .set_entity_table_row(swapped_location.archetype_row, location.table_row);
+            }
+
+            new_location
+        } else {
+            // The tables are the same
+            self.new_archetype
+                .as_mut()
+                .allocate(entity, location.table_row)
+        };
+
+        // SAFETY: The entity is valid and has been moved to the new location already.
+        unsafe {
+            world.entities.set(entity.index(), new_location);
+        }
+
+        (new_location, pre_remove_result)
+    }
+}
+
 // SAFETY: We have exclusive world access so our pointers can't be invalidated externally
 pub(crate) struct BundleSpawner<'w> {
     world: UnsafeWorldCell<'w>,

crates/bevy_ecs/src/world/entity_ref.rs

@@ -1,8 +1,8 @@
 use crate::{
-    archetype::{Archetype, ArchetypeId, Archetypes},
+    archetype::{Archetype, ArchetypeId},
     bundle::{
-        Bundle, BundleEffect, BundleFromComponents, BundleId, BundleInfo, BundleInserter,
+        Bundle, BundleEffect, BundleFromComponents, BundleInserter, BundleRemover, DynamicBundle,
-        DynamicBundle, InsertMode,
+        InsertMode,
     },
     change_detection::{MaybeLocation, MutUntyped},
     component::{
@@ -10,20 +10,17 @@ use crate::{
         StorageType,
     },
     entity::{
-        ContainsEntity, Entities, Entity, EntityCloner, EntityClonerBuilder, EntityEquivalent,
+        ContainsEntity, Entity, EntityCloner, EntityClonerBuilder, EntityEquivalent, EntityLocation,
-        EntityLocation,
     },
     event::Event,
     observer::Observer,
-    query::{Access, ReadOnlyQueryData},
+    query::{Access, DebugCheckedUnwrap, ReadOnlyQueryData},
     relationship::RelationshipHookMode,
-    removal_detection::RemovedComponentEvents,
     resource::Resource,
-    storage::Storages,
     system::IntoObserverSystem,
     world::{
-        error::EntityComponentError, unsafe_world_cell::UnsafeEntityCell, DeferredWorld, Mut, Ref,
+        error::EntityComponentError, unsafe_world_cell::UnsafeEntityCell, Mut, Ref, World,
-        World, ON_DESPAWN, ON_REMOVE, ON_REPLACE,
+        ON_DESPAWN, ON_REMOVE, ON_REPLACE,
     },
 };
 use alloc::vec::Vec;
@@ -1978,281 +1975,57 @@ impl<'w> EntityWorldMut<'w> {
     /// # Panics
     ///
     /// If the entity has been despawned while this `EntityWorldMut` is still alive.
-    // TODO: BundleRemover?
     #[must_use]
     #[track_caller]
     pub fn take<T: Bundle + BundleFromComponents>(&mut self) -> Option<T> {
         self.assert_not_despawned();
-        let world = &mut self.world;
-        let storages = &mut world.storages;
-        // SAFETY: These come from the same world.
-        let mut registrator =
-            unsafe { ComponentsRegistrator::new(&mut world.components, &mut world.component_ids) };
-        let bundle_id = world.bundles.register_info::<T>(&mut registrator, storages);
-        // SAFETY: We just ensured this bundle exists
-        let bundle_info = unsafe { world.bundles.get_unchecked(bundle_id) };
-        let old_location = self.location;
-        // SAFETY: `archetype_id` exists because it is referenced in the old `EntityLocation` which is valid,
-        // components exist in `bundle_info` because `Bundles::init_info` initializes a `BundleInfo` containing all components of the bundle type `T`
-        let new_archetype_id = unsafe {
-            bundle_info.remove_bundle_from_archetype(
-                &mut world.archetypes,
-                storages,
-                &registrator,
-                &world.observers,
-                old_location.archetype_id,
-                false,
-            )?
-        };
-
-        if new_archetype_id == old_location.archetype_id {
-            return None;
-        }
-
         let entity = self.entity;
-        // SAFETY: Archetypes and Bundles cannot be mutably aliased through DeferredWorld
+        let location = self.location;
-        let (old_archetype, bundle_info, mut deferred_world) = unsafe {
+
-            let bundle_info: *const BundleInfo = bundle_info;
+        let mut remover =
-            let world = world.as_unsafe_world_cell();
+            // SAFETY: The archetype id must be valid since this entity is in it.
-            (
+            unsafe { BundleRemover::new::<T>(self.world, self.location.archetype_id, true) }?;
-                &world.archetypes()[old_location.archetype_id],
+        // SAFETY: The passed location has the sane archetype as the remover, since they came from the same location.
-                &*bundle_info,
+        let (new_location, result) = unsafe {
-                world.into_deferred(),
+            remover.remove(
+                entity,
+                location,
+                MaybeLocation::caller(),
+                |sets, table, components, bundle_components| {
+                    let mut bundle_components = bundle_components.iter().copied();
+                    (
+                        false,
+                        T::from_components(&mut (sets, table), &mut |(sets, table)| {
+                            let component_id = bundle_components.next().unwrap();
+                            // SAFETY: the component existed to be removed, so its id must be valid.
+                            let component_info = components.get_info_unchecked(component_id);
+                            match component_info.storage_type() {
+                                StorageType::Table => {
+                                    table
+                                        .as_mut()
+                                        // SAFETY: The table must be valid if the component is in it.
+                                        .debug_checked_unwrap()
+                                        // SAFETY: The remover is cleaning this up.
+                                        .take_component(component_id, location.table_row)
+                                }
+                                StorageType::SparseSet => sets
+                                    .get_mut(component_id)
+                                    .unwrap()
+                                    .remove_and_forget(entity)
+                                    .unwrap(),
+                            }
+                        }),
+                    )
+                },
             )
         };
+        self.location = new_location;
 
-        // SAFETY: all bundle components exist in World
-        unsafe {
-            trigger_on_replace_and_on_remove_hooks_and_observers(
-                &mut deferred_world,
-                old_archetype,
-                entity,
-                bundle_info,
-                MaybeLocation::caller(),
-            );
-        }
-
-        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;
-        let mut bundle_components = bundle_info.iter_explicit_components();
-        // SAFETY: bundle components are iterated in order, which guarantees that the component type
-        // matches
-        let result = unsafe {
-            T::from_components(storages, &mut |storages| {
-                let component_id = bundle_components.next().unwrap();
-                // SAFETY:
-                // - entity location is valid
-                // - table row is removed below, without dropping the contents
-                // - `components` comes from the same world as `storages`
-                // - the component exists on the entity
-                take_component(
-                    storages,
-                    components,
-                    removed_components,
-                    component_id,
-                    entity,
-                    old_location,
-                )
-            })
-        };
-
-        #[expect(
-            clippy::undocumented_unsafe_blocks,
-            reason = "Needs to be documented; see #17345."
-        )]
-        unsafe {
-            Self::move_entity_from_remove::<false>(
-                entity,
-                &mut self.location,
-                old_location.archetype_id,
-                old_location,
-                entities,
-                archetypes,
-                storages,
-                new_archetype_id,
-            );
-        }
         self.world.flush();
         self.update_location();
         Some(result)
     }
 
-    /// # 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
-    ///
-    /// 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,
-    ) {
-        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
-        if let Some(swapped_entity) = remove_result.swapped_entity {
-            let swapped_location = entities.get(swapped_entity).unwrap();
-
-            entities.set(
-                swapped_entity.index(),
-                EntityLocation {
-                    archetype_id: swapped_location.archetype_id,
-                    archetype_row: old_location.archetype_row,
-                    table_id: swapped_location.table_id,
-                    table_row: swapped_location.table_row,
-                },
-            );
-        }
-        let old_table_row = remove_result.table_row;
-        let old_table_id = old_archetype.table_id();
-        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)
-        } else {
-            let (old_table, new_table) = storages
-                .tables
-                .get_2_mut(old_table_id, new_archetype.table_id());
-
-            let move_result = if DROP {
-                // SAFETY: old_table_row exists
-                unsafe { old_table.move_to_and_drop_missing_unchecked(old_table_row, new_table) }
-            } else {
-                // SAFETY: old_table_row exists
-                unsafe { old_table.move_to_and_forget_missing_unchecked(old_table_row, new_table) }
-            };
-
-            // 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) };
-
-            // if an entity was moved into this entity's table row, update its table row
-            if let Some(swapped_entity) = move_result.swapped_entity {
-                let swapped_location = entities.get(swapped_entity).unwrap();
-
-                entities.set(
-                    swapped_entity.index(),
-                    EntityLocation {
-                        archetype_id: swapped_location.archetype_id,
-                        archetype_row: swapped_location.archetype_row,
-                        table_id: swapped_location.table_id,
-                        table_row: old_location.table_row,
-                    },
-                );
-                archetypes[swapped_location.archetype_id]
-                    .set_entity_table_row(swapped_location.archetype_row, old_table_row);
-            }
-
-            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);
-        }
-    }
-
-    /// 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 {
-        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);
-
-        // 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");
-
-        if new_archetype_id == location.archetype_id {
-            return location;
-        }
-
-        // SAFETY: Archetypes and Bundles cannot be mutably aliased through DeferredWorld
-        let (old_archetype, bundle_info, mut deferred_world) = unsafe {
-            let bundle_info: *const BundleInfo = bundle_info;
-            let world = world.as_unsafe_world_cell();
-            (
-                &world.archetypes()[location.archetype_id],
-                &*bundle_info,
-                world.into_deferred(),
-            )
-        };
-
-        // SAFETY: all bundle components exist in World
-        unsafe {
-            trigger_on_replace_and_on_remove_hooks_and_observers(
-                &mut deferred_world,
-                old_archetype,
-                entity,
-                bundle_info,
-                caller,
-            );
-        }
-
-        let old_archetype = &world.archetypes[location.archetype_id];
-        for component_id in bundle_info.iter_explicit_components() {
-            if old_archetype.contains(component_id) {
-                world.removed_components.send(component_id, entity);
-
-                // Make sure to drop components stored in sparse sets.
-                // Dense components are dropped later in `move_to_and_drop_missing_unchecked`.
-                if let Some(StorageType::SparseSet) = old_archetype.get_storage_type(component_id) {
-                    world
-                        .storages
-                        .sparse_sets
-                        .get_mut(component_id)
-                        // Set exists because the component existed on the entity
-                        .unwrap()
-                        .remove(entity);
-                }
-            }
-        }
-
-        // 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
-    }
-
     /// Removes any components in the [`Bundle`] from the entity.
     ///
     /// See [`EntityCommands::remove`](crate::system::EntityCommands::remove) for more details.
@@ -2260,7 +2033,6 @@ impl<'w> EntityWorldMut<'w> {
     /// # Panics
     ///
     /// If the entity has been despawned while this `EntityWorldMut` is still alive.
-    // TODO: BundleRemover?
     #[track_caller]
     pub fn remove<T: Bundle>(&mut self) -> &mut Self {
         self.remove_with_caller::<T>(MaybeLocation::caller())
@@ -2269,18 +2041,25 @@ impl<'w> EntityWorldMut<'w> {
     #[inline]
     pub(crate) fn remove_with_caller<T: Bundle>(&mut self, caller: MaybeLocation) -> &mut Self {
         self.assert_not_despawned();
-        let storages = &mut self.world.storages;
-        // SAFETY: These come from the same world.
-        let mut registrator = unsafe {
-            ComponentsRegistrator::new(&mut self.world.components, &mut self.world.component_ids)
-        };
-        let bundle_info = self
-            .world
-            .bundles
-            .register_info::<T>(&mut registrator, storages);
 
-        // SAFETY: the `BundleInfo` is initialized above
+        let Some(mut remover) =
-        self.location = unsafe { self.remove_bundle(bundle_info, caller) };
+            // SAFETY: The archetype id must be valid since this entity is in it.
+            (unsafe { BundleRemover::new::<T>(self.world, self.location.archetype_id, false) })
+        else {
+            return self;
+        };
+        // SAFETY: The remover archetype came from the passed location and the removal can not fail.
+        let new_location = unsafe {
+            remover.remove(
+                self.entity,
+                self.location,
+                caller,
+                BundleRemover::empty_pre_remove,
+            )
+        }
+        .0;
+
+        self.location = new_location;
         self.world.flush();
         self.update_location();
         self
@@ -2302,16 +2081,31 @@ impl<'w> EntityWorldMut<'w> {
     ) -> &mut Self {
         self.assert_not_despawned();
         let storages = &mut self.world.storages;
+        let bundles = &mut self.world.bundles;
         // SAFETY: These come from the same world.
         let mut registrator = unsafe {
             ComponentsRegistrator::new(&mut self.world.components, &mut self.world.component_ids)
         };
-        let bundles = &mut self.world.bundles;
-
         let bundle_id = bundles.register_contributed_bundle_info::<T>(&mut registrator, storages);
 
-        // SAFETY: the dynamic `BundleInfo` is initialized above
+        // SAFETY: We just created the bundle, and the archetype is valid, since we are in it.
-        self.location = unsafe { self.remove_bundle(bundle_id, caller) };
+        let Some(mut remover) = (unsafe {
+            BundleRemover::new_with_id(self.world, self.location.archetype_id, bundle_id, false)
+        }) else {
+            return self;
+        };
+        // SAFETY: The remover archetype came from the passed location and the removal can not fail.
+        let new_location = unsafe {
+            remover.remove(
+                self.entity,
+                self.location,
+                caller,
+                BundleRemover::empty_pre_remove,
+            )
+        }
+        .0;
+
+        self.location = new_location;
         self.world.flush();
         self.update_location();
         self
@@ -2358,8 +2152,24 @@ 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
+        // SAFETY: We just created the bundle, and the archetype is valid, since we are in it.
-        self.location = unsafe { self.remove_bundle(remove_bundle, caller) };
+        let Some(mut remover) = (unsafe {
+            BundleRemover::new_with_id(self.world, self.location.archetype_id, remove_bundle, false)
+        }) else {
+            return self;
+        };
+        // SAFETY: The remover archetype came from the passed location and the removal can not fail.
+        let new_location = unsafe {
+            remover.remove(
+                self.entity,
+                self.location,
+                caller,
+                BundleRemover::empty_pre_remove,
+            )
+        }
+        .0;
+
+        self.location = new_location;
         self.world.flush();
         self.update_location();
         self
@@ -2393,8 +2203,24 @@ impl<'w> EntityWorldMut<'w> {
             component_id,
         );
 
-        // SAFETY: the `BundleInfo` for this `component_id` is initialized above
+        // SAFETY: We just created the bundle, and the archetype is valid, since we are in it.
-        self.location = unsafe { self.remove_bundle(bundle_id, caller) };
+        let Some(mut remover) = (unsafe {
+            BundleRemover::new_with_id(self.world, self.location.archetype_id, bundle_id, false)
+        }) else {
+            return self;
+        };
+        // SAFETY: The remover archetype came from the passed location and the removal can not fail.
+        let new_location = unsafe {
+            remover.remove(
+                self.entity,
+                self.location,
+                caller,
+                BundleRemover::empty_pre_remove,
+            )
+        }
+        .0;
+
+        self.location = new_location;
         self.world.flush();
         self.update_location();
         self
@@ -2419,9 +2245,24 @@ impl<'w> EntityWorldMut<'w> {
             component_ids,
         );
 
-        // SAFETY: the `BundleInfo` for this `bundle_id` is initialized above
+        // SAFETY: We just created the bundle, and the archetype is valid, since we are in it.
-        unsafe { self.remove_bundle(bundle_id, MaybeLocation::caller()) };
+        let Some(mut remover) = (unsafe {
+            BundleRemover::new_with_id(self.world, self.location.archetype_id, bundle_id, false)
+        }) else {
+            return self;
+        };
+        // SAFETY: The remover archetype came from the passed location and the removal can not fail.
+        let new_location = unsafe {
+            remover.remove(
+                self.entity,
+                self.location,
+                MaybeLocation::caller(),
+                BundleRemover::empty_pre_remove,
+            )
+        }
+        .0;
 
+        self.location = new_location;
         self.world.flush();
         self.update_location();
         self
@@ -2449,8 +2290,24 @@ impl<'w> EntityWorldMut<'w> {
             component_ids.as_slice(),
         );
 
-        // SAFETY: the `BundleInfo` for this `component_id` is initialized above
+        // SAFETY: We just created the bundle, and the archetype is valid, since we are in it.
-        self.location = unsafe { self.remove_bundle(bundle_id, caller) };
+        let Some(mut remover) = (unsafe {
+            BundleRemover::new_with_id(self.world, self.location.archetype_id, bundle_id, false)
+        }) else {
+            return self;
+        };
+        // SAFETY: The remover archetype came from the passed location and the removal can not fail.
+        let new_location = unsafe {
+            remover.remove(
+                self.entity,
+                self.location,
+                caller,
+                BundleRemover::empty_pre_remove,
+            )
+        }
+        .0;
+
+        self.location = new_location;
         self.world.flush();
         self.update_location();
         self
@@ -2956,46 +2813,6 @@ impl<'w> EntityWorldMut<'w> {
     }
 }
 
-/// # Safety
-/// All components in the archetype must exist in world
-unsafe fn trigger_on_replace_and_on_remove_hooks_and_observers(
-    deferred_world: &mut DeferredWorld,
-    archetype: &Archetype,
-    entity: Entity,
-    bundle_info: &BundleInfo,
-    caller: MaybeLocation,
-) {
-    let bundle_components_in_archetype = || {
-        bundle_info
-            .iter_explicit_components()
-            .filter(|component_id| archetype.contains(*component_id))
-    };
-    if archetype.has_replace_observer() {
-        deferred_world.trigger_observers(
-            ON_REPLACE,
-            entity,
-            bundle_components_in_archetype(),
-            caller,
-        );
-    }
-    deferred_world.trigger_on_replace(
-        archetype,
-        entity,
-        bundle_components_in_archetype(),
-        caller,
-        RelationshipHookMode::Run,
-    );
-    if archetype.has_remove_observer() {
-        deferred_world.trigger_observers(
-            ON_REMOVE,
-            entity,
-            bundle_components_in_archetype(),
-            caller,
-        );
-    }
-    deferred_world.trigger_on_remove(archetype, entity, bundle_components_in_archetype(), caller);
-}
-
 /// A view into a single entity and component in a world, which may either be vacant or occupied.
 ///
 /// This `enum` can only be constructed from the [`entry`] method on [`EntityWorldMut`].
@@ -4443,7 +4260,7 @@ where
 /// # Safety
 ///
 /// - [`OwningPtr`] and [`StorageType`] iterators must correspond to the
-///   [`BundleInfo`] used to construct [`BundleInserter`]
+///   [`BundleInfo`](crate::bundle::BundleInfo) used to construct [`BundleInserter`]
 /// - [`Entity`] must correspond to [`EntityLocation`]
 unsafe fn insert_dynamic_bundle<
     'a,
@@ -4491,50 +4308,6 @@ unsafe fn insert_dynamic_bundle<
     }
 }
 
-/// Moves component data out of storage.
-///
-/// This function leaves the underlying memory unchanged, but the component behind
-/// returned pointer is semantically owned by the caller and will not be dropped in its original location.
-/// Caller is responsible to drop component data behind returned pointer.
-///
-/// # Safety
-/// - `location.table_row` must be in bounds of column of component id `component_id`
-/// - `component_id` must be valid
-/// - `components` must come from the same world as `self`
-/// - The relevant table row **must be removed** by the caller once all components are taken, without dropping the value
-///
-/// # Panics
-/// Panics if the entity did not have the component.
-#[inline]
-pub(crate) unsafe fn take_component<'a>(
-    storages: &'a mut Storages,
-    components: &Components,
-    removed_components: &mut RemovedComponentEvents,
-    component_id: ComponentId,
-    entity: Entity,
-    location: EntityLocation,
-) -> OwningPtr<'a> {
-    // SAFETY: caller promises component_id to be valid
-    let component_info = unsafe { components.get_info_unchecked(component_id) };
-    removed_components.send(component_id, entity);
-    match component_info.storage_type() {
-        StorageType::Table => {
-            let table = &mut storages.tables[location.table_id];
-            // SAFETY:
-            // - archetypes only store valid table_rows
-            // - index is in bounds as promised by caller
-            // - promote is safe because the caller promises to remove the table row without dropping it immediately afterwards
-            unsafe { table.take_component(component_id, location.table_row) }
-        }
-        StorageType::SparseSet => storages
-            .sparse_sets
-            .get_mut(component_id)
-            .unwrap()
-            .remove_and_forget(entity)
-            .unwrap(),
-    }
-}
-
 /// Types that can be used to fetch components from an entity dynamically by
 /// [`ComponentId`]s.
 ///