Add basic system for specifying and validating archetype invariants

Cargo.toml

@@ -761,6 +761,16 @@ description = "Full guide to Bevy's ECS"
 category = "ECS (Entity Component System)"
 wasm = false
 
+[[example]]
+name = "archetype_invariants"
+path = "examples/ecs/archetype_invariants.rs"
+
+[package.metadata.example.archetype_invariants]
+name = "Archetype Invariants"
+description = "Assertions about valid combinations of components"
+category = "ECS (Entity Component System)"
+wasm = false
+
 [[example]]
 name = "component_change_detection"
 path = "examples/ecs/component_change_detection.rs"

crates/bevy_app/src/app.rs

@@ -3,13 +3,13 @@ pub use bevy_derive::AppLabel;
 use bevy_derive::{Deref, DerefMut};
 use bevy_ecs::{
     event::{Event, Events},
-    prelude::FromWorld,
+    prelude::{Bundle, FromWorld},
     schedule::{
         IntoSystemDescriptor, Schedule, ShouldRun, Stage, StageLabel, State, StateData, SystemSet,
         SystemStage,
     },
     system::Resource,
-    world::World,
+    world::{ArchetypeInvariant, UntypedArchetypeInvariant, World},
 };
 use bevy_utils::{tracing::debug, HashMap};
 use std::fmt::Debug;
@@ -141,6 +141,32 @@ impl App {
         (runner)(app);
     }
 
+    /// Adds a new [`ArchetypeInvariant`] to this app's [`World`].
+    ///
+    /// Whenever a new archetype invariant is added to a world, all existing archetypes are re-checked.
+    /// This may include empty archetypes- archetypes that contain no entities.
+    pub fn add_archetype_invariant<B1: Bundle, B2: Bundle>(
+        &mut self,
+        archetype_invariant: ArchetypeInvariant<B1, B2>,
+    ) -> &mut Self {
+        self.world.add_archetype_invariant(archetype_invariant);
+        self
+    }
+
+    /// Inserts a new [`UntypedArchetypeInvariant`] to this app's [`World`].
+    ///
+    /// Whenever a new archetype invariant is added to a world, all existing archetypes are re-checked.
+    /// This may include empty archetypes- archetypes that contain no entities.
+    /// Prefer [`add_archetype_invariant`](App::add_archetype_invariant) where possible.
+    pub fn add_untyped_archetype_invariant(
+        &mut self,
+        archetype_invariant: UntypedArchetypeInvariant,
+    ) -> &mut Self {
+        self.world
+            .add_untyped_archetype_invariant(archetype_invariant);
+        self
+    }
+
     /// Adds a [`Stage`] with the given `label` to the last position of the app's
     /// [`Schedule`].
     ///

crates/bevy_ecs/src/archetype.rs

@@ -3,9 +3,10 @@
 
 use crate::{
     bundle::BundleId,
-    component::{ComponentId, StorageType},
+    component::{ComponentId, Components, StorageType},
     entity::{Entity, EntityLocation},
     storage::{Column, SparseArray, SparseSet, SparseSetIndex, TableId},
+    world::ArchetypeInvariants,
 };
 use std::{
     collections::HashMap,
@@ -145,20 +146,23 @@ pub struct Archetype {
 }
 
 impl Archetype {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         id: ArchetypeId,
         table_id: TableId,
         table_components: Box<[ComponentId]>,
         sparse_set_components: Box<[ComponentId]>,
         table_archetype_components: Vec<ArchetypeComponentId>,
         sparse_set_archetype_components: Vec<ArchetypeComponentId>,
+        archetype_invariants: &ArchetypeInvariants,
+        components: &Components,
     ) -> Self {
-        let mut components =
+        let mut component_set =
             SparseSet::with_capacity(table_components.len() + sparse_set_components.len());
         for (component_id, archetype_component_id) in
             table_components.iter().zip(table_archetype_components)
         {
-            components.insert(
+            component_set.insert(
                 *component_id,
                 ArchetypeComponentInfo {
                     storage_type: StorageType::Table,
@@ -171,21 +175,24 @@ impl Archetype {
             .iter()
             .zip(sparse_set_archetype_components)
         {
-            components.insert(
+            component_set.insert(
                 *component_id,
                 ArchetypeComponentInfo {
                     storage_type: StorageType::SparseSet,
                     archetype_component_id,
                 },
             );
         }
+
+        archetype_invariants.test_archetype(component_set.indices(), components);
+
         Self {
             id,
             table_info: TableInfo {
                 id: table_id,
                 entity_rows: Default::default(),
             },
-            components,
+            components: component_set,
             table_components,
             sparse_set_components,
             unique_components: SparseSet::new(),
@@ -391,7 +398,13 @@ impl Default for Archetypes {
             archetype_ids: Default::default(),
             archetype_component_count: 0,
         };
-        archetypes.get_id_or_insert(TableId::empty(), Vec::new(), Vec::new());
+        archetypes.get_id_or_insert(
+            TableId::empty(),
+            Vec::new(),
+            Vec::new(),
+            &ArchetypeInvariants::default(),
+            &Components::default(),
+        );
 
         // adds the resource archetype. it is "special" in that it is inaccessible via a "hash",
         // which prevents entities from being added to it
@@ -402,6 +415,8 @@ impl Default for Archetypes {
             Box::new([]),
             Vec::new(),
             Vec::new(),
+            &ArchetypeInvariants::default(),
+            &Components::default(),
         ));
         archetypes
     }
@@ -488,6 +503,8 @@ impl Archetypes {
         table_id: TableId,
         table_components: Vec<ComponentId>,
         sparse_set_components: Vec<ComponentId>,
+        archetype_invariants: &ArchetypeInvariants,
+        components: &Components,
     ) -> ArchetypeId {
         let table_components = table_components.into_boxed_slice();
         let sparse_set_components = sparse_set_components.into_boxed_slice();
@@ -521,6 +538,8 @@ impl Archetypes {
                     sparse_set_components,
                     table_archetype_components,
                     sparse_set_archetype_components,
+                    archetype_invariants,
+                    components,
                 ));
                 id
             })

crates/bevy_ecs/src/bundle.rs

@@ -9,6 +9,7 @@ use crate::{
     component::{Component, ComponentId, ComponentTicks, Components, StorageType},
     entity::{Entities, Entity, EntityLocation},
     storage::{SparseSetIndex, SparseSets, Storages, Table},
+    world::ArchetypeInvariants,
 };
 use bevy_ecs_macros::all_tuples;
 use bevy_ptr::OwningPtr;
@@ -260,6 +261,7 @@ impl BundleInfo {
         &self.storage_types
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn get_bundle_inserter<'a, 'b>(
         &'b self,
         entities: &'a mut Entities,
@@ -268,9 +270,15 @@ impl BundleInfo {
         storages: &'a mut Storages,
         archetype_id: ArchetypeId,
         change_tick: u32,
+        archetype_invariants: &ArchetypeInvariants,
     ) -> BundleInserter<'a, 'b> {
-        let new_archetype_id =
-            self.add_bundle_to_archetype(archetypes, storages, components, archetype_id);
+        let new_archetype_id = self.add_bundle_to_archetype(
+            archetypes,
+            storages,
+            components,
+            archetype_id,
+            archetype_invariants,
+        );
         let archetypes_ptr = archetypes.archetypes.as_mut_ptr();
         if new_archetype_id == archetype_id {
             let archetype = &mut archetypes[archetype_id];
@@ -327,9 +335,15 @@ impl BundleInfo {
         components: &mut Components,
         storages: &'a mut Storages,
         change_tick: u32,
+        archetype_invariants: &ArchetypeInvariants,
     ) -> BundleSpawner<'a, 'b> {
-        let new_archetype_id =
-            self.add_bundle_to_archetype(archetypes, storages, components, ArchetypeId::EMPTY);
+        let new_archetype_id = self.add_bundle_to_archetype(
+            archetypes,
+            storages,
+            components,
+            ArchetypeId::EMPTY,
+            archetype_invariants,
+        );
         let (empty_archetype, archetype) =
             archetypes.get_2_mut(ArchetypeId::EMPTY, new_archetype_id);
         let table = &mut storages.tables[archetype.table_id()];
@@ -399,6 +413,7 @@ impl BundleInfo {
         storages: &mut Storages,
         components: &mut Components,
         archetype_id: ArchetypeId,
+        archetype_invariants: &ArchetypeInvariants,
     ) -> ArchetypeId {
         if let Some(add_bundle) = archetypes[archetype_id].edges().get_add_bundle(self.id) {
             return add_bundle.archetype_id;
@@ -461,8 +476,13 @@ impl BundleInfo {
                     new_sparse_set_components
                 };
             };
-            let new_archetype_id =
-                archetypes.get_id_or_insert(table_id, table_components, sparse_set_components);
+            let new_archetype_id = archetypes.get_id_or_insert(
+                table_id,
+                table_components,
+                sparse_set_components,
+                archetype_invariants,
+                components,
+            );
             // add an edge from the old archetype to the new archetype
             archetypes[archetype_id].edges_mut().insert_add_bundle(
                 self.id,

crates/bevy_ecs/src/component.rs

@@ -7,6 +7,7 @@ use crate::{
 };
 pub use bevy_ecs_macros::Component;
 use bevy_ptr::OwningPtr;
+use bevy_utils::get_short_name;
 use std::{
     alloc::Layout,
     any::{Any, TypeId},
@@ -251,6 +252,22 @@ impl SparseSetIndex for ComponentId {
     }
 }
 
+/// Returns the shortened type names of the provided [`ComponentId`]s.
+///
+/// Uses [`get_short_name`] to strip the module paths of the items, resulting in cleaner lists.
+pub fn display_component_id_types<'a, I: Iterator<Item = &'a ComponentId>>(
+    component_ids: I,
+    components: &Components,
+) -> String {
+    component_ids
+        .map(|id| match components.get_info(*id) {
+            Some(info) => get_short_name(info.name()),
+            None => format!("{:?}", id),
+        })
+        .reduce(|acc, s| format!("{}, {}", acc, s))
+        .unwrap_or_default()
+}
+
 pub struct ComponentDescriptor {
     name: Cow<'static, str>,
     // SAFETY: This must remain private. It must match the statically known StorageType of the
@@ -370,6 +387,9 @@ pub struct Components {
 }
 
 impl Components {
+    /// Adds a new component type to [`Components`].
+    ///
+    /// If the component type is already present, it simply returns its [`ComponentId`].
     #[inline]
     pub fn init_component<T: Component>(&mut self, storages: &mut Storages) -> ComponentId {
         let type_id = TypeId::of::<T>();
@@ -385,6 +405,7 @@ impl Components {
         ComponentId(*index)
     }
 
+    /// Adds a new component with the provided [`ComponentDescriptor`] to [`Components`].
     pub fn init_component_with_descriptor(
         &mut self,
         storages: &mut Storages,

crates/bevy_ecs/src/lib.rs

@@ -42,7 +42,10 @@ pub mod prelude {
             NonSendMut, ParallelCommands, ParamSet, Query, RemovedComponents, Res, ResMut,
             Resource, System, SystemParamFunction,
         },
-        world::{FromWorld, Mut, World},
+        world::{
+            ArchetypeInvariant, ArchetypeInvariantHelpers, ArchetypeStatement, FromWorld, Mut,
+            World,
+        },
     };
 }