Skip to content

Commit d6eed58

Browse files
bushrat011899bushrat011899
committed
WIP
- Moved most of the implementation out of the new example and into an `indexing` submodule of `bevy_ecs` - Added basic documentation to satisfy linting
1 parent c39cda1 commit d6eed58

File tree

4 files changed

+197
-151
lines changed

4 files changed

+197
-151
lines changed

‎crates/bevy_ecs/src/indexing.rs

+184
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
//! Provides an [`Index`] system parameter, allowing a user to lookup an [`Entity`]
2+
//! based on the value of one of its [`Components`][`Component`].
3+
4+
use crate as bevy_ecs;
5+
use bevy_ecs::{
6+
component::{Component, Tick},
7+
prelude::{Changed, Entity, Query, Ref, RemovedComponents, ResMut},
8+
query::ReadOnlyWorldQuery,
9+
system::{SystemChangeTick, SystemParam},
10+
};
11+
12+
use bevy_ecs_macros::Resource;
13+
14+
use bevy_utils::{default, EntityHashMap, EntityHashSet, HashMap};
15+
16+
use std::{hash::Hash, marker::PhantomData};
17+
18+
/// Describes how to transform a [`Component`] `Input` into an `Index` suitable for an [`Index`].
19+
pub trait Indexer {
20+
/// The input [`Component`] to index against.
21+
type Input: Component;
22+
23+
/// A type suitable for indexing the [`Component`] `Input`
24+
type Index: Hash + Eq + Clone + Sync + Send + 'static;
25+
26+
/// Generate an `Index` from the provided `Input`
27+
fn index(input: &Self::Input) -> Self::Index;
28+
}
29+
30+
/// A basic [`Indexer`] which directly uses the [`Component`] `T`'s value.
31+
pub struct SimpleIndexer<T>(PhantomData<T>);
32+
33+
impl<T> Indexer for SimpleIndexer<T>
34+
where
35+
T: Component + Hash + Eq + Clone,
36+
{
37+
type Input = T;
38+
39+
type Index = T;
40+
41+
fn index(input: &Self::Input) -> Self::Index {
42+
input.clone()
43+
}
44+
}
45+
46+
/// Stored data required for an [`Index`].
47+
#[derive(Resource)]
48+
pub struct IndexBacking<T, F = (), I = SimpleIndexer<T>>
49+
where
50+
I: Indexer,
51+
{
52+
forward: HashMap<I::Index, EntityHashSet<Entity>>,
53+
reverse: EntityHashMap<Entity, I::Index>,
54+
last_this_run: Option<Tick>,
55+
_phantom: PhantomData<fn(T, F, I)>,
56+
/// Used to return an empty `impl Iterator` from `get` on the `None` branch
57+
empty: EntityHashSet<Entity>,
58+
}
59+
60+
impl<T, F, I> Default for IndexBacking<T, F, I>
61+
where
62+
I: Indexer,
63+
{
64+
fn default() -> Self {
65+
Self {
66+
forward: default(),
67+
reverse: default(),
68+
last_this_run: default(),
69+
_phantom: PhantomData,
70+
empty: default(),
71+
}
72+
}
73+
}
74+
75+
impl<T, F, I> IndexBacking<T, F, I>
76+
where
77+
I: Indexer<Input = T>,
78+
{
79+
fn update(&mut self, entity: Entity, value: Option<&T>) -> Option<I::Index> {
80+
let value = value.map(|value| I::index(value));
81+
82+
let old = if let Some(ref value) = value {
83+
self.reverse.insert(entity, value.clone())
84+
} else {
85+
self.reverse.remove(&entity)
86+
};
87+
88+
if let Some(ref old) = old {
89+
if let Some(set) = self.forward.get_mut(old) {
90+
set.remove(&entity);
91+
92+
if set.is_empty() {
93+
self.forward.remove(old);
94+
}
95+
}
96+
}
97+
98+
if let Some(value) = value {
99+
self.forward.entry(value).or_default().insert(entity);
100+
};
101+
102+
old
103+
}
104+
105+
fn get(&self, value: &T) -> impl Iterator<Item = Entity> + '_ {
106+
self.forward
107+
.get(&I::index(value))
108+
.unwrap_or(&self.empty)
109+
.iter()
110+
.copied()
111+
}
112+
}
113+
114+
/// Allows for lookup of an [`Entity`] based on the [`Component`] `T`'s value.
115+
/// `F` allows this [`Index`] to only target a subset of all [entities](`Entity`) using a
116+
/// [`ReadOnlyWorldQuery`].
117+
/// `I` controls how the [`Component`] `T` will be used to create an indexable value using the [`Indexer`] trait.
118+
#[derive(SystemParam)]
119+
pub struct Index<'w, 's, T, F = (), I = SimpleIndexer<T>>
120+
where
121+
T: Component,
122+
I: Indexer + 'static,
123+
F: ReadOnlyWorldQuery + 'static,
124+
{
125+
changed: Query<'w, 's, (Entity, Ref<'static, T>), (Changed<T>, F)>,
126+
removed: RemovedComponents<'w, 's, T>,
127+
index: ResMut<'w, IndexBacking<T, F, I>>,
128+
this_run: SystemChangeTick,
129+
}
130+
131+
impl<'w, 's, T, F, I> Index<'w, 's, T, F, I>
132+
where
133+
T: Component,
134+
I: Indexer<Input = T> + 'static,
135+
F: ReadOnlyWorldQuery + 'static,
136+
{
137+
fn update_index_internal(&mut self) {
138+
let this_run = self.this_run.this_run();
139+
140+
// Remove old entires
141+
for entity in self.removed.read() {
142+
self.index.update(entity, None);
143+
}
144+
145+
// Update new and existing entries
146+
for (entity, component) in self.changed.iter() {
147+
self.index.update(entity, Some(component.as_ref()));
148+
}
149+
150+
self.index.last_this_run = Some(this_run);
151+
}
152+
153+
/// System to keep [`Index`] coarsely updated every frame
154+
pub fn update_index(mut index: Index<T, F, I>) {
155+
index.update_index_internal();
156+
}
157+
158+
fn ensure_updated(&mut self) {
159+
let this_run = self.this_run.this_run();
160+
161+
if self.index.last_this_run != Some(this_run) {
162+
self.update_index_internal();
163+
}
164+
}
165+
166+
/// Get
167+
pub fn get(&mut self, value: &T) -> impl Iterator<Item = Entity> + '_ {
168+
self.ensure_updated();
169+
170+
self.index.get(value)
171+
}
172+
173+
/// Iterate over [entities](`Entity`) grouped by their [Index](`Indexer::Index`)
174+
pub fn iter(
175+
&mut self,
176+
) -> impl Iterator<Item = (&I::Index, impl Iterator<Item = Entity> + '_)> + '_ {
177+
self.ensure_updated();
178+
179+
self.index
180+
.forward
181+
.iter()
182+
.map(|(index, entities)| (index, entities.iter().copied()))
183+
}
184+
}

‎crates/bevy_ecs/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ pub mod change_detection;
1010
pub mod component;
1111
pub mod entity;
1212
pub mod event;
13+
pub mod indexing;
1314
pub mod query;
1415
#[cfg(feature = "bevy_reflect")]
1516
pub mod reflect;

‎examples/README.md

+1
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,7 @@ Example | Description
225225
[Fixed Timestep](../examples/ecs/fixed_timestep.rs) | Shows how to create systems that run every fixed timestep, rather than every tick
226226
[Generic System](../examples/ecs/generic_system.rs) | Shows how to create systems that can be reused with different types
227227
[Hierarchy](../examples/ecs/hierarchy.rs) | Creates a hierarchy of parents and children entities
228+
[Indexing](../examples/ecs/indexing.rs) | Get access to entities via component values
228229
[Iter Combinations](../examples/ecs/iter_combinations.rs) | Shows how to iterate over combinations of query results
229230
[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.
230231
[One Shot Systems](../examples/ecs/one_shot_systems.rs) | Shows how to flexibly run systems without scheduling them

‎examples/ecs/indexing.rs

+11-151
Original file line numberDiff line numberDiff line change
@@ -2,160 +2,28 @@
22
#![allow(clippy::type_complexity)]
33

44
use bevy::{
5-
ecs::{
6-
component::Tick,
7-
system::{SystemChangeTick, SystemParam},
8-
},
5+
ecs::{indexing::*, query::ReadOnlyWorldQuery},
96
prelude::*,
10-
utils::{HashMap, HashSet},
117
};
12-
use bevy_internal::ecs::query::ReadOnlyWorldQuery;
138
use std::{hash::Hash, marker::PhantomData};
149

15-
pub trait Indexer {
16-
type Input: Component;
17-
type Index: Hash + Eq + Clone + 'static;
10+
pub struct IndexPlugin<T, F = (), I = SimpleIndexer<T>>(PhantomData<fn(T, F, I)>);
1811

19-
fn index(input: &Self::Input) -> Self::Index;
20-
}
21-
22-
pub struct SimpleIndexer<T>(PhantomData<T>);
23-
24-
impl<T> Indexer for SimpleIndexer<T> where T: Component + Hash + Eq + Clone {
25-
type Input = T;
26-
27-
type Index = T;
28-
29-
fn index(input: &Self::Input) -> Self::Index {
30-
input.clone()
31-
}
32-
}
33-
34-
#[derive(Resource)]
35-
struct IndexBacking<T, F = (), I = SimpleIndexer<T>> {
36-
forward: HashMap<T, HashSet<Entity>>,
37-
reverse: HashMap<Entity, T>,
38-
last_this_run: Option<Tick>,
39-
_phantom: PhantomData<fn(F, I)>,
40-
}
41-
42-
impl<T, F, I> Default for IndexBacking<T, F, I> {
12+
impl<T, F, I> Default for IndexPlugin<T, F, I> {
4313
fn default() -> Self {
44-
Self {
45-
forward: default(),
46-
reverse: default(),
47-
last_this_run: default(),
48-
_phantom: PhantomData,
49-
}
50-
}
51-
}
52-
53-
impl<T, F> IndexBacking<T, F> {
54-
fn update(&mut self, entity: Entity, value: Option<T>) -> Option<T>
55-
where
56-
T: Hash + Eq + Clone,
57-
{
58-
let old = if let Some(ref value) = value {
59-
self.reverse.insert(entity, value.clone())
60-
} else {
61-
self.reverse.remove(&entity)
62-
};
63-
64-
if let Some(ref old) = old {
65-
if let Some(set) = self.forward.get_mut(old) {
66-
set.remove(&entity);
67-
68-
if set.is_empty() {
69-
self.forward.remove(old);
70-
}
71-
}
72-
}
73-
74-
if let Some(value) = value {
75-
self.forward.entry(value).or_default().insert(entity);
76-
};
77-
78-
old
79-
}
80-
81-
fn get(&self, value: &T) -> Option<impl Iterator<Item = Entity> + '_>
82-
where
83-
T: Hash + Eq + Clone,
84-
{
85-
Some(self.forward.get(value)?.iter().copied())
14+
Self(PhantomData)
8615
}
8716
}
8817

89-
#[derive(SystemParam)]
90-
pub struct Index<'w, 's, T, F = ()>
91-
where
92-
T: Component + Hash + Eq,
93-
F: ReadOnlyWorldQuery + 'static,
94-
{
95-
changed: Query<'w, 's, (Entity, Ref<'static, T>), (Changed<T>, F)>,
96-
removed: RemovedComponents<'w, 's, T>,
97-
index: ResMut<'w, IndexBacking<T, F>>,
98-
this_run: SystemChangeTick,
99-
}
100-
101-
impl<'w, 's, T, F> Index<'w, 's, T, F>
18+
impl<T, I, F> Plugin for IndexPlugin<T, F, I>
10219
where
103-
T: Component + Hash + Eq + Clone,
20+
T: Component,
21+
I: Indexer<Input = T> + 'static,
10422
F: ReadOnlyWorldQuery + 'static,
10523
{
106-
fn update_index_internal(&mut self) {
107-
let this_run = self.this_run.this_run();
108-
109-
// Remove old entires
110-
for entity in self.removed.read() {
111-
self.index.update(entity, None);
112-
}
113-
114-
// Update new and existing entries
115-
for (entity, component) in self.changed.iter() {
116-
self.index.update(entity, Some(component.clone()));
117-
}
118-
119-
self.index.last_this_run = Some(this_run);
120-
}
121-
122-
fn update_index(mut index: Index<T>) {
123-
index.update_index_internal();
124-
}
125-
126-
fn ensure_updated(&mut self) {
127-
let this_run = self.this_run.this_run();
128-
129-
if self.index.last_this_run != Some(this_run) {
130-
self.update_index_internal();
131-
}
132-
}
133-
134-
pub fn get(&mut self, value: &T) -> Option<impl Iterator<Item = Entity> + '_> {
135-
self.ensure_updated();
136-
137-
self.index.get(value)
138-
}
139-
140-
pub fn iter(&mut self) -> impl Iterator<Item = Entity> + '_ {
141-
self.ensure_updated();
142-
143-
self.index.reverse.keys().copied()
144-
}
145-
}
146-
147-
pub struct IndexPlugin<T>(PhantomData<T>);
148-
149-
impl<T> Default for IndexPlugin<T> {
150-
fn default() -> Self {
151-
Self(PhantomData)
152-
}
153-
}
154-
155-
impl<T> Plugin for IndexPlugin<T> where T: Component + Hash + Eq + Clone {
15624
fn build(&self, app: &mut App) {
157-
app.init_resource::<IndexBacking<T>>()
158-
.add_systems(Update, Index::<T>::update_index);
25+
app.init_resource::<IndexBacking<T, F, I>>()
26+
.add_systems(Update, Index::<T, F, I>::update_index);
15927
}
16028
}
16129

@@ -187,18 +55,10 @@ fn main() {
18755
fn get_bodies_for_head(
18856
heads: Query<(Entity, &Player), With<Head>>,
18957
bodies: Query<Entity, With<Body>>,
190-
mut index: Index<Player>
58+
mut index: Index<Player>,
19159
) {
19260
for (head_entity, head_player) in heads.iter() {
193-
let Some(body_entities) = index.get(head_player) else {
194-
continue;
195-
};
196-
197-
for body_entity in body_entities {
198-
let Ok(body_entity) = bodies.get(body_entity) else {
199-
continue;
200-
};
201-
61+
for body_entity in index.get(head_player).flat_map(|entity| bodies.get(entity)) {
20262
info!("{head_player:?}: {head_entity:?} <-> {body_entity:?}");
20363
}
20464
}

0 commit comments

Comments
 (0)