RFC: Minimal Implementation of Indexed Components

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"

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<T>(PhantomData<T>);
+
+impl<T> Indexer for SimpleIndexer<T>
+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<T, F = (), I = SimpleIndexer<T>>
+where
+    I: Indexer,
+{
+    forward: HashMap<I::Index, EntityHashSet<Entity>>,
+    reverse: EntityHashMap<Entity, I::Index>,
+    last_this_run: Option<Tick>,
+    _phantom: PhantomData<fn(T, F, I)>,
+    /// Used to return an empty `impl Iterator` from `get` on the `None` branch
+    empty: EntityHashSet<Entity>,
+}
+
+impl<T, F, I> Default for IndexBacking<T, F, I>
+where
+    I: Indexer,
+{
+    fn default() -> Self {
+        Self {
+            forward: default(),
+            reverse: default(),
+            last_this_run: default(),
+            _phantom: PhantomData,
+            empty: default(),
+        }
+    }
+}
+
+impl<T, F, I> IndexBacking<T, F, I>
+where
+    I: Indexer<Input = T>,
+    I::Index: Hash + Clone + Eq,
+{
+    fn update(&mut self, entity: Entity, value: Option<&T>) -> Option<I::Index> {
+        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<I::Index> {
+        self.update(entity, Some(value))
+    }
+
+    fn remove_by_entity(&mut self, entity: Entity) -> Option<I::Index> {
+        self.update(entity, None)
+    }
+
+    fn get(&self, value: &T) -> impl Iterator<Item = Entity> + '_ {
+        self.get_by_index(&I::index(value))
+    }
+
+    fn get_by_index(&self, index: &I::Index) -> impl Iterator<Item = Entity> + '_ {
+        self.forward
+            .get(index)
+            .unwrap_or(&self.empty)
+            .iter()
+            .copied()
+    }
+
+    fn iter(
+        &mut self,
+    ) -> impl Iterator<Item = (&I::Index, impl Iterator<Item = Entity> + '_)> + '_ {
+        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<T>>
+where
+    T: Component,
+    I: Indexer + 'static,
+    F: ReadOnlyWorldQuery + 'static,
+    I::Index: Send + Sync + 'static,
+{
+    changed: Query<'w, 's, (Entity, Ref<'static, T>), (Changed<T>, F)>,
+    removed: RemovedComponents<'w, 's, T>,
+    index: ResMut<'w, IndexBacking<T, F, I>>,
+    change_tick: SystemChangeTick,
+}
+
+impl<'w, 's, T, F, I> Index<'w, 's, T, F, I>
+where
+    T: Component,
+    I: Indexer<Input = T> + '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<T, F, I>) {
+        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<Item = Entity> + '_ {
+        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<Item = Entity> + '_ {
+        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<Item = (&I::Index, impl Iterator<Item = Entity> + '_)> + '_ {
+        self.ensure_updated();
+
+        self.index.iter()
+    }
+}
+
+/// Starts indexing the [`Component`] `T`. This provides access to the [`Index`] system parameter.
+pub struct IndexPlugin<T, F = (), I = SimpleIndexer<T>>(PhantomData<fn(T, F, I)>);
+
+impl<T, F, I> Default for IndexPlugin<T, F, I> {
+    fn default() -> Self {
+        Self(PhantomData)
+    }
+}
+
+impl<T, I, F> Plugin for IndexPlugin<T, F, I>
+where
+    T: Component,
+    I: Indexer<Input = T> + 'static,
+    F: ReadOnlyWorldQuery + 'static,
+    I::Index: Hash + Clone + Eq + Send + Sync + 'static,
+{
+    fn build(&self, app: &mut App) {
+        app.init_resource::<IndexBacking<T, F, I>>()
+            .add_systems(Update, Index::<T, F, I>::update_index);
+    }
+}

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::*;
 

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

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::<Owner, With<Item>, 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<Entity, With<Player>>,
+    mut items_by_owner: Index<Owner, With<Item>, OwnerIndexer>,
+    mut printed: Local<bool>,
+) {
+    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;
+}