Query by Indexed Component Value

Cargo.toml

@@ -2040,6 +2040,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"

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<Planet, &Planet>) {
+    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::<Planet>();
+
+        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);
+    }
+}

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);
+    }
+}

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<Planet, (Entity, &Planet)>,
+    mut local: Local<u8>,
+    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::<Planet>();
+
+        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();
+    }
+}

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<u8>,
+    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();
+    }
+}

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();
+}

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,

crates/bevy_app/src/app.rs

@@ -9,15 +9,16 @@ 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},
     system::{IntoObserverSystem, SystemId, SystemInput},
 };
 use bevy_platform_support::collections::HashMap;
-use core::{fmt::Debug, num::NonZero, panic::AssertUnwindSafe};
+use core::{fmt::Debug, hash::Hash, num::NonZero, panic::AssertUnwindSafe};
 use log::debug;
 use thiserror::Error;
 
@@ -1324,6 +1325,93 @@ 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.
+    ///
+    /// See also [`add_index_with_options`](App::add_index_with_options).
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use bevy_app::prelude::*;
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # let mut app = App::new();
+    /// #[derive(Component, PartialEq, Eq, Hash, Clone)]
+    /// #[component(immutable)]
+    /// enum FavoriteColor {
+    ///     Red,
+    ///     Green,
+    ///     Blue,
+    /// }
+    ///
+    /// app.add_index::<FavoriteColor>();
+    ///
+    /// fn find_red_fans(mut query: QueryByIndex<FavoriteColor, Entity>) {
+    ///     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<C: Component<Mutability = Immutable> + Eq + Hash + Clone>(
+        &mut self,
+    ) -> &mut Self {
+        self.world_mut().add_index::<C>();
+        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 finer control over how `C` is indexed, which may
+    /// allow you to improve the ergonomics or performance of an index over recommended defaults.
+    ///
+    /// See also [`add_index`](App::add_index).
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use bevy_app::prelude::*;
+    /// # use bevy_ecs::{prelude::*, index::IndexOptions};
+    /// # 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_with_options(IndexOptions::<FavoriteColor> {
+    ///     // FavoriteColor only has 3 states, so can be totally addressed in 2 bits.
+    ///     address_space: 2,
+    ///     ..default()
+    /// });
+    ///
+    /// fn find_red_fans(mut query: QueryByIndex<FavoriteColor, Entity>) {
+    ///     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_with_options<C: Component<Mutability = Immutable>, S: IndexStorage<C>>(
+        &mut self,
+        options: IndexOptions<C, S>,
+    ) -> &mut Self {
+        self.world_mut().add_index_with_options(options);
+        self
+    }
 }
 
 type RunnerFn = Box<dyn FnOnce(App) -> AppExit>;