Reduce memory usage from Edges

crates/bevy_ecs/src/archetype.rs

@@ -25,7 +25,7 @@ use crate::{
     entity::{Entity, EntityLocation},
     event::Event,
     observer::Observers,
-    storage::{ImmutableSparseSet, SparseArray, SparseSet, TableId, TableRow},
+    storage::{ImmutableSparseSet, SparseSet, TableId, TableRow},
 };
 use alloc::{boxed::Box, vec::Vec};
 use bevy_platform::collections::{hash_map::Entry, HashMap};
@@ -193,6 +193,12 @@ impl BundleComponentStatus for SpawnBundleStatus {
     }
 }
 
+#[derive(Debug, Hash, PartialEq, Eq)]
+struct ArchetypeMove {
+    source: ArchetypeId,
+    bundle_id: BundleId,
+}
+
 /// Archetypes and bundles form a graph. Adding or removing a bundle moves
 /// an [`Entity`] to a new [`Archetype`].
 ///
@@ -209,9 +215,9 @@ impl BundleComponentStatus for SpawnBundleStatus {
 /// [`World`]: crate::world::World
 #[derive(Default)]
 pub struct Edges {
-    insert_bundle: SparseArray<BundleId, ArchetypeAfterBundleInsert>,
-    remove_bundle: SparseArray<BundleId, Option<ArchetypeId>>,
-    take_bundle: SparseArray<BundleId, Option<ArchetypeId>>,
+    insert_bundle: HashMap<ArchetypeMove, ArchetypeAfterBundleInsert>,
+    remove_bundle: HashMap<ArchetypeMove, Option<ArchetypeId>>,
+    take_bundle: HashMap<ArchetypeMove, Option<ArchetypeId>>,
 }
 
 impl Edges {
@@ -221,8 +227,12 @@ impl Edges {
     /// If this returns `None`, it means there has not been a transition from
     /// the source archetype via the provided bundle.
     #[inline]
-    pub fn get_archetype_after_bundle_insert(&self, bundle_id: BundleId) -> Option<ArchetypeId> {
-        self.get_archetype_after_bundle_insert_internal(bundle_id)
+    pub fn get_archetype_after_bundle_insert(
+        &self,
+        source: ArchetypeId,
+        bundle_id: BundleId,
+    ) -> Option<ArchetypeId> {
+        self.get_archetype_after_bundle_insert_internal(source, bundle_id)
             .map(|bundle| bundle.archetype_id)
     }
 
@@ -231,15 +241,17 @@ impl Edges {
     #[inline]
     pub(crate) fn get_archetype_after_bundle_insert_internal(
         &self,
+        source: ArchetypeId,
         bundle_id: BundleId,
     ) -> Option<&ArchetypeAfterBundleInsert> {
-        self.insert_bundle.get(bundle_id)
+        self.insert_bundle.get(&ArchetypeMove { source, bundle_id })
     }
 
     /// Caches the target archetype when inserting a bundle into the source archetype.
     #[inline]
     pub(crate) fn cache_archetype_after_bundle_insert(
         &mut self,
+        source: ArchetypeId,
         bundle_id: BundleId,
         archetype_id: ArchetypeId,
         bundle_status: impl Into<Box<[ComponentStatus]>>,
@@ -248,7 +260,7 @@ impl Edges {
         existing: impl Into<Box<[ComponentId]>>,
     ) {
         self.insert_bundle.insert(
-            bundle_id,
+            ArchetypeMove { source, bundle_id },
             ArchetypeAfterBundleInsert {
                 archetype_id,
                 bundle_status: bundle_status.into(),
@@ -270,19 +282,24 @@ impl Edges {
     #[inline]
     pub fn get_archetype_after_bundle_remove(
         &self,
+        source: ArchetypeId,
         bundle_id: BundleId,
     ) -> Option<Option<ArchetypeId>> {
-        self.remove_bundle.get(bundle_id).cloned()
+        self.remove_bundle
+            .get(&ArchetypeMove { source, bundle_id })
+            .cloned()
     }
 
     /// Caches the target archetype when removing a bundle from the source archetype.
     #[inline]
     pub(crate) fn cache_archetype_after_bundle_remove(
         &mut self,
+        source: ArchetypeId,
         bundle_id: BundleId,
         archetype_id: Option<ArchetypeId>,
     ) {
-        self.remove_bundle.insert(bundle_id, archetype_id);
+        self.remove_bundle
+            .insert(ArchetypeMove { source, bundle_id }, archetype_id);
     }
 
     /// Checks the cache for the target archetype when taking a bundle from the
@@ -299,9 +316,12 @@ impl Edges {
     #[inline]
     pub fn get_archetype_after_bundle_take(
         &self,
+        source: ArchetypeId,
         bundle_id: BundleId,
     ) -> Option<Option<ArchetypeId>> {
-        self.take_bundle.get(bundle_id).cloned()
+        self.take_bundle
+            .get(&ArchetypeMove { source, bundle_id })
+            .cloned()
     }
 
     /// Caches the target archetype when taking a bundle from the source archetype.
@@ -311,10 +331,12 @@ impl Edges {
     #[inline]
     pub(crate) fn cache_archetype_after_bundle_take(
         &mut self,
+        source: ArchetypeId,
         bundle_id: BundleId,
         archetype_id: Option<ArchetypeId>,
     ) {
-        self.take_bundle.insert(bundle_id, archetype_id);
+        self.take_bundle
+            .insert(ArchetypeMove { source, bundle_id }, archetype_id);
     }
 }
 
@@ -384,7 +406,6 @@ bitflags::bitflags! {
 pub struct Archetype {
     id: ArchetypeId,
     table_id: TableId,
-    edges: Edges,
     entities: Vec<ArchetypeEntity>,
     components: ImmutableSparseSet<ComponentId, ArchetypeComponentInfo>,
     pub(crate) flags: ArchetypeFlags,
@@ -446,7 +467,6 @@ impl Archetype {
             table_id,
             entities: Vec::new(),
             components: archetype_components.into_immutable(),
-            edges: Default::default(),
             flags,
         }
     }
@@ -538,20 +558,6 @@ impl Archetype {
         self.components.len()
     }
 
-    /// Fetches an immutable reference to the archetype's [`Edges`], a cache of
-    /// archetypal relationships.
-    #[inline]
-    pub fn edges(&self) -> &Edges {
-        &self.edges
-    }
-
-    /// Fetches a mutable reference to the archetype's [`Edges`], a cache of
-    /// archetypal relationships.
-    #[inline]
-    pub(crate) fn edges_mut(&mut self) -> &mut Edges {
-        &mut self.edges
-    }
-
     /// Fetches the row in the [`Table`] where the components for the entity at `index`
     /// is stored.
     ///

crates/bevy_ecs/src/bundle/insert.rs

@@ -5,7 +5,7 @@ use core::ptr::NonNull;
 use crate::{
     archetype::{
         Archetype, ArchetypeAfterBundleInsert, ArchetypeCreated, ArchetypeId, Archetypes,
-        ComponentStatus,
+        ComponentStatus, Edges,
     },
     bundle::{ArchetypeMoveType, Bundle, BundleId, BundleInfo, DynamicBundle, InsertMode},
     change_detection::MaybeLocation,
@@ -63,6 +63,7 @@ impl<'w> BundleInserter<'w> {
         let bundle_id = bundle_info.id();
         let (new_archetype_id, is_new_created) = bundle_info.insert_bundle_into_archetype(
             &mut world.archetypes,
+            &mut world.edges,
             &mut world.storages,
             &world.components,
             &world.observers,
@@ -73,9 +74,9 @@ impl<'w> BundleInserter<'w> {
             let archetype = &mut world.archetypes[archetype_id];
             // SAFETY: The edge is assured to be initialized when we called insert_bundle_into_archetype
             let archetype_after_insert = unsafe {
-                archetype
-                    .edges()
-                    .get_archetype_after_bundle_insert_internal(bundle_id)
+                world
+                    .edges
+                    .get_archetype_after_bundle_insert_internal(archetype_id, bundle_id)
                     .debug_checked_unwrap()
             };
             let table_id = archetype.table_id();
@@ -94,9 +95,9 @@ impl<'w> BundleInserter<'w> {
                 world.archetypes.get_2_mut(archetype_id, new_archetype_id);
             // SAFETY: The edge is assured to be initialized when we called insert_bundle_into_archetype
             let archetype_after_insert = unsafe {
-                archetype
-                    .edges()
-                    .get_archetype_after_bundle_insert_internal(bundle_id)
+                world
+                    .edges
+                    .get_archetype_after_bundle_insert_internal(archetype_id, bundle_id)
                     .debug_checked_unwrap()
             };
             let table_id = archetype.table_id();
@@ -420,14 +421,14 @@ impl BundleInfo {
     pub(crate) unsafe fn insert_bundle_into_archetype(
         &self,
         archetypes: &mut Archetypes,
+        edges: &mut Edges,
         storages: &mut Storages,
         components: &Components,
         observers: &Observers,
         archetype_id: ArchetypeId,
     ) -> (ArchetypeId, bool) {
-        if let Some(archetype_after_insert_id) = archetypes[archetype_id]
-            .edges()
-            .get_archetype_after_bundle_insert(self.id)
+        if let Some(archetype_after_insert_id) =
+            edges.get_archetype_after_bundle_insert(archetype_id, self.id)
         {
             return (archetype_after_insert_id, false);
         }
@@ -473,9 +474,9 @@ impl BundleInfo {
         }
 
         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 insert this bundle.
             edges.cache_archetype_after_bundle_insert(
+                archetype_id,
                 self.id,
                 archetype_id,
                 bundle_status,
@@ -528,16 +529,15 @@ impl BundleInfo {
             );
 
             // Add an edge from the old archetype to the new archetype.
-            archetypes[archetype_id]
-                .edges_mut()
-                .cache_archetype_after_bundle_insert(
-                    self.id,
-                    new_archetype_id,
-                    bundle_status,
-                    added_required_components,
-                    added,
-                    existing,
-                );
+            edges.cache_archetype_after_bundle_insert(
+                archetype_id,
+                self.id,
+                new_archetype_id,
+                bundle_status,
+                added_required_components,
+                added,
+                existing,
+            );
             (new_archetype_id, is_new_created)
         }
     }

crates/bevy_ecs/src/bundle/remove.rs

@@ -3,7 +3,7 @@ use bevy_ptr::ConstNonNull;
 use core::ptr::NonNull;
 
 use crate::{
-    archetype::{Archetype, ArchetypeCreated, ArchetypeId, Archetypes},
+    archetype::{Archetype, ArchetypeCreated, ArchetypeId, Archetypes, Edges},
     bundle::{Bundle, BundleId, BundleInfo},
     change_detection::MaybeLocation,
     component::{ComponentId, Components, ComponentsRegistrator, StorageType},
@@ -66,6 +66,7 @@ impl<'w> BundleRemover<'w> {
         let (new_archetype_id, is_new_created) = unsafe {
             bundle_info.remove_bundle_from_archetype(
                 &mut world.archetypes,
+                &mut world.edges,
                 &mut world.storages,
                 &world.components,
                 &world.observers,
@@ -314,6 +315,7 @@ impl BundleInfo {
     pub(crate) unsafe fn remove_bundle_from_archetype(
         &self,
         archetypes: &mut Archetypes,
+        edges: &mut Edges,
         storages: &mut Storages,
         components: &Components,
         observers: &Observers,
@@ -323,11 +325,10 @@ impl BundleInfo {
         // Check the archetype graph to see if the bundle has been
         // removed from this archetype in the past.
         let archetype_after_remove_result = {
-            let edges = archetypes[archetype_id].edges();
             if intersection {
-                edges.get_archetype_after_bundle_remove(self.id())
+                edges.get_archetype_after_bundle_remove(archetype_id, self.id())
             } else {
-                edges.get_archetype_after_bundle_take(self.id())
+                edges.get_archetype_after_bundle_take(archetype_id, self.id())
             }
         };
         let (result, is_new_created) = if let Some(result) = archetype_after_remove_result {
@@ -354,9 +355,7 @@ impl BundleInfo {
                     } else if !intersection {
                         // A component in the bundle was not present in the entity's archetype, so this
                         // removal is invalid. Cache the result in the archetype graph.
-                        current_archetype
-                            .edges_mut()
-                            .cache_archetype_after_bundle_take(self.id(), None);
+                        edges.cache_archetype_after_bundle_take(archetype_id, self.id(), None);
                         return (None, false);
                     }
                 }
@@ -394,16 +393,12 @@ impl BundleInfo {
             );
             (Some(new_archetype_id), is_new_created)
         };
-        let current_archetype = &mut archetypes[archetype_id];
+
         // Cache the result in an edge.
         if intersection {
-            current_archetype
-                .edges_mut()
-                .cache_archetype_after_bundle_remove(self.id(), result);
+            edges.cache_archetype_after_bundle_remove(archetype_id, self.id(), result);
         } else {
-            current_archetype
-                .edges_mut()
-                .cache_archetype_after_bundle_take(self.id(), result);
+            edges.cache_archetype_after_bundle_take(archetype_id, self.id(), result);
         }
         (result, is_new_created)
     }

crates/bevy_ecs/src/bundle/spawner.rs

@@ -49,6 +49,7 @@ impl<'w> BundleSpawner<'w> {
         let bundle_info = world.bundles.get_unchecked(bundle_id);
         let (new_archetype_id, is_new_created) = bundle_info.insert_bundle_into_archetype(
             &mut world.archetypes,
+            &mut world.edges,
             &mut world.storages,
             &world.components,
             &world.observers,

crates/bevy_ecs/src/world/mod.rs

@@ -37,7 +37,7 @@ pub use identifier::WorldId;
 pub use spawn_batch::*;
 
 use crate::{
-    archetype::{ArchetypeId, Archetypes},
+    archetype::{ArchetypeId, Archetypes, Edges},
     bundle::{
         Bundle, BundleEffect, BundleInfo, BundleInserter, BundleSpawner, Bundles, InsertMode,
         NoBundleEffect,
@@ -97,6 +97,7 @@ pub struct World {
     pub(crate) components: Components,
     pub(crate) component_ids: ComponentIds,
     pub(crate) archetypes: Archetypes,
+    pub(crate) edges: Edges,
     pub(crate) storages: Storages,
     pub(crate) bundles: Bundles,
     pub(crate) observers: Observers,
@@ -115,6 +116,7 @@ impl Default for World {
             entities: Entities::new(),
             components: Default::default(),
             archetypes: Archetypes::new(),
+            edges: Default::default(),
             storages: Default::default(),
             bundles: Default::default(),
             observers: Observers::default(),
@@ -222,6 +224,13 @@ impl World {
         &self.archetypes
     }
 
+    /// Fetches an immutable reference to the [`Edges`], a cache of
+    /// archetypal relationships.
+    #[inline]
+    pub fn edges(&self) -> &Edges {
+        &self.edges
+    }
+
     /// Retrieves this world's [`Components`] collection.
     #[inline]
     pub fn components(&self) -> &Components {