Fix sparse insert (#1748)

Removing the checks on this line https://github.com/bevyengine/bevy/blob/main/crates/bevy_sprite/src/frustum_culling.rs#L64 and running the "many_sprites" example revealed two corner case bugs in bevy_ecs. The first, a simple and honest missed line introduced in #1471. The other, an insidious monster that has been there since the ECS v2 rewrite, just waiting for the time to strike:

1. #1471 accidentally removed the "insert" line for sparse set components with the "mutated" bundle state. Re-adding it fixes the problem. I did a slight refactor here to make the implementation simpler and remove a branch.
2. The other issue is nastier. ECS v2 added an "archetype graph". When determining what components were added/mutated during an archetype change, we read the FromBundle edge (which encodes this state) on the "new" archetype.  The problem is that unlike "add edges" which are guaranteed to be unique for a given ("graph node", "bundle id") pair, FromBundle edges are not necessarily unique:

```rust
// OLD_ARCHETYPE -> NEW_ARCHETYPE

// [] -> [usize]
e.insert(2usize);
// [usize] -> [usize, i32]
e.insert(1i32);
// [usize, i32] -> [usize, i32]
e.insert(1i32);
// [usize, i32] -> [usize]
e.remove::<i32>();
// [usize] -> [usize, i32]
e.insert(1i32);
```

Note that the second `e.insert(1i32)` command has a different "archetype graph edge" than the first, but they both lead to the same "new archetype".

The fix here is simple: just remove FromBundle edges because they are broken and store the information in the "add edges", which are guaranteed to be unique.

FromBundle edges were added to cut down on the number of archetype accesses / make the archetype access patterns nicer. But benching this change resulted in no significant perf changes and the addition of get_2_mut() for archetypes resolves the access pattern issue.

Author: cart

crates/bevy_ecs/src/archetype.rs

@@ -41,45 +41,34 @@ pub enum ComponentStatus {
     Mutated,
 }
 
-pub struct FromBundle {
+pub struct AddBundle {
     pub archetype_id: ArchetypeId,
     pub bundle_status: Vec<ComponentStatus>,
 }
 
 #[derive(Default)]
 pub struct Edges {
-    pub add_bundle: SparseArray<BundleId, ArchetypeId>,
+    pub add_bundle: SparseArray<BundleId, AddBundle>,
     pub remove_bundle: SparseArray<BundleId, Option<ArchetypeId>>,
     pub remove_bundle_intersection: SparseArray<BundleId, Option<ArchetypeId>>,
-    pub from_bundle: SparseArray<BundleId, FromBundle>,
 }
 
 impl Edges {
     #[inline]
-    pub fn get_add_bundle(&self, bundle_id: BundleId) -> Option<ArchetypeId> {
-        self.add_bundle.get(bundle_id).cloned()
+    pub fn get_add_bundle(&self, bundle_id: BundleId) -> Option<&AddBundle> {
+        self.add_bundle.get(bundle_id)
     }
 
     #[inline]
-    pub fn set_add_bundle(&mut self, bundle_id: BundleId, archetype_id: ArchetypeId) {
-        self.add_bundle.insert(bundle_id, archetype_id);
-    }
-
-    #[inline]
-    pub fn get_from_bundle(&self, bundle_id: BundleId) -> Option<&FromBundle> {
-        self.from_bundle.get(bundle_id)
-    }
-
-    #[inline]
-    pub fn set_from_bundle(
+    pub fn set_add_bundle(
         &mut self,
         bundle_id: BundleId,
         archetype_id: ArchetypeId,
         bundle_status: Vec<ComponentStatus>,
     ) {
-        self.from_bundle.insert(
+        self.add_bundle.insert(
             bundle_id,
-            FromBundle {
+            AddBundle {
                 archetype_id,
                 bundle_status,
             },
@@ -458,6 +447,21 @@ impl Archetypes {
         self.archetypes.get_mut(id.index())
     }
 
+    #[inline]
+    pub(crate) fn get_2_mut(
+        &mut self,
+        a: ArchetypeId,
+        b: ArchetypeId,
+    ) -> (&mut Archetype, &mut Archetype) {
+        if a.index() > b.index() {
+            let (b_slice, a_slice) = self.archetypes.split_at_mut(a.index());
+            (&mut a_slice[0], &mut b_slice[b.index()])
+        } else {
+            let (a_slice, b_slice) = self.archetypes.split_at_mut(b.index());
+            (&mut a_slice[a.index()], &mut b_slice[0])
+        }
+    }
+
     #[inline]
     pub fn iter(&self) -> impl Iterator<Item = &Archetype> {
         self.archetypes.iter()

crates/bevy_ecs/src/bundle.rs

@@ -159,21 +159,7 @@ impl BundleInfo {
                 }
                 StorageType::SparseSet => {
                     let sparse_set = sparse_sets.get_mut(component_id).unwrap();
-                    match component_status {
-                        ComponentStatus::Added => {
-                            sparse_set.insert(
-                                entity,
-                                component_ptr,
-                                ComponentTicks::new(change_tick),
-                            );
-                        }
-                        ComponentStatus::Mutated => {
-                            sparse_set
-                                .get_ticks(entity)
-                                .unwrap()
-                                .set_changed(change_tick);
-                        }
-                    }
+                    sparse_set.insert(entity, component_ptr, change_tick);
                 }
             }
             bundle_component += 1;

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -119,23 +119,18 @@ impl ComponentSparseSet {
     /// # Safety
     /// The `value` pointer must point to a valid address that matches the `Layout`
     ///  inside the `ComponentInfo` given when constructing this sparse set.
-    pub unsafe fn insert(
-        &mut self,
-        entity: Entity,
-        value: *mut u8,
-        component_ticks: ComponentTicks,
-    ) {
+    pub unsafe fn insert(&mut self, entity: Entity, value: *mut u8, change_tick: u32) {
         let dense = &mut self.dense;
         let entities = &mut self.entities;
         let ticks_list = self.ticks.get_mut();
         let dense_index = *self.sparse.get_or_insert_with(entity, move || {
-            ticks_list.push(component_ticks);
+            ticks_list.push(ComponentTicks::new(change_tick));
             entities.push(entity);
             dense.push_uninit()
         });
         // SAFE: dense_index exists thanks to the call above
         self.dense.set_unchecked(dense_index, value);
-        *(*self.ticks.get()).get_unchecked_mut(dense_index) = component_ticks;
+        ((*self.ticks.get()).get_unchecked_mut(dense_index)).set_changed(change_tick);
     }
 
     #[inline]

crates/bevy_ecs/src/world/entity_ref.rs

@@ -195,7 +195,7 @@ impl<'w> EntityMut<'w> {
         let bundle_info = self.world.bundles.init_info::<T>(components);
         let current_location = self.location;
 
-        let new_location = unsafe {
+        let (archetype, bundle_status, archetype_index) = unsafe {
             // SAFE: component ids in `bundle_info` and self.location are valid
             let new_archetype_id = add_bundle_to_archetype(
                 archetypes,
@@ -205,61 +205,66 @@ impl<'w> EntityMut<'w> {
                 bundle_info,
             );
             if new_archetype_id == current_location.archetype_id {
-                current_location
+                let archetype = &archetypes[current_location.archetype_id];
+                let edge = archetype.edges().get_add_bundle(bundle_info.id).unwrap();
+                (archetype, &edge.bundle_status, current_location.index)
             } else {
-                let old_table_row;
-                let old_table_id;
-                {
+                let (old_table_row, old_table_id) = {
                     let old_archetype = &mut archetypes[current_location.archetype_id];
                     let result = old_archetype.swap_remove(current_location.index);
                     if let Some(swapped_entity) = result.swapped_entity {
-                        // SAFE: entity is live and is contained in an archetype that exists
                         entities.meta[swapped_entity.id as usize].location = current_location;
                     }
-                    old_table_row = result.table_row;
-                    old_table_id = old_archetype.table_id()
-                }
-                let new_archetype = &mut archetypes[new_archetype_id];
+                    (result.table_row, old_archetype.table_id())
+                };
+
+                let new_table_id = archetypes[new_archetype_id].table_id();
 
-                if old_table_id == new_archetype.table_id() {
-                    new_archetype.allocate(entity, old_table_row)
+                let new_location = if old_table_id == new_table_id {
+                    archetypes[new_archetype_id].allocate(entity, old_table_row)
                 } else {
-                    let (old_table, new_table) = storages
-                        .tables
-                        .get_2_mut(old_table_id, new_archetype.table_id());
+                    let (old_table, new_table) =
+                        storages.tables.get_2_mut(old_table_id, new_table_id);
                     // PERF: store "non bundle" components in edge, then just move those to avoid
                     // redundant copies
                     let move_result =
                         old_table.move_to_superset_unchecked(old_table_row, new_table);
 
-                    let new_location = new_archetype.allocate(entity, move_result.new_row);
+                    let new_location =
+                        archetypes[new_archetype_id].allocate(entity, move_result.new_row);
                     // if an entity was moved into this entity's table spot, update its table row
                     if let Some(swapped_entity) = move_result.swapped_entity {
                         let swapped_location = entities.get(swapped_entity).unwrap();
                         archetypes[swapped_location.archetype_id]
                             .set_entity_table_row(swapped_location.index, old_table_row);
                     }
                     new_location
-                }
+                };
+
+                self.location = new_location;
+                entities.meta[self.entity.id as usize].location = new_location;
+                let (old_archetype, new_archetype) =
+                    archetypes.get_2_mut(current_location.archetype_id, new_archetype_id);
+                let edge = old_archetype
+                    .edges()
+                    .get_add_bundle(bundle_info.id)
+                    .unwrap();
+                (&*new_archetype, &edge.bundle_status, new_location.index)
 
                 // Sparse set components are intentionally ignored here. They don't need to move
             }
         };
-        self.location = new_location;
-        entities.meta[self.entity.id as usize].location = new_location;
 
-        let archetype = &archetypes[new_location.archetype_id];
         let table = &storages.tables[archetype.table_id()];
-        let table_row = archetype.entity_table_row(new_location.index);
-        let from_bundle = archetype.edges().get_from_bundle(bundle_info.id).unwrap();
+        let table_row = archetype.entity_table_row(archetype_index);
         // SAFE: table row is valid
         unsafe {
             bundle_info.write_components(
                 &mut storages.sparse_sets,
                 entity,
                 table,
                 table_row,
-                &from_bundle.bundle_status,
+                bundle_status,
                 bundle,
                 change_tick,
             )
@@ -650,11 +655,11 @@ pub(crate) unsafe fn add_bundle_to_archetype(
     archetype_id: ArchetypeId,
     bundle_info: &BundleInfo,
 ) -> ArchetypeId {
-    if let Some(archetype_id) = archetypes[archetype_id]
+    if let Some(add_bundle) = archetypes[archetype_id]
         .edges()
         .get_add_bundle(bundle_info.id)
     {
-        return archetype_id;
+        return add_bundle.archetype_id;
     }
     let mut new_table_components = Vec::new();
     let mut new_sparse_set_components = Vec::new();
@@ -680,8 +685,7 @@ pub(crate) unsafe fn add_bundle_to_archetype(
     if new_table_components.is_empty() && new_sparse_set_components.is_empty() {
         let edges = current_archetype.edges_mut();
         // the archetype does not change when we add this bundle
-        edges.set_add_bundle(bundle_info.id, archetype_id);
-        edges.set_from_bundle(bundle_info.id, archetype_id, bundle_status);
+        edges.set_add_bundle(bundle_info.id, archetype_id, bundle_status);
         archetype_id
     } else {
         let table_id;
@@ -718,11 +722,7 @@ pub(crate) unsafe fn add_bundle_to_archetype(
         let new_archetype_id =
             archetypes.get_id_or_insert(table_id, table_components, sparse_set_components);
         // add an edge from the old archetype to the new archetype
-        archetypes[archetype_id]
-            .edges_mut()
-            .set_add_bundle(bundle_info.id, new_archetype_id);
-        // add a "from bundle" edge from the new archetype to the old archetype
-        archetypes[new_archetype_id].edges_mut().set_from_bundle(
+        archetypes[archetype_id].edges_mut().set_add_bundle(
             bundle_info.id,
             new_archetype_id,
             bundle_status,

crates/bevy_ecs/src/world/spawn_batch.rs

@@ -1,5 +1,5 @@
 use crate::{
-    archetype::{Archetype, ArchetypeId},
+    archetype::{Archetype, ArchetypeId, ComponentStatus},
     bundle::{Bundle, BundleInfo},
     entity::{Entities, Entity},
     storage::{SparseSets, Table},
@@ -17,6 +17,7 @@ where
     table: &'w mut Table,
     sparse_sets: &'w mut SparseSets,
     bundle_info: &'w BundleInfo,
+    bundle_status: &'w [ComponentStatus],
     change_tick: u32,
 }
 
@@ -46,11 +47,17 @@ where
                 bundle_info,
             )
         };
-        let archetype = &mut world.archetypes[archetype_id];
+        let (empty_archetype, archetype) = world
+            .archetypes
+            .get_2_mut(ArchetypeId::empty(), archetype_id);
         let table = &mut world.storages.tables[archetype.table_id()];
         archetype.reserve(length);
         table.reserve(length);
         world.entities.reserve(length as u32);
+        let edge = empty_archetype
+            .edges()
+            .get_add_bundle(bundle_info.id())
+            .unwrap();
         Self {
             inner: iter,
             entities: &mut world.entities,
@@ -59,6 +66,7 @@ where
             sparse_sets: &mut world.storages.sparse_sets,
             bundle_info,
             change_tick: *world.change_tick.get_mut(),
+            bundle_status: &edge.bundle_status,
         }
     }
 }
@@ -88,17 +96,12 @@ where
         unsafe {
             let table_row = self.table.allocate(entity);
             let location = self.archetype.allocate(entity, table_row);
-            let from_bundle = self
-                .archetype
-                .edges()
-                .get_from_bundle(self.bundle_info.id)
-                .unwrap();
             self.bundle_info.write_components(
                 self.sparse_sets,
                 entity,
                 self.table,
                 table_row,
-                &from_bundle.bundle_status,
+                self.bundle_status,
                 bundle,
                 self.change_tick,
             );