diff --git a/Cargo.toml b/Cargo.toml index bce48e8e4636d..7af78a16781b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1344,6 +1344,17 @@ description = "Creates a hierarchy of parents and children entities" category = "ECS (Entity Component System)" wasm = false +[[example]] +name = "indexing" +path = "examples/ecs/indexing.rs" +doc-scrape-examples = true + +[package.metadata.example.indexing] +name = "Indexing" +description = "Get access to entities via component values" +category = "ECS (Entity Component System)" +wasm = false + [[example]] name = "iter_combinations" path = "examples/ecs/iter_combinations.rs" diff --git a/crates/bevy_hierarchy/src/indexing.rs b/crates/bevy_hierarchy/src/indexing.rs new file mode 100644 index 0000000000000..4a6968242263a --- /dev/null +++ b/crates/bevy_hierarchy/src/indexing.rs @@ -0,0 +1,234 @@ +use std::{hash::Hash, marker::PhantomData}; + +use bevy_app::{App, Plugin, Update}; + +use bevy_ecs::{ + component::{Component, Tick}, + prelude::{Changed, Entity, Query, Ref, RemovedComponents, ResMut}, + query::ReadOnlyWorldQuery, + system::{Resource, SystemChangeTick, SystemParam}, +}; + +use bevy_utils::{default, EntityHashMap, EntityHashSet, HashMap}; + +/// Describes how to transform an `Input` into an `Index` suitable for an [`Index`]. +pub trait Indexer { + /// The input to index against. + type Input; + + /// A type suitable for indexing the `Input` + type Index; + + /// Generate an `Index` from the provided `Input` + fn index(input: &Self::Input) -> Self::Index; +} + +/// A basic [`Indexer`] which directly uses `T`'s value. +pub struct SimpleIndexer(PhantomData); + +impl Indexer for SimpleIndexer +where + T: Clone, +{ + type Input = T; + + type Index = T; + + fn index(input: &Self::Input) -> Self::Index { + input.clone() + } +} + +/// Stored data required for an [`Index`]. +#[derive(Resource)] +struct IndexBacking> +where + I: Indexer, +{ + forward: HashMap>, + reverse: EntityHashMap, + last_this_run: Option, + _phantom: PhantomData, + /// Used to return an empty `impl Iterator` from `get` on the `None` branch + empty: EntityHashSet, +} + +impl Default for IndexBacking +where + I: Indexer, +{ + fn default() -> Self { + Self { + forward: default(), + reverse: default(), + last_this_run: default(), + _phantom: PhantomData, + empty: default(), + } + } +} + +impl IndexBacking +where + I: Indexer, + I::Index: Hash + Clone + Eq, +{ + fn update(&mut self, entity: Entity, value: Option<&T>) -> Option { + let value = value.map(|value| I::index(value)); + + if self.reverse.get(&entity) == value.as_ref() { + // Return early since the value is already up-to-date + return None; + } + + let old = if let Some(ref value) = value { + self.reverse.insert(entity, value.clone()) + } else { + self.reverse.remove(&entity) + }; + + if let Some(ref old) = old { + if let Some(set) = self.forward.get_mut(old) { + set.remove(&entity); + + if set.is_empty() { + self.forward.remove(old); + } + } + } + + if let Some(value) = value { + self.forward.entry(value).or_default().insert(entity); + }; + + old + } + + fn insert(&mut self, entity: Entity, value: &T) -> Option { + self.update(entity, Some(value)) + } + + fn remove_by_entity(&mut self, entity: Entity) -> Option { + self.update(entity, None) + } + + fn get(&self, value: &T) -> impl Iterator + '_ { + self.get_by_index(&I::index(value)) + } + + fn get_by_index(&self, index: &I::Index) -> impl Iterator + '_ { + self.forward + .get(index) + .unwrap_or(&self.empty) + .iter() + .copied() + } + + fn iter( + &mut self, + ) -> impl Iterator + '_)> + '_ { + self.forward + .iter() + .map(|(index, entities)| (index, entities.iter().copied())) + } +} + +/// Allows for lookup of an [`Entity`] based on the [`Component`] `T`'s value. +/// `F` allows this [`Index`] to only target a subset of all [entities](`Entity`) using a +/// [`ReadOnlyWorldQuery`]. +/// `I` controls how the [`Component`] `T` will be used to create an indexable value using the [`Indexer`] trait. +#[derive(SystemParam)] +pub struct Index<'w, 's, T, F = (), I = SimpleIndexer> +where + T: Component, + I: Indexer + 'static, + F: ReadOnlyWorldQuery + 'static, + I::Index: Send + Sync + 'static, +{ + changed: Query<'w, 's, (Entity, Ref<'static, T>), (Changed, F)>, + removed: RemovedComponents<'w, 's, T>, + index: ResMut<'w, IndexBacking>, + change_tick: SystemChangeTick, +} + +impl<'w, 's, T, F, I> Index<'w, 's, T, F, I> +where + T: Component, + I: Indexer + 'static, + F: ReadOnlyWorldQuery + 'static, + I::Index: Hash + Clone + Eq + Send + Sync + 'static, +{ + fn update_index_internal(&mut self) { + let this_run = self.change_tick.this_run(); + + // Remove old entires + for entity in self.removed.read() { + self.index.remove_by_entity(entity); + } + + // Update new and existing entries + for (entity, component) in self.changed.iter() { + self.index.insert(entity, component.as_ref()); + } + + self.index.last_this_run = Some(this_run); + } + + /// System to keep [`Index`] coarsely updated every frame + fn update_index(mut index: Index) { + index.update_index_internal(); + } + + fn ensure_updated(&mut self) { + let this_run = self.change_tick.this_run(); + + if self.index.last_this_run != Some(this_run) { + self.update_index_internal(); + } + } + + /// Get all [entities](`Entity`) with a [`Component`] of `value`. + pub fn get(&mut self, value: &T) -> impl Iterator + '_ { + self.ensure_updated(); + + self.index.get(value) + } + + /// Get all [entities](`Entity`) with an `index`. + pub fn get_by_index(&mut self, index: &I::Index) -> impl Iterator + '_ { + self.ensure_updated(); + + self.index.get_by_index(index) + } + + /// Iterate over [entities](`Entity`) grouped by their [Index](`Indexer::Index`) + pub fn iter( + &mut self, + ) -> impl Iterator + '_)> + '_ { + self.ensure_updated(); + + self.index.iter() + } +} + +/// Starts indexing the [`Component`] `T`. This provides access to the [`Index`] system parameter. +pub struct IndexPlugin>(PhantomData); + +impl Default for IndexPlugin { + fn default() -> Self { + Self(PhantomData) + } +} + +impl Plugin for IndexPlugin +where + T: Component, + I: Indexer + 'static, + F: ReadOnlyWorldQuery + 'static, + I::Index: Hash + Clone + Eq + Send + Sync + 'static, +{ + fn build(&self, app: &mut App) { + app.init_resource::>() + .add_systems(Update, Index::::update_index); + } +} diff --git a/crates/bevy_hierarchy/src/lib.rs b/crates/bevy_hierarchy/src/lib.rs index f156086b1a926..4dbe423cb89f9 100644 --- a/crates/bevy_hierarchy/src/lib.rs +++ b/crates/bevy_hierarchy/src/lib.rs @@ -10,6 +10,9 @@ pub use components::*; mod hierarchy; pub use hierarchy::*; +mod indexing; +pub use indexing::*; + mod child_builder; pub use child_builder::*; diff --git a/examples/README.md b/examples/README.md index 0450ef88159b8..338405ee46fda 100644 --- a/examples/README.md +++ b/examples/README.md @@ -225,6 +225,7 @@ Example | Description [Fixed Timestep](../examples/ecs/fixed_timestep.rs) | Shows how to create systems that run every fixed timestep, rather than every tick [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 +[Indexing](../examples/ecs/indexing.rs) | Get access to entities via component values [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. [One Shot Systems](../examples/ecs/one_shot_systems.rs) | Shows how to flexibly run systems without scheduling them diff --git a/examples/ecs/indexing.rs b/examples/ecs/indexing.rs new file mode 100644 index 0000000000000..b62484b9b0a6c --- /dev/null +++ b/examples/ecs/indexing.rs @@ -0,0 +1,93 @@ +//! This example demonstrates how to access `Component` data through an `Index`. +#![allow(clippy::type_complexity)] + +use bevy::{ + hierarchy::{Index, IndexPlugin, Indexer}, + prelude::*, +}; + +/// Flag for an inventory item. +#[derive(Component)] +struct Item; + +/// Represents something with an owner. Similar to parent-child relationships, but distinct. +#[derive(Component)] +struct Owner(Entity); + +/// Flag for a player. +#[derive(Component)] +struct Player; + +/// Flag for an NPC. +#[derive(Component)] +struct Npc; + +/// Index [`Owner`] by the contained [`Entity`] +struct OwnerIndexer; + +impl Indexer for OwnerIndexer { + type Input = Owner; + + type Index = Entity; + + fn index(input: &Self::Input) -> Self::Index { + input.0 + } +} + +fn main() { + App::new() + .add_plugins(DefaultPlugins) + // This will only index owned items + .add_plugins(IndexPlugin::, OwnerIndexer>::default()) + .add_systems(Startup, setup) + .add_systems(FixedUpdate, print_player_items) + .run(); +} + +fn setup(mut commands: Commands) { + // Spawn a single player with 10 items, all in their possession. + let mut player = commands.spawn(Player); + let player_id = player.id(); + + player.with_children(|builder| { + for _ in 0..10 { + builder.spawn((Item, Owner(builder.parent_entity()))); + } + }); + + // Spawn 100 NPCs with 10 items each. + for _ in 0..100 { + commands.spawn(Npc).with_children(|builder| { + for _ in 0..10 { + builder.spawn((Item, Owner(builder.parent_entity()))); + } + }); + } + + // This NPC is a thief! They're holding one of the player's items + commands.spawn(Npc).with_children(|builder| { + builder.spawn((Item, Owner(player_id))); + }); +} + +fn print_player_items( + player: Query>, + mut items_by_owner: Index, OwnerIndexer>, + mut printed: Local, +) { + if *printed { + return; + } + + // This is a single-player "game" + let player = player.single(); + + // With an index, their isn't a need for an "Owned" component, akin to "Children" for "Parent". + // Instead, we can ask the index itself for what entities are "Owned" by a particular entity. + let item_count = items_by_owner.get_by_index(&player).count(); + + info!("Player owns {item_count} item(s)"); + + *printed = true; +}