diff --git a/Cargo.toml b/Cargo.toml index 9ed144a12cdd5..9519a5ff7e1a1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -294,6 +294,10 @@ path = "examples/ecs/system_param.rs" name = "timers" path = "examples/ecs/timers.rs" +[[example]] +name = "indexing" +path = "examples/ecs/indexing.rs" + # Games [[example]] name = "alien_cake_addict" diff --git a/crates/bevy_app/src/app_builder.rs b/crates/bevy_app/src/app_builder.rs index 0acd946adb5e3..09365bf11caf7 100644 --- a/crates/bevy_app/src/app_builder.rs +++ b/crates/bevy_app/src/app_builder.rs @@ -6,6 +6,8 @@ use crate::{ use bevy_ecs::{ component::{Component, ComponentDescriptor}, event::Events, + prelude::Index, + query::indexing::index_maintanance_system, schedule::{ RunOnce, Schedule, Stage, StageLabel, State, SystemDescriptor, SystemSet, SystemStage, }, @@ -208,6 +210,22 @@ impl AppBuilder { .add_system_set_to_stage(stage, State::::get_driver()) } + pub fn add_index(&mut self) -> &mut Self { + self.insert_resource(Index::::default()) + .add_system(index_maintanance_system::.exclusive_system()) + .add_startup_system_to_stage( + StartupStage::PostStartup, + index_maintanance_system::.exclusive_system(), + ) + } + + pub fn add_index_sync_at( + &mut self, + label: L, + ) -> &mut Self { + self.add_system_to_stage(label, index_maintanance_system::.exclusive_system()) + } + pub fn add_default_stages(&mut self) -> &mut Self { self.add_stage(CoreStage::First, SystemStage::parallel()) .add_stage( diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 4db81495b3866..3d520753564a5 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -18,7 +18,7 @@ pub mod prelude { bundle::Bundle, entity::Entity, event::{EventReader, EventWriter}, - query::{Added, ChangeTrackers, Changed, Or, QueryState, With, WithBundle, Without}, + query::{Added, ChangeTrackers, Changed, Or, Index, Indexed, QueryState, With, WithBundle, Without}, schedule::{ AmbiguitySetLabel, ExclusiveSystemDescriptorCoercion, ParallelSystemDescriptorCoercion, RunCriteria, RunCriteriaDescriptorCoercion, RunCriteriaLabel, RunCriteriaPiping, diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index 0350651f6808f..d1a933e0e5485 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -7,8 +7,11 @@ use crate::{ world::{Mut, World}, }; use bevy_ecs_macros::all_tuples; +use bevy_utils::{HashMap, HashSet}; use std::{ + hash::Hash, marker::PhantomData, + ops::{Deref, DerefMut}, ptr::{self, NonNull}, }; @@ -409,6 +412,10 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch { last_change_tick: u32, change_tick: u32, ) -> Self { + assert!( + world.get_resource::>().is_none(), + "You can't directly mutate an indexed component. Use `Indexed` instead of `&mut T`." + ); let mut value = Self { storage_type: state.storage_type, table_components: NonNull::dangling(), @@ -493,6 +500,150 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch { } } +pub struct Index { + forward: HashMap>, +} + +pub mod indexing { + use crate::prelude::{Added, Query, ResMut}; + + use super::*; + pub fn index_maintanance_system( + query: Query<(&T, Entity), Added>, + mut index: ResMut>, + ) { + for (t, e) in query.iter() { + index.forward.entry(t.clone()).or_default().insert(e); + } + } + + impl Index { + pub fn get(&self, t: &T) -> Option<&HashSet> { + self.forward.get(t) + } + } + + impl Default for Index { + fn default() -> Self { + Self { + forward: Default::default(), + } + } + } +} + +pub struct Indexed<'a, T: Component + Eq + Hash + Clone> { + inner: Mut<'a, T>, + entity: Entity, +} + +pub struct IndexedFetch(WriteFetch, EntityFetch); + +impl<'w, T: Component + Eq + Hash + Clone> Fetch<'w> for IndexedFetch { + type Item = Indexed<'w, T>; + + type State = (WriteState, EntityState); + + unsafe fn init( + world: &World, + state: &Self::State, + last_change_tick: u32, + change_tick: u32, + ) -> Self { + let (a, b) = <(WriteFetch, EntityFetch) as Fetch<'w>>::init( + world, + state, + last_change_tick, + change_tick, + ); + Self(a, b) + } + + fn is_dense(&self) -> bool { + self.0.is_dense() + } + + unsafe fn set_archetype( + &mut self, + state: &Self::State, + archetype: &Archetype, + tables: &Tables, + ) { + self.0.set_archetype(&state.0, archetype, tables); + self.1.set_archetype(&state.1, archetype, tables); + } + + unsafe fn set_table(&mut self, state: &Self::State, table: &Table) { + self.0.set_table(&state.0, table); + self.1.set_table(&state.1, table); + } + + unsafe fn archetype_fetch(&mut self, archetype_index: usize) -> Self::Item { + let inner = self.0.archetype_fetch(archetype_index); + let entity = self.1.archetype_fetch(archetype_index); + + Indexed { inner, entity } + } + + unsafe fn table_fetch(&mut self, table_row: usize) -> Self::Item { + let inner = self.0.table_fetch(table_row); + let entity = self.1.table_fetch(table_row); + Indexed { inner, entity } + } +} + +impl<'a, T: Component + Eq + Hash + Clone> WorldQuery for Indexed<'a, T> { + type Fetch = IndexedFetch; + type State = (WriteState, EntityState); +} + +impl<'a, T: Component + Eq + Hash + Clone> Indexed<'a, T> { + pub fn deref<'b: 'a>(&'b mut self, index: &'b mut Index) -> ActiveIndexed<'b, T> { + ActiveIndexed { + orig_value: self.inner.clone(), + indexed: self, + index, + } + } +} +pub struct ActiveIndexed<'a, T: Component + Eq + Hash + Clone> { + indexed: &'a mut Indexed<'a, T>, + orig_value: T, + index: &'a mut Index, +} + +impl<'a, T: Component + Eq + Hash + Clone> Deref for ActiveIndexed<'a, T> { + type Target = Mut<'a, T>; + + fn deref(&self) -> &Self::Target { + &self.indexed.inner + } +} + +impl<'a, T: Component + Eq + Hash + Clone> DerefMut for ActiveIndexed<'a, T> { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.indexed.inner + } +} + +impl<'a, T: Component + Eq + Hash + Clone> Drop for ActiveIndexed<'a, T> { + fn drop(&mut self) { + // Sync back to the index + let new_value = self.value.clone(); + let index = &mut self.index; + index + .forward + .get_mut(&self.orig_value) + .unwrap() + .remove(&self.indexed.entity); + index + .forward + .entry(new_value) + .or_default() + .insert(self.indexed.entity); + } +} + impl WorldQuery for Option { type Fetch = OptionFetch; type State = OptionState; diff --git a/examples/ecs/indexing.rs b/examples/ecs/indexing.rs new file mode 100644 index 0000000000000..55e0828bfa732 --- /dev/null +++ b/examples/ecs/indexing.rs @@ -0,0 +1,138 @@ +use bevy::prelude::*; + +#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)] +struct Position { + x: i32, + y: i32, +} +#[derive(Debug)] +struct Selected(Position); +struct Number(i32); + +fn main() { + App::build() + .add_plugins(DefaultPlugins) + .add_index::() + .init_resource::() + .insert_resource(Selected(Position { x: 0, y: 0 })) + .add_startup_system(setup.system()) + .add_system(update.system()) + .run(); +} + +fn setup(mut commands: Commands, materials: Res, asset_server: Res) { + commands.spawn(UiCameraBundle::default()); + let font = asset_server.load("fonts/FiraSans-Bold.ttf"); + for x in 0..8 { + for y in 0..8 { + commands + .spawn(NodeBundle { + style: Style { + position_type: PositionType::Absolute, + position: Rect { + left: Val::Px(100. * x as f32), + top: Val::Px(100. * y as f32), + ..Default::default() + }, + size: Size::new(Val::Px(100.), Val::Px(100.)), + ..Default::default() + }, + material: if (x + y) % 2 == 1 { + materials.white.clone() + } else { + materials.black.clone() + }, + ..Default::default() + }) + .with_children(|parent| { + parent + .spawn(TextBundle { + text: Text::with_section( + "0", + TextStyle { + font: font.clone(), + font_size: 90., + color: Color::RED, + }, + TextAlignment { + vertical: VerticalAlign::Center, + horizontal: HorizontalAlign::Center, + }, + ), + ..Default::default() + }) + .with(Position { x, y }) + .with(Number(0)); + }); + } + } +} + +struct Materials { + black: Handle, + white: Handle, + green: Handle, +} + +impl FromWorld for Materials { + fn from_world(world: &mut World) -> Self { + let mut materials = world.get_resource_mut::>().unwrap(); + Self { + white: materials.add(Color::WHITE.into()), + black: materials.add(Color::BLACK.into()), + green: materials.add(Color::GREEN.into()), + } + } +} + +fn update( + mut query: Query<(&Position, &Parent, &mut Number, &mut Text)>, + mut parent_query: Query<&mut Handle>, + index: Res>, + mut selected: ResMut, + input: Res>, + mats: Res, +) { + let old_selected = selected.0; + let selected = &mut selected.0; + let mut increment = 0; + for c in input.get_just_pressed() { + match c { + KeyCode::Left => selected.x -= 1, + KeyCode::Right => selected.x += 1, + KeyCode::Up => selected.y -= 1, + KeyCode::Down => selected.y += 1, + KeyCode::Return => increment += 1, + KeyCode::Back => increment -= 1, + _ => (), + } + } + selected.x = selected.x.clamp(0, 7); + selected.y = selected.y.clamp(0, 7); + dbg!(&selected); + if let Some((_, parent, mut num, mut text)) = index + .get(&selected) + .and_then(|i| i.iter().next()) + .and_then(|e| query.get_mut(*e).ok()) + { + let mut mat = parent_query.get_mut(parent.0).unwrap(); + num.0 += increment; + text.sections[0].value = num.0.to_string(); + + if *selected != old_selected { + *mat = mats.green.clone(); + if let Some((pos, parent, _, _)) = index + .get(&old_selected) + .and_then(|i| i.iter().next()) + .and_then(|e| query.get_mut(*e).ok()) + { + let mut mat = parent_query.get_mut(parent.0).unwrap(); + *mat = if (pos.x + pos.y) % 2 == 1 { + mats.white.clone() + } else { + mats.black.clone() + }; + } + } + } +}