diff --git a/Cargo.toml b/Cargo.toml index 715d1e683aa03..a075b79cc90d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2051,6 +2051,18 @@ description = "Demonstrates the creation and utility of immutable components" category = "ECS (Entity Component System)" wasm = false +[[example]] +name = "index" +path = "examples/ecs/index.rs" +doc-scrape-examples = true +required-features = ["bevy_dev_tools"] + +[package.metadata.example.index] +name = "Indexes" +description = "Demonstrates querying by component value using indexing" +category = "ECS (Entity Component System)" +wasm = true + [[example]] name = "iter_combinations" path = "examples/ecs/iter_combinations.rs" diff --git a/benches/benches/bevy_ecs/index/index_iter_indexed.rs b/benches/benches/bevy_ecs/index/index_iter_indexed.rs new file mode 100644 index 0000000000000..22dc56f8989a4 --- /dev/null +++ b/benches/benches/bevy_ecs/index/index_iter_indexed.rs @@ -0,0 +1,38 @@ +use bevy_ecs::{prelude::*, system::SystemId}; +use core::hint::black_box; +use glam::*; + +const PLANETS: u16 = 1_000; +const SPAWNS: usize = 1_000_000; + +#[derive(Component, Copy, Clone, PartialEq, Eq, Hash)] +#[component(immutable)] +struct Planet(u16); + +fn find_planet_zeroes_indexed(query: QueryByIndex) { + let mut query = query.at(&Planet(0)); + for planet in query.query().iter() { + let _ = black_box(planet); + } +} + +pub struct Benchmark(World, SystemId); + +impl Benchmark { + pub fn new() -> Self { + let mut world = World::new(); + + world.add_index(IndexOptions::::default()); + + world.spawn_batch((0..PLANETS).map(Planet).cycle().take(SPAWNS)); + + let id = world.register_system(find_planet_zeroes_indexed); + + Self(world, id) + } + + #[inline(never)] + pub fn run(&mut self) { + let _ = self.0.run_system(self.1); + } +} diff --git a/benches/benches/bevy_ecs/index/index_iter_naive.rs b/benches/benches/bevy_ecs/index/index_iter_naive.rs new file mode 100644 index 0000000000000..c691d1bb38847 --- /dev/null +++ b/benches/benches/bevy_ecs/index/index_iter_naive.rs @@ -0,0 +1,35 @@ +use bevy_ecs::{prelude::*, system::SystemId}; +use core::hint::black_box; +use glam::*; + +const PLANETS: u16 = 1_000; +const SPAWNS: usize = 1_000_000; + +#[derive(Component, Copy, Clone, PartialEq, Eq, Hash)] +#[component(immutable)] +struct Planet(u16); + +fn find_planet_zeroes_naive(query: Query<&Planet>) { + for planet in query.iter().filter(|&&planet| planet == Planet(0)) { + let _ = black_box(planet); + } +} + +pub struct Benchmark(World, SystemId); + +impl Benchmark { + pub fn new() -> Self { + let mut world = World::new(); + + world.spawn_batch((0..PLANETS).map(Planet).cycle().take(SPAWNS)); + + let id = world.register_system(find_planet_zeroes_naive); + + Self(world, id) + } + + #[inline(never)] + pub fn run(&mut self) { + let _ = self.0.run_system(self.1); + } +} diff --git a/benches/benches/bevy_ecs/index/index_update_indexed.rs b/benches/benches/bevy_ecs/index/index_update_indexed.rs new file mode 100644 index 0000000000000..002deff143b0f --- /dev/null +++ b/benches/benches/bevy_ecs/index/index_update_indexed.rs @@ -0,0 +1,47 @@ +use bevy_ecs::{prelude::*, system::SystemId}; +use glam::*; + +const PLANETS: u8 = 16; +const SPAWNS: usize = 10_000; + +#[derive(Component, Copy, Clone, PartialEq, Eq, Hash)] +#[component(immutable)] +struct Planet(u8); + +fn increment_planet_zeroes_indexed( + query: QueryByIndex, + mut local: Local, + mut commands: Commands, +) { + let target = Planet(*local); + let next_planet = Planet(target.0 + 1); + + let mut query = query.at(&target); + for (entity, _planet) in query.query().iter() { + commands.entity(entity).insert(next_planet); + } + + *local += 1; +} + +pub struct Benchmark(World, SystemId); + +impl Benchmark { + pub fn new() -> Self { + let mut world = World::new(); + + world.add_index(IndexOptions::::default()); + + world.spawn_batch((0..PLANETS).map(Planet).cycle().take(SPAWNS)); + + let id = world.register_system(increment_planet_zeroes_indexed); + + Self(world, id) + } + + #[inline(never)] + pub fn run(&mut self) { + let _ = self.0.run_system(self.1); + self.0.flush(); + } +} diff --git a/benches/benches/bevy_ecs/index/index_update_naive.rs b/benches/benches/bevy_ecs/index/index_update_naive.rs new file mode 100644 index 0000000000000..8b46fedf25174 --- /dev/null +++ b/benches/benches/bevy_ecs/index/index_update_naive.rs @@ -0,0 +1,44 @@ +use bevy_ecs::{prelude::*, system::SystemId}; +use glam::*; + +const PLANETS: u8 = 16; +const SPAWNS: usize = 10_000; + +#[derive(Component, Copy, Clone, PartialEq, Eq, Hash)] +#[component(immutable)] +struct Planet(u8); + +fn increment_planet_zeroes_naive( + query: Query<(Entity, &Planet)>, + mut local: Local, + mut commands: Commands, +) { + let target = Planet(*local); + let next_planet = Planet(target.0 + 1); + + for (entity, _planet) in query.iter().filter(|(_, planet)| **planet == target) { + commands.entity(entity).insert(next_planet); + } + + *local += 1; +} + +pub struct Benchmark(World, SystemId); + +impl Benchmark { + pub fn new() -> Self { + let mut world = World::new(); + + world.spawn_batch((0..PLANETS).map(Planet).cycle().take(SPAWNS)); + + let id = world.register_system(increment_planet_zeroes_naive); + + Self(world, id) + } + + #[inline(never)] + pub fn run(&mut self) { + let _ = self.0.run_system(self.1); + self.0.flush(); + } +} diff --git a/benches/benches/bevy_ecs/index/mod.rs b/benches/benches/bevy_ecs/index/mod.rs new file mode 100644 index 0000000000000..d7beaf330c6a4 --- /dev/null +++ b/benches/benches/bevy_ecs/index/mod.rs @@ -0,0 +1,38 @@ +mod index_iter_indexed; +mod index_iter_naive; +mod index_update_indexed; +mod index_update_naive; + +use criterion::{criterion_group, Criterion}; + +criterion_group!(benches, index_iter, index_update,); + +fn index_iter(c: &mut Criterion) { + let mut group = c.benchmark_group("index_iter"); + group.warm_up_time(core::time::Duration::from_millis(500)); + group.measurement_time(core::time::Duration::from_secs(4)); + group.bench_function("naive", |b| { + let mut bench = index_iter_naive::Benchmark::new(); + b.iter(move || bench.run()); + }); + group.bench_function("indexed", |b| { + let mut bench = index_iter_indexed::Benchmark::new(); + b.iter(move || bench.run()); + }); + group.finish(); +} + +fn index_update(c: &mut Criterion) { + let mut group = c.benchmark_group("index_update"); + group.warm_up_time(core::time::Duration::from_millis(500)); + group.measurement_time(core::time::Duration::from_secs(4)); + group.bench_function("naive", |b| { + let mut bench = index_update_naive::Benchmark::new(); + b.iter(move || bench.run()); + }); + group.bench_function("indexed", |b| { + let mut bench = index_update_indexed::Benchmark::new(); + b.iter(move || bench.run()); + }); + group.finish(); +} diff --git a/benches/benches/bevy_ecs/main.rs b/benches/benches/bevy_ecs/main.rs index 4a025ab829369..cb2cd034b0b92 100644 --- a/benches/benches/bevy_ecs/main.rs +++ b/benches/benches/bevy_ecs/main.rs @@ -11,6 +11,7 @@ mod empty_archetypes; mod entity_cloning; mod events; mod fragmentation; +mod index; mod iteration; mod observers; mod param; @@ -23,6 +24,7 @@ criterion_main!( empty_archetypes::benches, entity_cloning::benches, events::benches, + index::benches, iteration::benches, fragmentation::benches, observers::benches, diff --git a/crates/bevy_app/src/app.rs b/crates/bevy_app/src/app.rs index 11602ff5ad719..6fb7e609ffcdb 100644 --- a/crates/bevy_app/src/app.rs +++ b/crates/bevy_app/src/app.rs @@ -9,8 +9,9 @@ use alloc::{ }; pub use bevy_derive::AppLabel; use bevy_ecs::{ - component::RequiredComponentsError, + component::{Immutable, RequiredComponentsError}, event::{event_update_system, EventCursor}, + index::{IndexOptions, IndexStorage}, intern::Interned, prelude::*, schedule::{ScheduleBuildSettings, ScheduleLabel}, @@ -1324,6 +1325,52 @@ impl App { self.world_mut().add_observer(observer); self } + + /// Creates and maintains an index for the provided immutable [`Component`] `C` using observers. + /// + /// This allows querying by component _value_ to be very performant, at the expense of a small + /// amount of overhead when modifying the indexed component. + /// + /// The [options](IndexOptions) provided allow control over how `C` is indexed, which may + /// allow you to improve the ergonomics or performance of an index over recommended defaults. + /// + /// # Examples + /// + /// ```rust + /// # use bevy_app::prelude::*; + /// # use bevy_ecs::prelude::*; + /// # use bevy_utils::default; + /// # + /// # let mut app = App::new(); + /// #[derive(Component, PartialEq, Eq, Hash, Clone)] + /// #[component(immutable)] + /// enum FavoriteColor { + /// Red, + /// Green, + /// Blue, + /// } + /// + /// app.add_index(IndexOptions:: { + /// // FavoriteColor only has 3 states, so can be totally addressed in 2 bits. + /// address_space: 2, + /// ..default() + /// }); + /// + /// fn find_red_fans(mut query: QueryByIndex) { + /// let mut lens = query.at(&FavoriteColor::Red); + /// for entity in lens.query().iter() { + /// println!("{entity:?} likes the color Red!"); + /// } + /// } + /// # app.add_systems(Update, find_red_fans); + /// ``` + pub fn add_index, S: IndexStorage>( + &mut self, + options: IndexOptions, + ) -> &mut Self { + self.world_mut().add_index(options); + self + } } type RunnerFn = Box AppExit>; @@ -1413,7 +1460,7 @@ impl Termination for AppExit { #[cfg(test)] mod tests { - use core::{iter, marker::PhantomData}; + use core::{hash::Hash, iter, marker::PhantomData}; use std::sync::Mutex; use bevy_ecs::{ diff --git a/crates/bevy_ecs/src/index/mod.rs b/crates/bevy_ecs/src/index/mod.rs new file mode 100644 index 0000000000000..6374a173f9a42 --- /dev/null +++ b/crates/bevy_ecs/src/index/mod.rs @@ -0,0 +1,450 @@ +//! Provides indexing support for the ECS. +//! +//! # Background +//! +//! The most common way of querying for data within the [`World`] is with [`Query`] as a system parameter. +//! This requires specifying all the parameters of your query up-front in the type-signature of system. +//! This is problematic when you don't want to query for _all_ entities with a particular set of components, +//! and instead want entities who have particular _values_ for a given component. +//! +//! Consider a `Planet` component that marks which planet an entity is on. +//! We _could_ create a unique marking component for each planet: +//! +//! ```rust +//! # use bevy_ecs::prelude::*; +//! #[derive(Component)] +//! struct Earth; +//! +//! #[derive(Component)] +//! struct Mars; +//! +//! // ... +//! ``` +//! +//! But what if the list of planets isn't knowable at compile-time and is instead controlled at runtime? +//! This would require something like: +//! +//! ```rust +//! # use bevy_ecs::prelude::*; +//! #[derive(Component, PartialEq, Eq)] +//! struct Planet(&'static str); +//! ``` +//! +//! This lets us create planets at runtime (maybe the player is the one creating them!). +//! But how do we query for this runtime-compatible `Planet`? +//! The naive approach would be to query for the `Planet` component and `filter` for a particular value. +//! +//! ```rust +//! # use bevy_ecs::prelude::*; +//! # #[derive(Component, PartialEq, Eq)] +//! # struct Planet(&'static str); +//! fn get_earthlings(mut query: Query<(Entity, &Planet)>) { +//! let earthlings = query.iter().filter(|(_, planet)| **planet == Planet("Earth")); +//! +//! for earthling in earthlings { +//! // ... +//! } +//! } +//! ``` +//! +//! The problem here is that our `get_earthlings` system reserves access to and iterates through _every_ +//! entity on _every_ planet! +//! If you have a lot of planets and a lot of entities, that's a massive bottleneck. +//! +//! _There must be a better way!_ +//! +//! # Query By Index +//! +//! Instead of filtering by value in the body of a system, we can instead use [`QueryByIndex`] and treat +//! our `Planet` as an indexable component. +//! +//! First, we need to modify `Planet` to include implementations for `Clone` and `Hash`, and to mark it as +//! an immutable component: +//! +//! ```rust +//! # use bevy_ecs::prelude::*; +//! #[derive(Component, PartialEq, Eq, Hash, Clone)] +//! #[component(immutable)] +//! struct Planet(&'static str); +//! ``` +//! +//! Next, we need to inform the world that we want `Planet` to be indexed: +//! +//! ```rust +//! # use bevy_ecs::prelude::*; +//! # #[derive(Component, PartialEq, Eq, Hash, Clone)] +//! # #[component(immutable)] +//! # struct Planet(&'static str); +//! # let mut world = World::new(); +//! world.add_index(IndexOptions::::default()); +//! ``` +//! +//! This sets up the necessary mechanisms behind the scenes to track `Planet` components and make +//! querying by value as performant as possible. +//! +//! Now we can use [`QueryByIndex`] instead of [`Query`] in our `get_earthlings` system: +//! +//! ```rust +//! # use bevy_ecs::prelude::*; +//! # #[derive(Component, PartialEq, Eq, Hash, Clone)] +//! # #[component(immutable)] +//! # struct Planet(&'static str); +//! fn get_earthlings(mut query: QueryByIndex) { +//! let mut earthlings = query.at(&Planet("Earth")); +//! +//! for earthling in &earthlings.query() { +//! // ... +//! } +//! } +//! ``` +//! +//! While this may look similar, the way this information is loaded from the ECS is completely different. +//! Instead of loading archetypes, then the entities, and comparing to our value, we first check an +//! index for our value and only load the archetypes with that value. +//! This gives us the same iteration performance as if we had created all those planet +//! marker components at compile time. +//! +//! # Drawbacks +//! +//! Indexing by a component value isn't free unfortunately. If it was, it would be enabled by default! +//! +//! ## Fragmentation +//! +//! To provide the maximum iteration speed, the indexable component is fragmented, meaning each unique +//! value is stored in its own archetype. +//! Archetypes are reused when values are no longer in use; +//! and so the cost paid scales with the maximum number of unique values alive _simultaneously_. +//! This makes iterating through a subset of the total archetypes faster, but decreases the performance +//! of iterating all archetypes by a small amount. +//! +//! This also has the potential to multiply the number of unused [`Archetypes`](crate::archetype::Archetype). +//! Since Bevy does not currently have a mechanism for cleaning up unused [`Archetypes`](crate::archetype::Archetype), +//! this can present itself like a memory leak. +//! If you find your application consuming substantially more memory when using indexing, please +//! [open an issue on GitHub](https://github.com/bevyengine/bevy/issues/new/choose) to help us +//! improve memory performance in real-world applications. +//! +//! ## Mutation Overhead +//! +//! The index is maintained continuously to ensure it is always valid. +//! This is great for usability, but means all mutations of indexed components will carry a small but +//! existent overhead. + +mod query_by_index; +mod storage; + +pub use query_by_index::*; +pub use storage::*; + +use crate::{ + self as bevy_ecs, + component::{ + Component, ComponentCloneBehavior, ComponentDescriptor, ComponentId, Immutable, StorageType, + }, + entity::Entity, + prelude::Trigger, + system::{Commands, Query, ResMut}, + world::{OnInsert, OnReplace, World}, +}; +use alloc::{boxed::Box, format, vec::Vec}; +use bevy_ecs_macros::Resource; +use bevy_platform_support::{collections::HashMap, hash::FixedHasher, sync::Arc}; +use bevy_ptr::OwningPtr; +use core::{alloc::Layout, hash::Hash, marker::PhantomData, ptr::NonNull}; +use thiserror::Error; + +/// This [`Resource`] is responsible for managing a value-to-[`ComponentId`] mapping, allowing +/// [`QueryByIndex`] to simply filter by [`ComponentId`] on a standard [`Query`]. +#[derive(Resource)] +struct Index> { + /// Maps `C` values to an index within [`slots`](Index::slots). + /// + /// We use a `Box>` instead of a concrete type parameter to ensure two indexes + /// for the same component `C` never exist. + mapping: Box>, + /// A collection of ZST dynamic [`Component`]s which (in combination) uniquely address a _value_ + /// of `C` within the [`World`]. + /// + /// We use an [`Arc`] to allow for cheap cloning by [`QueryByIndex`]. + markers: Arc<[ComponentId]>, + /// A list of liveness counts. + /// Once a value hits zero, it is free for reuse. + /// If no values are zero, you must append to the end of the list. + slots: Vec, + /// Slots with an active count of zero should be put here for later reuse. + spare_slots: Vec, +} + +/// Internal state for a [slot](Index::slots) within an [`Index`]. +struct IndexState { + /// A count of how many living [entities](Entity) exist with the world. + /// + /// Once this value reaches zero, this slot can be re-allocated for a different + /// value. + active: usize, +} + +/// Errors returned by [`track_entity`](Index::track_entity). +#[derive(Error, Debug)] +enum TrackEntityError { + /// The total address space allocated for this index has been exhausted. + #[error("address space exhausted")] + AddressSpaceExhausted, + /// An entity set to be tracked did not contain a suitable value. + #[error("entity was expected to have the indexable component but it was not found")] + EntityMissingValue, +} + +impl> Index { + fn new>(world: &mut World, options: IndexOptions) -> Self { + let bits = options + .address_space + .min(size_of::().saturating_mul(8) as u8) as u16; + + let markers = (0..bits) + .map(|bit| Self::alloc_new_marker(world, bit, options.marker_storage)) + .collect(); + + Self { + mapping: Box::new(options.index_storage), + markers, + slots: Vec::new(), + spare_slots: Vec::new(), + } + } + + fn track_entity(&mut self, world: &mut World, entity: Entity) -> Result<(), TrackEntityError> { + let Some(value) = world.get::(entity) else { + return Err(TrackEntityError::EntityMissingValue); + }; + + let slot_index = match self.mapping.get(value) { + Some(index) => { + self.slots[index].active += 1; + index + } + None => { + let spare_slot = self.spare_slots.pop(); + + match spare_slot { + Some(index) => { + self.slots[index].active += 1; + self.mapping.insert(value, index); + index + } + None => { + if self.slots.len() >= 1 << self.markers.len() { + return Err(TrackEntityError::AddressSpaceExhausted); + } + + let index = self.slots.len(); + self.mapping.insert(value, index); + self.slots.push(IndexState { active: 1 }); + + index + } + } + } + }; + + let ids = self.ids_for(slot_index); + + let zsts = core::iter::repeat_with(|| { + // SAFETY: + // - NonNull::dangling() is appropriate for a ZST + unsafe { OwningPtr::new(NonNull::dangling()) } + }) + .take(ids.len()); + + // SAFETY: + // - ids are from the same world + // - OwningPtr is valid for the entire lifetime of the application + unsafe { + world.entity_mut(entity).insert_by_ids(&ids, zsts); + } + + Ok(()) + } + + /// Observer for [`OnInsert`] events for the indexed [`Component`] `C`. + fn on_insert(trigger: Trigger, mut commands: Commands) { + let entity = trigger.target(); + + commands.queue(move |world: &mut World| { + world.resource_scope::(|world, mut index| { + if let Err(error) = index.track_entity(world, entity) { + match error { + TrackEntityError::AddressSpaceExhausted => { + log::error!( + "Entity {:?} could not be indexed by component {} as the total addressable space ({} bits) has been exhausted. Consider increasing the address space using `IndexOptions::address_space`.", + entity, + disqualified::ShortName::of::(), + index.markers.len(), + ); + }, + TrackEntityError::EntityMissingValue => { + // Swallow error. + // This was likely caused by the component `C` being removed + // before deferred commands were applied in response to the insertion. + }, + } + } + }); + }); + } + + /// Observer for [`OnReplace`] events for the indexed [`Component`] `C`. + fn on_replace( + trigger: Trigger, + query: Query<&C>, + mut index: ResMut, + mut commands: Commands, + ) { + let entity = trigger.target(); + + let value = query.get(entity).unwrap(); + + let slot_index = index.mapping.get(value).unwrap(); + + let slot = &mut index.slots[slot_index]; + + slot.active = slot.active.saturating_sub(1); + + // On removal, we check if this was the last usage of this marker. + // If so, we can recycle it for a different value + if slot.active == 0 { + index.mapping.remove(value); + index.spare_slots.push(slot_index); + } + + let ids = index.ids_for(slot_index); + + commands.queue(move |world: &mut World| { + let ids = ids; + // The old marker is no longer applicable since the value has changed/been removed. + world.entity_mut(entity).remove_by_ids(&ids); + }); + } + + /// Creates a new marker component for this index. + /// It represents a ZST and is not tied to a particular value. + /// This allows moving entities into new archetypes based on the indexed value. + fn alloc_new_marker(world: &mut World, bit: u16, storage_type: StorageType) -> ComponentId { + // SAFETY: + // - ZST is Send + Sync + // - No drop function provided or required + let descriptor = unsafe { + ComponentDescriptor::new_with_layout( + format!("{} Marker #{}", disqualified::ShortName::of::(), bit), + storage_type, + Layout::new::<()>(), + None, + false, + ComponentCloneBehavior::default(), + ) + }; + + world.register_component_with_descriptor(descriptor) + } + + /// Gets the [`ComponentId`]s of all markers that _must_ be included on an [`Entity`] allocated + /// to a particular slot index. + fn ids_for(&self, index: usize) -> Vec { + self.markers + .iter() + .enumerate() + .filter_map(|(i, &id)| (index & (1 << i) > 0).then_some(id)) + .collect::>() + } +} + +/// Options when configuring an index for a given indexable component `C`. +pub struct IndexOptions< + C: Component, + S: IndexStorage = HashMap, +> { + /// Marker components will be added to indexed entities to allow for efficient lookups. + /// This controls the [`StorageType`] that will be used with these markers. + /// + /// - [`Table`](StorageType::Table) is faster for querying + /// - [`SpareSet`](StorageType::SparseSet) is more memory efficient + /// + /// Ensure you benchmark both options appropriately if you are experiencing performance issues. + /// + /// This defaults to [`SparseSet`](StorageType::SparseSet). + pub marker_storage: StorageType, + /// Marker components are combined into a unique address for each distinct value of the indexed + /// component. + /// This controls how many markers will be used to create that unique address. + /// Note that a value greater than 32 will be reduced down to 32. + /// Bevy's [`World`] only supports 2^32 entities alive at any one moment in time, so an address space + /// of 32 is sufficient to uniquely refer to every individual [`Entity`] in the entire [`World`]. + /// + /// Selecting a value lower than the default value may lead to a panic at runtime or entities + /// missing from the index if the address space is exhausted. + /// + /// This defaults to [`size_of`] + pub address_space: u8, + /// A storage backend for this index. + /// For certain indexing strategies and [`Component`] types, you may be able to greatly + /// optimize the utility and performance of an index by creating a custom backend. + /// + /// See [`IndexStorage`] for details around the implementation of a custom backend. + /// + /// This defaults to [`HashMap`]. + pub index_storage: S, + #[doc(hidden)] + pub _phantom: PhantomData, +} + +impl + Eq + Hash + Clone> Default + for IndexOptions> +{ + fn default() -> Self { + Self { + marker_storage: StorageType::SparseSet, + address_space: size_of::().saturating_mul(8) as u8, + index_storage: HashMap::with_hasher(FixedHasher), + _phantom: PhantomData, + } + } +} + +impl, S: IndexStorage> IndexOptions { + /// Performs initial setup for an index. + // Note that this is placed here instead of inlined into `World` to allow most + // most of the indexing internals to stay private. + #[inline] + pub(crate) fn setup_index(self, world: &mut World) { + if world.get_resource::>().is_none() { + let mut index = Index::::new(world, self); + + world.query::<(Entity, &C)>() + .iter(world) + .map(|(entity, _)| entity) + .collect::>() + .into_iter() + .for_each(|entity| { + if let Err(error) = index.track_entity(world, entity) { + match error { + TrackEntityError::AddressSpaceExhausted => { + log::error!( + "Entity {:?} could not be indexed by component {} as the total addressable space ({} bits) has been exhausted. Consider increasing the address space using `IndexOptions::address_space`.", + entity, + disqualified::ShortName::of::(), + index.markers.len(), + ); + }, + TrackEntityError::EntityMissingValue => { + unreachable!(); + }, + } + } + }); + + world.insert_resource(index); + world.add_observer(Index::::on_insert); + world.add_observer(Index::::on_replace); + } + } +} diff --git a/crates/bevy_ecs/src/index/query_by_index.rs b/crates/bevy_ecs/src/index/query_by_index.rs new file mode 100644 index 0000000000000..ac8522ee9b840 --- /dev/null +++ b/crates/bevy_ecs/src/index/query_by_index.rs @@ -0,0 +1,312 @@ +use alloc::vec::Vec; + +use crate::{ + archetype::Archetype, + component::{ComponentId, Immutable, Tick}, + prelude::Component, + query::{QueryBuilder, QueryData, QueryFilter, QueryState, With}, + system::{Query, QueryLens, Res, SystemMeta, SystemParam}, + world::{unsafe_world_cell::UnsafeWorldCell, World}, +}; + +use super::Index; + +/// This system parameter allows querying by an indexable component value. +/// +/// # Examples +/// +/// ```rust +/// # use bevy_ecs::prelude::*; +/// # let mut world = World::new(); +/// #[derive(Component, PartialEq, Eq, Hash, Clone)] +/// #[component(immutable)] +/// struct Player(u8); +/// +/// // Indexing is opt-in through `World::add_index` +/// world.add_index(IndexOptions::::default()); +/// # for i in 0..6 { +/// # for _ in 0..(i + 1) { +/// # world.spawn(Player(i)); +/// # } +/// # } +/// # +/// # world.flush(); +/// +/// fn find_all_player_one_entities(by_player: QueryByIndex) { +/// let mut lens = by_player.at(&Player(0)); +/// +/// for entity in lens.query().iter() { +/// println!("{entity:?} belongs to Player 1!"); +/// } +/// # assert_eq!(( +/// # by_player.at(&Player(0)).query().iter().count(), +/// # by_player.at(&Player(1)).query().iter().count(), +/// # by_player.at(&Player(2)).query().iter().count(), +/// # by_player.at(&Player(3)).query().iter().count(), +/// # by_player.at(&Player(4)).query().iter().count(), +/// # by_player.at(&Player(5)).query().iter().count(), +/// # ), (1, 2, 3, 4, 5, 6)); +/// } +/// # world.run_system_cached(find_all_player_one_entities); +/// ``` +pub struct QueryByIndex< + 'world, + 'state, + C: Component, + D: QueryData + 'static, + F: QueryFilter + 'static = (), +> { + world: UnsafeWorldCell<'world>, + state: &'state QueryByIndexState, + last_run: Tick, + this_run: Tick, + index: Res<'world, Index>, +} + +impl, D: QueryData, F: QueryFilter> + QueryByIndex<'_, '_, C, D, F> +{ + /// Return a [`QueryLens`] returning entities with a component `C` of the provided value. + /// + /// # Examples + /// + /// ```rust + /// # use bevy_ecs::prelude::*; + /// # let mut world = World::new(); + /// #[derive(Component, PartialEq, Eq, Hash, Clone)] + /// #[component(immutable)] + /// enum FavoriteColor { + /// Red, + /// Green, + /// Blue, + /// } + /// + /// world.add_index(IndexOptions::::default()); + /// + /// fn find_red_fans(mut by_color: QueryByIndex) { + /// let mut lens = by_color.at(&FavoriteColor::Red); + /// + /// for entity in lens.query().iter() { + /// println!("{entity:?} likes the color Red!"); + /// } + /// } + /// ``` + pub fn at_mut(&mut self, value: &C) -> QueryLens<'_, D, (F, With)> + where + QueryState)>: Clone, + { + let state = self.state.primary_query_state.clone(); + + // SAFETY: We have registered all of the query's world accesses, + // so the caller ensures that `world` has permission to access any + // world data that the query needs. + unsafe { + QueryLens::new( + self.world, + self.filter_for_value(value, state), + self.last_run, + self.this_run, + ) + } + } + + /// Return a read-only [`QueryLens`] returning entities with a component `C` of the provided value. + pub fn at(&self, value: &C) -> QueryLens<'_, D::ReadOnly, (F, With)> + where + QueryState)>: Clone, + { + let state = self.state.primary_query_state.as_readonly().clone(); + + // SAFETY: We have registered all of the query's world accesses, + // so the caller ensures that `world` has permission to access any + // world data that the query needs. + unsafe { + QueryLens::new( + self.world, + self.filter_for_value(value, state), + self.last_run, + self.this_run, + ) + } + } + + fn filter_for_value( + &self, + value: &C, + mut state: QueryState, + ) -> QueryState { + match self.index.mapping.get(value) { + Some(index) => { + state = (0..self.index.markers.len()) + .map(|i| (i, 1 << i)) + .take_while(|&(_, mask)| mask <= self.index.slots.len()) + .map(|(i, mask)| { + if index & mask > 0 { + &self.state.with_states[i] + } else { + &self.state.without_states[i] + } + }) + .fold(state, |state, filter| { + state.join_filtered(self.world, filter) + }); + } + None => { + // Create a no-op filter by joining two conflicting filters together. + let filter = &self.state.with_states[0]; + state = state.join_filtered(self.world, filter); + + let filter = &self.state.without_states[0]; + state = state.join_filtered(self.world, filter); + } + } + + state + } +} + +/// The [`SystemParam::State`] type for [`QueryByIndex`], which caches information needed for queries. +/// Hidden from documentation as it is purely an implementation detail for [`QueryByIndex`] that may +/// change at any time to better suit [`QueryByIndex`]s needs. +#[doc(hidden)] +pub struct QueryByIndexState, D: QueryData, F: QueryFilter> { + primary_query_state: QueryState)>, + index_state: ComponentId, + + // TODO: Instead of storing 1 QueryState per marker component, it would be nice to instead + // track all marker components and then "somehow" create a `With`/`Without` + // query filter from that. + /// A list of [`QueryState`]s which each include a filter condition of `With`. + /// Since the marking component is dynamic, it is missing from the type signature of the state. + /// Note that this includes `With` purely for communicative purposes, `With` is a + /// strict subset of `With`. + with_states: Vec>>, + /// A list of [`QueryState`]s which each include a filter condition of `Without`. + /// Since the marking component is dynamic, it is missing from the type signature of the state. + /// Note that this includes `With` to limit the scope of this `QueryState`. + without_states: Vec>>, // No, With is not a typo +} + +impl, D: QueryData + 'static, F: QueryFilter + 'static> + QueryByIndexState +{ + fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self { + let Some(index) = world.get_resource::>() else { + panic!( + "Index not setup prior to usage. Please call `app.add_index(IndexOptions::<{}>::default())` during setup", + disqualified::ShortName::of::(), + ); + }; + + let ids = index.markers.clone(); + + let primary_query_state = + )> as SystemParam>::init_state(world, system_meta); + let index_state = > as SystemParam>::init_state(world, system_meta); + + let with_states = ids + .iter() + .map(|&id| QueryBuilder::new(world).with_id(id).build()) + .collect::>(); + + let without_states = ids + .iter() + .map(|&id| QueryBuilder::new(world).without_id(id).build()) + .collect::>(); + + Self { + primary_query_state, + index_state, + without_states, + with_states, + } + } + + unsafe fn new_archetype(&mut self, archetype: &Archetype, system_meta: &mut SystemMeta) { + )> as SystemParam>::new_archetype( + &mut self.primary_query_state, + archetype, + system_meta, + ); + + for state in self + .with_states + .iter_mut() + .chain(self.without_states.iter_mut()) + { + > as SystemParam>::new_archetype(state, archetype, system_meta); + } + } + + #[inline] + unsafe fn validate_param(&self, system_meta: &SystemMeta, world: UnsafeWorldCell) -> bool { + let mut valid = true; + + valid &= )> as SystemParam>::validate_param( + &self.primary_query_state, + system_meta, + world, + ); + valid &= + > as SystemParam>::validate_param(&self.index_state, system_meta, world); + + for state in self.with_states.iter().chain(self.without_states.iter()) { + valid &= > as SystemParam>::validate_param(state, system_meta, world); + } + + valid + } +} + +// SAFETY: We rely on the known-safe implementations of `SystemParam` for `Res` and `Query`. +unsafe impl, D: QueryData + 'static, F: QueryFilter + 'static> + SystemParam for QueryByIndex<'_, '_, C, D, F> +{ + type State = QueryByIndexState; + type Item<'w, 's> = QueryByIndex<'w, 's, C, D, F>; + + fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State { + Self::State::init_state(world, system_meta) + } + + unsafe fn new_archetype( + state: &mut Self::State, + archetype: &Archetype, + system_meta: &mut SystemMeta, + ) { + Self::State::new_archetype(state, archetype, system_meta); + } + + #[inline] + unsafe fn validate_param( + state: &Self::State, + system_meta: &SystemMeta, + world: UnsafeWorldCell, + ) -> bool { + Self::State::validate_param(state, system_meta, world) + } + + unsafe fn get_param<'world, 'state>( + state: &'state mut Self::State, + system_meta: &SystemMeta, + world: UnsafeWorldCell<'world>, + change_tick: Tick, + ) -> Self::Item<'world, 'state> { + state.primary_query_state.validate_world(world.id()); + + let index = > as SystemParam>::get_param( + &mut state.index_state, + system_meta, + world, + change_tick, + ); + + QueryByIndex { + world, + state, + last_run: system_meta.last_run, + this_run: change_tick, + index, + } + } +} diff --git a/crates/bevy_ecs/src/index/storage.rs b/crates/bevy_ecs/src/index/storage.rs new file mode 100644 index 0000000000000..b76ab2e46075c --- /dev/null +++ b/crates/bevy_ecs/src/index/storage.rs @@ -0,0 +1,59 @@ +use alloc::collections::BTreeMap; +use core::hash::{BuildHasher, Hash}; + +use bevy_platform_support::collections::HashMap; + +use crate::{component::Immutable, prelude::Component}; + +/// A storage backend for indexing the [`Component`] `C`. +/// +/// Typically, you would use a [`HashMap`] for this purpose, but you may be able to further optimize +/// the indexing performance of a particular component with a custom implementation of this backend. +/// +/// For example, the component `C` may be uniquely identifiable with a subset of its data, or implement +/// traits like [`Hash`] or [`Ord`] which would allow for specialized storage options. +pub trait IndexStorage>: 'static + Send + Sync { + /// Get the identifier of the provided value for `C`. + /// + /// Returns [`None`] if no identifier for the given value has been provided. + fn get(&self, value: &C) -> Option; + /// Sets the identifier of the provided value for `C`. + fn insert(&mut self, value: &C, index: usize); + /// Removes the identifier of the provided value for `C`. + fn remove(&mut self, value: &C); +} + +impl IndexStorage for HashMap +where + C: Component + Eq + Hash + Clone, + S: BuildHasher + Send + Sync + 'static, +{ + fn get(&self, value: &C) -> Option { + self.get(value).copied() + } + + fn insert(&mut self, value: &C, index: usize) { + self.insert(C::clone(value), index); + } + + fn remove(&mut self, value: &C) { + self.remove(value); + } +} + +impl IndexStorage for BTreeMap +where + C: Component + Ord + Clone, +{ + fn get(&self, value: &C) -> Option { + self.get(value).copied() + } + + fn insert(&mut self, value: &C, index: usize) { + self.insert(C::clone(value), index); + } + + fn remove(&mut self, value: &C) { + self.remove(value); + } +} diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 57d173b975ddd..33da191d94a51 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -43,6 +43,7 @@ pub mod entity_disabling; pub mod event; pub mod hierarchy; pub mod identifier; +pub mod index; pub mod intern; pub mod label; pub mod name; @@ -78,6 +79,7 @@ pub mod prelude { entity::{Entity, EntityBorrow, EntityMapper}, event::{Event, EventMutator, EventReader, EventWriter, Events}, hierarchy::{ChildOf, ChildSpawner, ChildSpawnerCommands, Children}, + index::{IndexOptions, QueryByIndex}, name::{Name, NameOrEntity}, observer::{Observer, Trigger}, query::{Added, AnyOf, Changed, Has, Or, QueryBuilder, QueryState, With, Without}, diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 01fe7845656aa..7ae55557f1233 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -85,6 +85,28 @@ pub struct QueryState { par_iter_span: Span, } +impl Clone for QueryState +where + D::State: Clone, + F::State: Clone, +{ + fn clone(&self) -> Self { + Self { + world_id: self.world_id, + archetype_generation: self.archetype_generation, + matched_tables: self.matched_tables.clone(), + matched_archetypes: self.matched_archetypes.clone(), + component_access: self.component_access.clone(), + matched_storage_ids: self.matched_storage_ids.clone(), + is_dense: self.is_dense, + fetch_state: self.fetch_state.clone(), + filter_state: self.filter_state.clone(), + #[cfg(feature = "trace")] + par_iter_span: self.par_iter_span.clone(), + } + } +} + impl fmt::Debug for QueryState { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("QueryState") diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index ecca479d3bf9c..362652c69d9c6 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -2098,11 +2098,35 @@ pub struct QueryLens<'w, Q: QueryData, F: QueryFilter = ()> { impl<'w, Q: QueryData, F: QueryFilter> QueryLens<'w, Q, F> { /// Create a [`Query`] from the underlying [`QueryState`]. pub fn query(&mut self) -> Query<'w, '_, Q, F> { - Query { - world: self.world, - state: &self.state, - last_run: self.last_run, - this_run: self.this_run, + // SAFETY: construction of a `QueryLens` requires ensuring the provided parameters will + // uphold the safety invariants of `Query::new` + unsafe { Query::new(self.world, &self.state, self.last_run, self.this_run) } + } + + /// Creates a new [`QueryLens`]. + /// + /// # Panics + /// + /// This will panic if the world used to create `state` is not `world`. + /// + /// # Safety + /// + /// `QueryLens` can be used to construct a `Query` by internally calling `Query::new`. + /// All safety invariants of `Query::new` must be upheld when calling `QueryLens::new`. + #[inline] + pub(crate) unsafe fn new( + world: UnsafeWorldCell<'w>, + state: QueryState, + last_run: Tick, + this_run: Tick, + ) -> Self { + state.validate_world(world.id()); + + Self { + world, + state, + last_run, + this_run, } } } diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index b7ecf5b6f3312..ed82c0d0be87f 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -36,11 +36,12 @@ use crate::{ change_detection::{MutUntyped, TicksMut}, component::{ Component, ComponentDescriptor, ComponentHooks, ComponentId, ComponentInfo, ComponentTicks, - Components, Mutable, RequiredComponents, RequiredComponentsError, Tick, + Components, Immutable, Mutable, RequiredComponents, RequiredComponentsError, Tick, }, entity::{AllocAtWithoutReplacement, Entities, Entity, EntityLocation}, entity_disabling::{DefaultQueryFilters, Disabled}, event::{Event, EventId, Events, SendBatchIds}, + index::{IndexOptions, IndexStorage}, observer::Observers, query::{DebugCheckedUnwrap, QueryData, QueryFilter, QueryState}, removal_detection::RemovedComponentEvents, @@ -3668,6 +3669,20 @@ impl World { } } +// Methods relating to component indexing +impl World { + /// Create and track an index for `C`. + /// This is required to use the [`QueryByIndex`](crate::index::QueryByIndex) system parameter. + pub fn add_index, S: IndexStorage>( + &mut self, + options: IndexOptions, + ) -> &mut Self { + options.setup_index(self); + + self + } +} + impl fmt::Debug for World { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { // SAFETY: `UnsafeWorldCell` requires that this must only access metadata. diff --git a/examples/README.md b/examples/README.md index dc562b2bf9fe4..5a6d999111381 100644 --- a/examples/README.md +++ b/examples/README.md @@ -314,6 +314,7 @@ Example | Description [Generic System](../examples/ecs/generic_system.rs) | Shows how to create systems that can be reused with different types [Hierarchy](../examples/ecs/hierarchy.rs) | Creates a hierarchy of parents and children entities [Immutable Components](../examples/ecs/immutable_components.rs) | Demonstrates the creation and utility of immutable components +[Indexes](../examples/ecs/index.rs) | Demonstrates querying by component value using indexing [Iter Combinations](../examples/ecs/iter_combinations.rs) | Shows how to iterate over combinations of query results [Nondeterministic System Order](../examples/ecs/nondeterministic_system_order.rs) | Systems run in parallel, but their order isn't always deterministic. Here's how to detect and fix this. [Observer Propagation](../examples/ecs/observer_propagation.rs) | Demonstrates event propagation with observers diff --git a/examples/ecs/index.rs b/examples/ecs/index.rs new file mode 100644 index 0000000000000..bb5dfa770ed80 --- /dev/null +++ b/examples/ecs/index.rs @@ -0,0 +1,204 @@ +//! Demonstrates how to efficiently query for a component _value_ using automatically updated indexes. +//! +//! This pattern can be substantially faster than iterating over and filtering for matching components, +//! at the expense of slower mutations and component insertion. +//! +//! See the [`bevy::ecs::index`] module docs for in-depth information. + +use bevy::{dev_tools::fps_overlay::FpsOverlayPlugin, prelude::*}; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +// To query by component value, first we need to ensure our component is suitable +// for indexing. +// +// The hard requirements are: +// * Immutability +// * Eq + Hash + Clone + +fn main() { + App::new() + .add_plugins(( + DefaultPlugins.set(WindowPlugin { + primary_window: Some(Window { + title: "Conway's Game of Life".into(), + name: Some("conway.app".into()), + resolution: (680., 700.).into(), + ..default() + }), + ..default() + }), + FpsOverlayPlugin::default(), + )) + .add_index(IndexOptions::::default()) + .insert_resource(Time::::from_seconds(0.1)) + .add_systems(Startup, setup) + .add_systems(Update, randomly_revive) + .add_systems(FixedUpdate, (spread_livelihood, update_state).chain()) + .run(); +} + +#[derive(Component, PartialEq, Eq, Hash, Clone, Copy, Debug)] +#[component(immutable, storage = "SparseSet")] +struct Alive; + +#[derive(Component, PartialEq, Eq, Hash, Clone, Copy, Debug)] +#[component(immutable)] +struct Position(i8, i8); + +/// To increase cache performance, we group a 3x3 square of cells into a chunk, and index against that. +/// If you instead indexed against the [`Position`] directly, you would have a single archetype per tile, +/// massively decreasing query performance. +#[derive(Component, PartialEq, Eq, Hash, Clone, Copy, Debug)] +#[component(immutable)] +struct Chunk(i8, i8); + +impl From for Chunk { + fn from(Position(x, y): Position) -> Self { + Chunk(x / 3, y / 3) + } +} + +#[derive(Component, PartialEq, Eq, Hash, Clone, Copy, Debug)] +struct LivingNeighbors(u8); + +#[derive(Resource)] +struct LivingHandles { + alive_material: Handle, + dead_material: Handle, +} + +#[derive(Resource)] +struct SeededRng(ChaCha8Rng); + +fn setup( + mut commands: Commands, + mut meshes: ResMut>, + mut materials: ResMut>, +) { + commands.spawn(Camera2d); + + let rect = meshes.add(Rectangle::new(10., 10.)); + let alive = materials.add(Color::BLACK); + let dead = materials.add(Color::WHITE); + + commands.insert_resource(LivingHandles { + alive_material: alive.clone(), + dead_material: dead.clone(), + }); + + // We're seeding the PRNG here to make this example deterministic for testing purposes. + // This isn't strictly required in practical use unless you need your app to be deterministic. + let seeded_rng = ChaCha8Rng::seed_from_u64(19878367467712); + commands.insert_resource(SeededRng(seeded_rng)); + + // Spawn the cells + commands.spawn_batch( + (-30..30) + .flat_map(|x| (-30..30).map(move |y| Position(x, y))) + .map(move |position| { + ( + position, + Chunk::from(position), + LivingNeighbors(0), + Mesh2d(rect.clone()), + MeshMaterial2d(dead.clone()), + Transform::from_xyz( + 11. * position.0 as f32 + 5.5, + 11. * position.1 as f32 + 5.5 - 20., + 0., + ), + ) + }), + ); +} + +fn randomly_revive( + mut commands: Commands, + mut rng: ResMut, + handles: Res, + mut query: Query< + (Entity, &mut MeshMaterial2d), + (With, Without), + >, + keyboard_input: Res>, +) { + if keyboard_input.just_pressed(KeyCode::Space) { + for (entity, mut material) in query.iter_mut() { + if rng.0.gen::() < 0.25 { + commands.entity(entity).insert(Alive); + material.0 = handles.alive_material.clone(); + } + } + } +} + +fn spread_livelihood( + mut neighbors_by_chunk: QueryByIndex, + living: Query<&Position, With>, +) { + /// In Conway's Game of Life, we consider the adjacent cells (including diagonals) as our neighbors: + /// + /// ```no_run + /// O O O O O + /// O N N N O + /// O N X N O + /// O N N N O + /// O O O O O + /// ``` + /// + /// In the above diagram, if `X` denotes a particular cell, `N` denotes neighbors, while `O` denotes a non-neighbor cell. + const NEIGHBORS: [Position; 8] = [ + // Position(0, 0), // Excluded, as this is us! + Position(0, 1), + Position(0, -1), + Position(1, 0), + Position(1, 1), + Position(1, -1), + Position(-1, 0), + Position(-1, 1), + Position(-1, -1), + ]; + + for this in living.iter() { + for delta in NEIGHBORS { + let other_pos = Position(this.0 + delta.0, this.1 + delta.1); + let mut lens = neighbors_by_chunk.at_mut(&Chunk::from(other_pos)); + let mut query = lens.query(); + + let Some((_, mut count)) = query.iter_mut().find(|(&pos, _)| pos == other_pos) else { + continue; + }; + + count.0 += 1; + } + } +} + +fn update_state( + mut commands: Commands, + mut query: Query<( + Entity, + &mut LivingNeighbors, + Has, + &mut MeshMaterial2d, + )>, + handles: Res, +) { + for (entity, mut count, alive, mut color) in query.iter_mut() { + match count.0 { + 0 | 1 | 3.. if alive => { + commands.entity(entity).remove::(); + color.0 = handles.dead_material.clone(); + } + 3 if !alive => { + commands.entity(entity).insert(Alive); + color.0 = handles.alive_material.clone(); + } + _ => {} + } + + // Reset for next frame + count.0 = 0; + } +}