Skip to content

Use Component Column ids in tables to skip component maps in queries #19063

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Use Component Column ids in tables to skip component maps in queries #19063

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
98d510a
add data
ElliottjPierce May 2, 2025
0e683b2
fixed warnings
ElliottjPierce May 2, 2025
e58a3da
add tables meta
ElliottjPierce May 4, 2025
31cb633
add getters to table
ElliottjPierce May 4, 2025
2877893
skip maps for &T
ElliottjPierce May 4, 2025
d199b21
other queries work
ElliottjPierce May 4, 2025
e1a62cd
yup. that was a dumb bug
ElliottjPierce May 4, 2025
d737ff1
skip maps for filters
ElliottjPierce May 4, 2025
1e6d0da
fix ci
ElliottjPierce May 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
187 changes: 111 additions & 76 deletions crates/bevy_ecs/src/query/fetch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use crate::{
component::{Component, ComponentId, Components, Mutable, StorageType, Tick},
entity::{Entities, Entity, EntityLocation},
query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery},
storage::{ComponentSparseSet, Table, TableRow},
storage::{ComponentSparseSet, Table, TableRow, TablesComponentMeta},
world::{
unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept,
FilteredEntityMut, FilteredEntityRef, Mut, Ref, World,
Expand Down Expand Up @@ -1102,7 +1102,10 @@ pub struct ReadFetch<'w, T: Component> {
components: StorageSwitch<
T,
// T::STORAGE_TYPE = StorageType::Table
Option<ThinSlicePtr<'w, UnsafeCell<T>>>,
(
Option<&'w TablesComponentMeta>,
Option<ThinSlicePtr<'w, UnsafeCell<T>>>,
),
// T::STORAGE_TYPE = StorageType::SparseSet
Option<&'w ComponentSparseSet>,
>,
Expand Down Expand Up @@ -1137,7 +1140,15 @@ unsafe impl<T: Component> WorldQuery for &T {
) -> ReadFetch<'w, T> {
ReadFetch {
components: StorageSwitch::new(
|| None,
|| {
let tables = &world.storages().tables;
(
tables
.get_tables_component_id(component_id)
.map(|id| tables.get_component_storage_meta(id).debug_checked_unwrap()),
None,
)
},
|| {
// SAFETY: The underlying type associated with `component_id` is `T`,
// which we are allowed to access since we registered it in `update_archetype_component_access`.
Expand Down Expand Up @@ -1174,17 +1185,23 @@ unsafe impl<T: Component> WorldQuery for &T {
#[inline]
unsafe fn set_table<'w>(
fetch: &mut ReadFetch<'w, T>,
&component_id: &ComponentId,
_component_id: &ComponentId,
table: &'w Table,
) {
let table_data = Some(
let column = fetch
.components
.table
.0
// SAFETY: If this table didn't have this component, this would not be being called.
.map(|meta| meta.column_in(table.id()).debug_checked_unwrap());
let table_data = column.map(|column| {
table
.get_data_slice_for(component_id)
.debug_checked_unwrap()
.into(),
);
// SAFETY: We just got the column from this table for this component.
.get_data_slice_for(column)
.into()
});
// SAFETY: set_table is only called when T::STORAGE_TYPE = StorageType::Table
unsafe { fetch.components.set_table(table_data) };
fetch.components.table.1 = table_data;
}

fn update_component_access(
Expand Down Expand Up @@ -1233,8 +1250,8 @@ unsafe impl<T: Component> QueryData for &T {
) -> Self::Item<'w> {
fetch.components.extract(
|table| {
// SAFETY: set_table was previously called
let table = unsafe { table.debug_checked_unwrap() };
// SAFETY: set_table was previously called and this entity must have this component, so it must not be none.
let table = unsafe { table.1.debug_checked_unwrap() };
// SAFETY: Caller ensures `table_row` is in range.
let item = unsafe { table.get(table_row.as_usize()) };
item.deref()
Expand All @@ -1261,12 +1278,15 @@ pub struct RefFetch<'w, T: Component> {
components: StorageSwitch<
T,
// T::STORAGE_TYPE = StorageType::Table
Option<(
ThinSlicePtr<'w, UnsafeCell<T>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
MaybeLocation<ThinSlicePtr<'w, UnsafeCell<&'static Location<'static>>>>,
)>,
(
Option<&'w TablesComponentMeta>,
Option<(
ThinSlicePtr<'w, UnsafeCell<T>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
MaybeLocation<ThinSlicePtr<'w, UnsafeCell<&'static Location<'static>>>>,
)>,
),
// T::STORAGE_TYPE = StorageType::SparseSet
// Can be `None` when the component has never been inserted
Option<&'w ComponentSparseSet>,
Expand Down Expand Up @@ -1304,7 +1324,15 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> {
) -> RefFetch<'w, T> {
RefFetch {
components: StorageSwitch::new(
|| None,
|| {
let tables = &world.storages().tables;
(
tables
.get_tables_component_id(component_id)
.map(|id| tables.get_component_storage_meta(id).debug_checked_unwrap()),
None,
)
},
|| {
// SAFETY: The underlying type associated with `component_id` is `T`,
// which we are allowed to access since we registered it in `update_archetype_component_access`.
Expand Down Expand Up @@ -1343,20 +1371,29 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> {
#[inline]
unsafe fn set_table<'w>(
fetch: &mut RefFetch<'w, T>,
&component_id: &ComponentId,
_component_id: &ComponentId,
table: &'w Table,
) {
let column = table.get_column(component_id).debug_checked_unwrap();
let table_data = Some((
column.get_data_slice(table.entity_count()).into(),
column.get_added_ticks_slice(table.entity_count()).into(),
column.get_changed_ticks_slice(table.entity_count()).into(),
column
.get_changed_by_slice(table.entity_count())
.map(Into::into),
));
let column = fetch
.components
.table
.0
// SAFETY: If this table didn't have this component, this would not be being called.
.map(|meta| meta.column_in(table.id()).debug_checked_unwrap());
let table_data = column.map(|column| {
// SAFETY: We just got the column for *this* table.
let column = table.get_column_by_id(column);
(
column.get_data_slice(table.entity_count()).into(),
column.get_added_ticks_slice(table.entity_count()).into(),
column.get_changed_ticks_slice(table.entity_count()).into(),
column
.get_changed_by_slice(table.entity_count())
.map(Into::into),
)
});
// SAFETY: set_table is only called when T::STORAGE_TYPE = StorageType::Table
unsafe { fetch.components.set_table(table_data) };
fetch.components.table.1 = table_data;
}

fn update_component_access(
Expand Down Expand Up @@ -1405,9 +1442,9 @@ unsafe impl<'__w, T: Component> QueryData for Ref<'__w, T> {
) -> Self::Item<'w> {
fetch.components.extract(
|table| {
// SAFETY: set_table was previously called
// SAFETY: set_table was previously called and this entity must have this component, so it must not be none.
let (table_components, added_ticks, changed_ticks, callers) =
unsafe { table.debug_checked_unwrap() };
unsafe { table.1.debug_checked_unwrap() };

// SAFETY: The caller ensures `table_row` is in range.
let component = unsafe { table_components.get(table_row.as_usize()) };
Expand Down Expand Up @@ -1456,12 +1493,15 @@ pub struct WriteFetch<'w, T: Component> {
components: StorageSwitch<
T,
// T::STORAGE_TYPE = StorageType::Table
Option<(
ThinSlicePtr<'w, UnsafeCell<T>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
MaybeLocation<ThinSlicePtr<'w, UnsafeCell<&'static Location<'static>>>>,
)>,
(
Option<&'w TablesComponentMeta>,
Option<(
ThinSlicePtr<'w, UnsafeCell<T>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
ThinSlicePtr<'w, UnsafeCell<Tick>>,
MaybeLocation<ThinSlicePtr<'w, UnsafeCell<&'static Location<'static>>>>,
)>,
),
// T::STORAGE_TYPE = StorageType::SparseSet
// Can be `None` when the component has never been inserted
Option<&'w ComponentSparseSet>,
Expand Down Expand Up @@ -1499,7 +1539,15 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T {
) -> WriteFetch<'w, T> {
WriteFetch {
components: StorageSwitch::new(
|| None,
|| {
let tables = &world.storages().tables;
(
tables
.get_tables_component_id(component_id)
.map(|id| tables.get_component_storage_meta(id).debug_checked_unwrap()),
None,
)
},
|| {
// SAFETY: The underlying type associated with `component_id` is `T`,
// which we are allowed to access since we registered it in `update_archetype_component_access`.
Expand Down Expand Up @@ -1538,20 +1586,29 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T {
#[inline]
unsafe fn set_table<'w>(
fetch: &mut WriteFetch<'w, T>,
&component_id: &ComponentId,
_component_id: &ComponentId,
table: &'w Table,
) {
let column = table.get_column(component_id).debug_checked_unwrap();
let table_data = Some((
column.get_data_slice(table.entity_count()).into(),
column.get_added_ticks_slice(table.entity_count()).into(),
column.get_changed_ticks_slice(table.entity_count()).into(),
column
.get_changed_by_slice(table.entity_count())
.map(Into::into),
));
let column = fetch
.components
.table
.0
// SAFETY: If this table didn't have this component, this would not be being called.
.map(|meta| meta.column_in(table.id()).debug_checked_unwrap());
let table_data = column.map(|column| {
// SAFETY: We just got the column for *this* table.
let column = table.get_column_by_id(column);
(
column.get_data_slice(table.entity_count()).into(),
column.get_added_ticks_slice(table.entity_count()).into(),
column.get_changed_ticks_slice(table.entity_count()).into(),
column
.get_changed_by_slice(table.entity_count())
.map(Into::into),
)
});
// SAFETY: set_table is only called when T::STORAGE_TYPE = StorageType::Table
unsafe { fetch.components.set_table(table_data) };
fetch.components.table.1 = table_data;
}

fn update_component_access(
Expand Down Expand Up @@ -1600,9 +1657,9 @@ unsafe impl<'__w, T: Component<Mutability = Mutable>> QueryData for &'__w mut T
) -> Self::Item<'w> {
fetch.components.extract(
|table| {
// SAFETY: set_table was previously called
// SAFETY: set_table was previously called and this entity must have this component, so it must not be none.
let (table_components, added_ticks, changed_ticks, callers) =
unsafe { table.debug_checked_unwrap() };
unsafe { table.1.debug_checked_unwrap() };

// SAFETY: The caller ensures `table_row` is in range.
let component = unsafe { table_components.get(table_row.as_usize()) };
Expand Down Expand Up @@ -2427,9 +2484,9 @@ unsafe impl<T: ?Sized> ReadOnlyQueryData for PhantomData<T> {}
/// [`StorageType`] of a given component.
pub(super) union StorageSwitch<C: Component, T: Copy, S: Copy> {
/// The table variant. Requires the component to be a table component.
table: T,
pub(super) table: T,
/// The sparse set variant. Requires the component to be a sparse set component.
sparse_set: S,
pub(super) sparse_set: S,
_marker: PhantomData<C>,
}

Expand All @@ -2445,28 +2502,6 @@ impl<C: Component, T: Copy, S: Copy> StorageSwitch<C, T, S> {
}
}

/// Creates a new [`StorageSwitch`] using a table variant.
///
/// # Panics
///
/// This will panic on debug builds if `C` is not a table component.
///
/// # Safety
///
/// `C` must be a table component.
#[inline]
pub unsafe fn set_table(&mut self, table: T) {
match C::STORAGE_TYPE {
StorageType::Table => self.table = table,
_ => {
#[cfg(debug_assertions)]
unreachable!();
#[cfg(not(debug_assertions))]
core::hint::unreachable_unchecked()
}
}
}

/// Fetches the internal value from the variant that corresponds to the
/// component's [`StorageType`].
pub fn extract<R>(&self, table: impl FnOnce(T) -> R, sparse_set: impl FnOnce(S) -> R) -> R {
Expand Down
Loading