Prototype Indexing

Cargo.toml

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

crates/bevy_app/src/app_builder.rs

@@ -1,3 +1,5 @@
+use std::hash::Hash;
+
 use crate::{
     app::{App, AppExit},
     event::Events,
@@ -6,6 +8,8 @@ use crate::{
 };
 use bevy_ecs::{
     component::Component,
+    prelude::Index,
+    query::indexing::index_maintanance_system,
     schedule::{
         RunOnce, Schedule, Stage, StageLabel, StateStage, SystemDescriptor, SystemSet, SystemStage,
     },
@@ -210,6 +214,22 @@ impl AppBuilder {
         })
     }
 
+    pub fn add_index<T: Component + Hash + Eq + Clone>(&mut self) -> &mut Self {
+        self.insert_resource(Index::<T>::default())
+            .add_system(index_maintanance_system::<T>.exclusive_system())
+            .add_startup_system_to_stage(
+                StartupStage::PostStartup,
+                index_maintanance_system::<T>.exclusive_system(),
+            )
+    }
+
+    pub fn add_index_sync_at<T: Component + Hash + Eq + Clone, L: StageLabel>(
+        &mut self,
+        label: L,
+    ) -> &mut Self {
+        self.add_system_to_stage(label, index_maintanance_system::<T>.exclusive_system())
+    }
+
     pub fn add_default_stages(&mut self) -> &mut Self {
         self.add_stage(
             CoreStage::Startup,

crates/bevy_ecs/src/lib.rs

@@ -16,7 +16,10 @@ pub mod prelude {
     pub use crate::{
         bundle::Bundle,
         entity::Entity,
-        query::{Added, Changed, Flags, Mutated, Or, QueryState, With, WithBundle, Without},
+        query::{
+            Added, Changed, Flags, Index, Indexed, Mutated, Or, QueryState, With, WithBundle,
+            Without,
+        },
         schedule::{
             AmbiguitySetLabel, ExclusiveSystemDescriptorCoercion, ParallelSystemDescriptorCoercion,
             Schedule, Stage, StageLabel, State, StateStage, SystemLabel, SystemStage,

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},
 };
 
@@ -365,6 +368,10 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
     }
 
     unsafe fn init(world: &World, state: &Self::State) -> Self {
+        assert!(
+            world.get_resource::<Index<T>>().is_none(),
+            "You can't directly mutate an indexed component. Use `Indexed<T>` instead of `&mut T`."
+        );
         let mut value = Self {
             storage_type: state.storage_type,
             table_components: NonNull::dangling(),
@@ -440,6 +447,140 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
     }
 }
 
+pub struct Index<T: Component> {
+    forward: HashMap<T, HashSet<Entity>>,
+}
+
+pub mod indexing {
+    use crate::prelude::{Added, Query, ResMut};
+
+    use super::*;
+    pub fn index_maintanance_system<T: Component + Eq + Hash + Clone>(
+        query: Query<(&T, Entity), Added<T>>,
+        mut index: ResMut<Index<T>>,
+    ) {
+        for (t, e) in query.iter() {
+            index.forward.entry(t.clone()).or_default().insert(e);
+        }
+    }
+
+    impl<T: Component + Eq + Hash + Clone> Index<T> {
+        pub fn get(&self, t: &T) -> Option<&HashSet<Entity>> {
+            self.forward.get(t)
+        }
+    }
+
+    impl<T: Component + Eq + Hash + Clone> Default for Index<T> {
+        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<T>(WriteFetch<T>, EntityFetch);
+
+impl<'w, T: Component + Eq + Hash + Clone> Fetch<'w> for IndexedFetch<T> {
+    type Item = Indexed<'w, T>;
+
+    type State = (WriteState<T>, EntityState);
+
+    unsafe fn init(world: &World, state: &Self::State) -> Self {
+        let (a, b) = <(WriteFetch<T>, EntityFetch) as Fetch<'w>>::init(world, state);
+        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<T>;
+    type State = (WriteState<T>, EntityState);
+}
+
+impl<'a, T: Component + Eq + Hash + Clone> Indexed<'a, T> {
+    pub fn deref<'b: 'a>(&'b mut self, index: &'b mut Index<T>) -> 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<T>,
+}
+
+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<T: WorldQuery> WorldQuery for Option<T> {
     type Fetch = OptionFetch<T::Fetch>;
     type State = OptionState<T::State>;

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::<Position>()
+        .init_resource::<Materials>()
+        .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<Materials>, asset_server: Res<AssetServer>) {
+    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<ColorMaterial>,
+    white: Handle<ColorMaterial>,
+    green: Handle<ColorMaterial>,
+}
+
+impl FromWorld for Materials {
+    fn from_world(world: &mut World) -> Self {
+        let mut materials = world.get_resource_mut::<Assets<ColorMaterial>>().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<ColorMaterial>>,
+    index: Res<Index<Position>>,
+    mut selected: ResMut<Selected>,
+    input: Res<Input<KeyCode>>,
+    mats: Res<Materials>,
+) {
+    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()
+                };
+            }
+        }
+    }
+}