diff --git a/Cargo.toml b/Cargo.toml index d039936d83779..951fc9dced9a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4877,3 +4877,14 @@ name = "Pan Camera" description = "Example Pan-Camera Styled Camera Controller for 2D scenes" category = "Camera" wasm = true + +[[example]] +name = "contiguous_query" +path = "examples/ecs/contiguous_query.rs" +doc-scrape-examples = true + +[package.metadata.example.contiguous_query] +name = "Contiguous Query" +description = "Demonstrates contiguous queries" +category = "ECS (Entity Component System)" +wasm = false diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_contiguous.rs b/benches/benches/bevy_ecs/iteration/iter_simple_contiguous.rs new file mode 100644 index 0000000000000..c40d59e90c3db --- /dev/null +++ b/benches/benches/bevy_ecs/iteration/iter_simple_contiguous.rs @@ -0,0 +1,47 @@ +use bevy_ecs::prelude::*; +use glam::*; + +#[derive(Component, Copy, Clone)] +struct Transform(Mat4); + +#[derive(Component, Copy, Clone)] +struct Position(Vec3); + +#[derive(Component, Copy, Clone)] +struct Rotation(Vec3); + +#[derive(Component, Copy, Clone)] +struct Velocity(Vec3); + +pub struct Benchmark<'w>(World, QueryState<(&'w Velocity, &'w mut Position)>); + +impl<'w> Benchmark<'w> { + pub fn new() -> Self { + let mut world = World::new(); + + world.spawn_batch(core::iter::repeat_n( + ( + Transform(Mat4::from_scale(Vec3::ONE)), + Position(Vec3::X), + Rotation(Vec3::X), + Velocity(Vec3::X), + ), + 10_000, + )); + + let query = world.query::<(&Velocity, &mut Position)>(); + Self(world, query) + } + + #[inline(never)] + pub fn run(&mut self) { + let mut iter = self.1.iter_mut(&mut self.0); + for (velocity, (position, mut ticks)) in iter.as_contiguous_iter().unwrap() { + for (v, p) in velocity.iter().zip(position.iter_mut()) { + p.0 += v.0; + } + // to match the iter_simple benchmark + ticks.mark_all_as_updated(); + } + } +} diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_contiguous_avx2.rs b/benches/benches/bevy_ecs/iteration/iter_simple_contiguous_avx2.rs new file mode 100644 index 0000000000000..837c92be8c190 --- /dev/null +++ b/benches/benches/bevy_ecs/iteration/iter_simple_contiguous_avx2.rs @@ -0,0 +1,65 @@ +use bevy_ecs::prelude::*; +use glam::*; + +#[derive(Component, Copy, Clone)] +struct Transform(Mat4); + +#[derive(Component, Copy, Clone)] +struct Position(Vec3); + +#[derive(Component, Copy, Clone)] +struct Rotation(Vec3); + +#[derive(Component, Copy, Clone)] +struct Velocity(Vec3); + +pub struct Benchmark<'w>(World, QueryState<(&'w Velocity, &'w mut Position)>); + +impl<'w> Benchmark<'w> { + pub fn supported() -> bool { + is_x86_feature_detected!("avx2") + } + + pub fn new() -> Option { + if !Self::supported() { + return None; + } + + let mut world = World::new(); + + world.spawn_batch(core::iter::repeat_n( + ( + Transform(Mat4::from_scale(Vec3::ONE)), + Position(Vec3::X), + Rotation(Vec3::X), + Velocity(Vec3::X), + ), + 10_000, + )); + + let query = world.query::<(&Velocity, &mut Position)>(); + Some(Self(world, query)) + } + + #[inline(never)] + pub fn run(&mut self) { + /// # Safety + /// avx2 must be supported + #[target_feature(enable = "avx2")] + fn exec(position: &mut [Position], velocity: &[Velocity]) { + for i in 0..position.len() { + position[i].0 += velocity[i].0; + } + } + + let mut iter = self.1.iter_mut(&mut self.0); + for (velocity, (position, mut ticks)) in iter.as_contiguous_iter().unwrap() { + // SAFETY: checked in new + unsafe { + exec(position, velocity); + } + // to match the iter_simple benchmark + ticks.mark_all_as_updated(); + } + } +} diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_no_detection.rs b/benches/benches/bevy_ecs/iteration/iter_simple_no_detection.rs new file mode 100644 index 0000000000000..7381d3a5db67e --- /dev/null +++ b/benches/benches/bevy_ecs/iteration/iter_simple_no_detection.rs @@ -0,0 +1,42 @@ +use bevy_ecs::prelude::*; +use glam::*; + +#[derive(Component, Copy, Clone)] +struct Transform(Mat4); + +#[derive(Component, Copy, Clone)] +struct Position(Vec3); + +#[derive(Component, Copy, Clone)] +struct Rotation(Vec3); + +#[derive(Component, Copy, Clone)] +struct Velocity(Vec3); + +pub struct Benchmark<'w>(World, QueryState<(&'w Velocity, &'w mut Position)>); + +impl<'w> Benchmark<'w> { + pub fn new() -> Self { + let mut world = World::new(); + + world.spawn_batch(core::iter::repeat_n( + ( + Transform(Mat4::from_scale(Vec3::ONE)), + Position(Vec3::X), + Rotation(Vec3::X), + Velocity(Vec3::X), + ), + 10_000, + )); + + let query = world.query::<(&Velocity, &mut Position)>(); + Self(world, query) + } + + #[inline(never)] + pub fn run(&mut self) { + for (velocity, mut position) in self.1.iter_mut(&mut self.0) { + position.bypass_change_detection().0 += velocity.0; + } + } +} diff --git a/benches/benches/bevy_ecs/iteration/iter_simple_no_detection_contiguous.rs b/benches/benches/bevy_ecs/iteration/iter_simple_no_detection_contiguous.rs new file mode 100644 index 0000000000000..ca8209bff2b5f --- /dev/null +++ b/benches/benches/bevy_ecs/iteration/iter_simple_no_detection_contiguous.rs @@ -0,0 +1,45 @@ +use bevy_ecs::prelude::*; +use glam::*; + +#[derive(Component, Copy, Clone)] +struct Transform(Mat4); + +#[derive(Component, Copy, Clone)] +struct Position(Vec3); + +#[derive(Component, Copy, Clone)] +struct Rotation(Vec3); + +#[derive(Component, Copy, Clone)] +struct Velocity(Vec3); + +pub struct Benchmark<'w>(World, QueryState<(&'w Velocity, &'w mut Position)>); + +impl<'w> Benchmark<'w> { + pub fn new() -> Self { + let mut world = World::new(); + + world.spawn_batch(core::iter::repeat_n( + ( + Transform(Mat4::from_scale(Vec3::ONE)), + Position(Vec3::X), + Rotation(Vec3::X), + Velocity(Vec3::X), + ), + 10_000, + )); + + let query = world.query::<(&Velocity, &mut Position)>(); + Self(world, query) + } + + #[inline(never)] + pub fn run(&mut self) { + let mut iter = self.1.iter_mut(&mut self.0); + for (velocity, (position, _ticks)) in iter.as_contiguous_iter().unwrap() { + for (v, p) in velocity.iter().zip(position.iter_mut()) { + p.0 += v.0; + } + } + } +} diff --git a/benches/benches/bevy_ecs/iteration/mod.rs b/benches/benches/bevy_ecs/iteration/mod.rs index b296c5ce0b091..7867507f62f67 100644 --- a/benches/benches/bevy_ecs/iteration/mod.rs +++ b/benches/benches/bevy_ecs/iteration/mod.rs @@ -8,11 +8,16 @@ mod iter_frag_sparse; mod iter_frag_wide; mod iter_frag_wide_sparse; mod iter_simple; +mod iter_simple_contiguous; +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +mod iter_simple_contiguous_avx2; mod iter_simple_foreach; mod iter_simple_foreach_hybrid; mod iter_simple_foreach_sparse_set; mod iter_simple_foreach_wide; mod iter_simple_foreach_wide_sparse_set; +mod iter_simple_no_detection; +mod iter_simple_no_detection_contiguous; mod iter_simple_sparse_set; mod iter_simple_system; mod iter_simple_wide; @@ -40,6 +45,27 @@ fn iter_simple(c: &mut Criterion) { let mut bench = iter_simple::Benchmark::new(); b.iter(move || bench.run()); }); + group.bench_function("base_contiguous", |b| { + let mut bench = iter_simple_contiguous::Benchmark::new(); + b.iter(move || bench.run()); + }); + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + { + if iter_simple_contiguous_avx2::Benchmark::supported() { + group.bench_function("base_contiguous_avx2", |b| { + let mut bench = iter_simple_contiguous_avx2::Benchmark::new().unwrap(); + b.iter(move || bench.run()); + }); + } + } + group.bench_function("base_no_detection", |b| { + let mut bench = iter_simple_no_detection::Benchmark::new(); + b.iter(move || bench.run()); + }); + group.bench_function("base_no_detection_contiguous", |b| { + let mut bench = iter_simple_no_detection_contiguous::Benchmark::new(); + b.iter(move || bench.run()); + }); group.bench_function("wide", |b| { let mut bench = iter_simple_wide::Benchmark::new(); b.iter(move || bench.run()); diff --git a/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs b/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs index 9c2c5832e5f26..ce4ea266db9f6 100644 --- a/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs +++ b/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs @@ -12,7 +12,7 @@ struct MutableUnmarked { #[derive(QueryData)] #[query_data(mut)] -//~^ ERROR: invalid attribute, expected `mutable` or `derive` +//~^ ERROR: invalid attribute, expected `mutable`, `derive` or `contiguous` struct MutableInvalidAttribute { a: &'static mut Foo, } diff --git a/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.stderr b/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.stderr index ec71c112a6a05..8884a6d184a0b 100644 --- a/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.stderr +++ b/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.stderr @@ -1,4 +1,4 @@ -error: invalid attribute, expected `mutable` or `derive` +error: invalid attribute, expected `mutable`, `derive` or `contiguous` --> tests/ui/world_query_derive.rs:14:14 | 14 | #[query_data(mut)] diff --git a/crates/bevy_ecs/macros/src/query_data.rs b/crates/bevy_ecs/macros/src/query_data.rs index 1e9dea94bbca3..f89e639763103 100644 --- a/crates/bevy_ecs/macros/src/query_data.rs +++ b/crates/bevy_ecs/macros/src/query_data.rs @@ -3,7 +3,8 @@ use proc_macro::TokenStream; use proc_macro2::{Ident, Span}; use quote::{format_ident, quote}; use syn::{ - parse_macro_input, parse_quote, punctuated::Punctuated, token::Comma, DeriveInput, Meta, + parse_macro_input, parse_quote, punctuated::Punctuated, token::Comma, Attribute, DeriveInput, + Fields, ImplGenerics, Member, Meta, Type, TypeGenerics, Visibility, WhereClause, }; use crate::{ @@ -15,11 +16,15 @@ use crate::{ struct QueryDataAttributes { pub is_mutable: bool, + pub is_contiguous_mutable: bool, + pub is_contiguous_immutable: bool, + pub derive_args: Punctuated, } static MUTABLE_ATTRIBUTE_NAME: &str = "mutable"; static DERIVE_ATTRIBUTE_NAME: &str = "derive"; +static CONTIGUOUS_ATTRIBUTE_NAME: &str = "contiguous"; mod field_attr_keywords { syn::custom_keyword!(ignore); @@ -27,6 +32,94 @@ mod field_attr_keywords { pub static QUERY_DATA_ATTRIBUTE_NAME: &str = "query_data"; +fn contiguous_item_struct( + path: &syn::Path, + fields: &Fields, + derive_macro_call: &proc_macro2::TokenStream, + struct_name: &Ident, + visibility: &Visibility, + item_struct_name: &Ident, + field_types: &Vec, + user_impl_generics_with_world_and_state: &ImplGenerics, + field_attrs: &Vec>, + field_visibilities: &Vec, + field_members: &Vec, + user_ty_generics: &TypeGenerics, + user_ty_generics_with_world_and_state: &TypeGenerics, + user_where_clauses_with_world_and_state: Option<&WhereClause>, +) -> proc_macro2::TokenStream { + let item_attrs = quote! { + #[doc = concat!( + "Automatically generated [`ContiguousQueryData`](", + stringify!(#path), + "::fetch::ContiguousQueryData) item type for [`", + stringify!(#struct_name), + "`], returned when iterating over contiguous query results", + )] + #[automatically_derived] + }; + + match fields { + Fields::Named(_) => quote! { + #derive_macro_call + #item_attrs + #visibility struct #item_struct_name #user_impl_generics_with_world_and_state #user_where_clauses_with_world_and_state { + #(#(#field_attrs)* #field_visibilities #field_members: <#field_types as #path::query::ContiguousQueryData>::Contiguous<'__w, '__s>,)* + } + }, + Fields::Unnamed(_) => quote! { + #derive_macro_call + #item_attrs + #visibility struct #item_struct_name #user_impl_generics_with_world_and_state #user_where_clauses_with_world_and_state ( + #( #field_visibilities <#field_types as #path::query::ContiguousQueryData>::Contiguous<'__w, '__s>, )* + ) + }, + Fields::Unit => quote! { + #item_attrs + #visibility type #item_struct_name #user_ty_generics_with_world_and_state = #struct_name #user_ty_generics; + }, + } +} + +fn contiguous_query_data_impl( + path: &syn::Path, + struct_name: &Ident, + contiguous_item_struct_name: &Ident, + field_types: &Vec, + user_impl_generics: &ImplGenerics, + user_ty_generics: &TypeGenerics, + user_ty_generics_with_world_and_state: &TypeGenerics, + field_members: &Vec, + field_aliases: &Vec, + user_where_clauses: Option<&WhereClause>, +) -> proc_macro2::TokenStream { + quote! { + // SAFETY: Individual `fetch_contiguous` are called. + unsafe impl #user_impl_generics #path::query::ContiguousQueryData for #struct_name #user_ty_generics #user_where_clauses { + type Contiguous<'__w, '__s> = #contiguous_item_struct_name #user_ty_generics_with_world_and_state; + + unsafe fn fetch_contiguous<'__w, '__s>( + _state: &'__s ::State, + _fetch: &mut ::Fetch<'__w>, + _entities: &'__w [#path::entity::Entity], + _offset: usize, + ) -> Self::Contiguous<'__w, '__s> { + #contiguous_item_struct_name { + #( + #field_members: + <#field_types>::fetch_contiguous( + &_state.#field_aliases, + &mut _fetch.#field_aliases, + _entities, + _offset, + ), + )* + } + } + } + } +} + pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { let tokens = input.clone(); @@ -48,8 +141,24 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { attributes.derive_args.push(Meta::Path(meta.path)); Ok(()) }) + } else if meta.path.is_ident(CONTIGUOUS_ATTRIBUTE_NAME) { + meta.parse_nested_meta(|meta| { + if meta.path.is_ident("all") { + attributes.is_contiguous_mutable = true; + attributes.is_contiguous_immutable = true; + Ok(()) + } else if meta.path.is_ident("mutable") { + attributes.is_contiguous_mutable = true; + Ok(()) + } else if meta.path.is_ident("immutable") { + attributes.is_contiguous_immutable = true; + Ok(()) + } else { + Err(meta.error("invalid target, expected `all`, `mutable` or `immutable`")) + } + }) } else { - Err(meta.error(format_args!("invalid attribute, expected `{MUTABLE_ATTRIBUTE_NAME}` or `{DERIVE_ATTRIBUTE_NAME}`"))) + Err(meta.error(format_args!("invalid attribute, expected `{MUTABLE_ATTRIBUTE_NAME}`, `{DERIVE_ATTRIBUTE_NAME}` or `{CONTIGUOUS_ATTRIBUTE_NAME}`"))) } }); @@ -94,6 +203,19 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { } else { item_struct_name.clone() }; + let contiguous_item_struct_name = if attributes.is_contiguous_mutable { + Ident::new(&format!("{struct_name}ContiguousItem"), Span::call_site()) + } else { + item_struct_name.clone() + }; + let read_only_contiguous_item_struct_name = if attributes.is_contiguous_immutable { + Ident::new( + &format!("{struct_name}ReadOnlyContiguousItem"), + Span::call_site(), + ) + } else { + item_struct_name.clone() + }; let fetch_struct_name = Ident::new(&format!("{struct_name}Fetch"), Span::call_site()); let fetch_struct_name = ensure_no_collision(fetch_struct_name, tokens.clone()); @@ -124,7 +246,7 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { .members() .map(|m| format_ident!("field{}", m)) .collect(); - let field_types: Vec = fields.iter().map(|f| f.ty.clone()).collect(); + let field_types: Vec = fields.iter().map(|f| f.ty.clone()).collect(); let read_only_field_types = field_types .iter() .map(|ty| parse_quote!(<#ty as #path::query::QueryData>::ReadOnly)) @@ -167,6 +289,43 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { user_where_clauses_with_world, ); + let (mutable_contiguous_item_struct, mutable_contiguous_impl) = + if attributes.is_contiguous_mutable { + let contiguous_item_struct = contiguous_item_struct( + &path, + fields, + &derive_macro_call, + &struct_name, + &visibility, + &contiguous_item_struct_name, + &field_types, + &user_impl_generics_with_world_and_state, + &field_attrs, + &field_visibilities, + &field_members, + &user_ty_generics, + &user_ty_generics_with_world_and_state, + user_where_clauses_with_world_and_state, + ); + + let contiguous_impl = contiguous_query_data_impl( + &path, + &struct_name, + &contiguous_item_struct_name, + &field_types, + &user_impl_generics, + &user_ty_generics, + &user_ty_generics_with_world_and_state, + &field_members, + &field_aliases, + user_where_clauses, + ); + + (contiguous_item_struct, contiguous_impl) + } else { + (quote! {}, quote! {}) + }; + let (read_only_struct, read_only_impl) = if attributes.is_mutable { // If the query is mutable, we need to generate a separate readonly version of some things let readonly_item_struct = item_struct( @@ -226,6 +385,43 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { (quote! {}, quote! {}) }; + let (read_only_contiguous_item_struct, read_only_contiguous_impl) = + if attributes.is_mutable && attributes.is_contiguous_immutable { + let contiguous_item_struct = contiguous_item_struct( + &path, + fields, + &derive_macro_call, + &read_only_struct_name, + &visibility, + &read_only_contiguous_item_struct_name, + &read_only_field_types, + &user_impl_generics_with_world_and_state, + &field_attrs, + &field_visibilities, + &field_members, + &user_ty_generics, + &user_ty_generics_with_world_and_state, + user_where_clauses_with_world_and_state, + ); + + let contiguous_impl = contiguous_query_data_impl( + &path, + &read_only_struct_name, + &read_only_contiguous_item_struct_name, + &read_only_field_types, + &user_impl_generics, + &user_ty_generics, + &user_ty_generics_with_world_and_state, + &field_members, + &field_aliases, + user_where_clauses, + ); + + (contiguous_item_struct, contiguous_impl) + } else { + (quote! {}, quote! {}) + }; + let data_impl = { let read_only_data_impl = if attributes.is_mutable { quote! { @@ -391,6 +587,10 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { #read_only_struct + #mutable_contiguous_item_struct + + #read_only_contiguous_item_struct + const _: () = { #[doc(hidden)] #[doc = concat!( @@ -412,6 +612,10 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream { #data_impl #read_only_data_impl + + #mutable_contiguous_impl + + #read_only_contiguous_impl }; #[allow(dead_code)] diff --git a/crates/bevy_ecs/src/change_detection/params.rs b/crates/bevy_ecs/src/change_detection/params.rs index e66d62864a604..bfa1646f23d66 100644 --- a/crates/bevy_ecs/src/change_detection/params.rs +++ b/crates/bevy_ecs/src/change_detection/params.rs @@ -3,8 +3,9 @@ use crate::{ ptr::PtrMut, resource::Resource, }; -use bevy_ptr::{Ptr, UnsafeCellDeref}; +use bevy_ptr::{Ptr, ThinSlicePtr, UnsafeCellDeref}; use core::{ + cell::UnsafeCell, ops::{Deref, DerefMut}, panic::Location, }; @@ -463,6 +464,82 @@ impl<'w, T: ?Sized> Mut<'w, T> { } } +/// Used by [`Mut`] for [`crate::query::ContiguousQueryData`] to allow marking component's changes +pub struct ContiguousComponentTicks<'w, const MUTABLE: bool> { + added: ThinSlicePtr<'w, UnsafeCell>, + changed: ThinSlicePtr<'w, UnsafeCell>, + changed_by: MaybeLocation>>>, + count: usize, + last_run: Tick, + this_run: Tick, +} + +impl<'w> ContiguousComponentTicks<'w, true> { + /// Returns mutable changed ticks slice + pub fn get_changed_ticks_mut(&mut self) -> &mut [Tick] { + // SAFETY: `changed` slice is `self.count` long, aliasing rules are uphold by `new`. + unsafe { self.changed.as_mut_slice_unchecked(self.count) } + } + + /// Marks all components as updated + pub fn mark_all_as_updated(&mut self) { + let this_run = self.this_run; + + for i in 0..self.count { + // SAFETY: `changed_by` slice is `self.count` long, aliasing rules are uphold by `new` + self.changed_by + .map(|v| unsafe { v.get_unchecked(i).deref_mut() }) + .assign(MaybeLocation::caller()); + } + + for t in self.get_changed_ticks_mut() { + *t = this_run; + } + } +} + +impl<'w, const MUTABLE: bool> ContiguousComponentTicks<'w, MUTABLE> { + pub(crate) unsafe fn new( + added: ThinSlicePtr<'w, UnsafeCell>, + changed: ThinSlicePtr<'w, UnsafeCell>, + changed_by: MaybeLocation>>>, + count: usize, + last_run: Tick, + this_run: Tick, + ) -> Self { + Self { + added, + changed, + count, + changed_by, + last_run, + this_run, + } + } + + /// Returns immutable changed ticks slice + pub fn get_changed_ticks(&self) -> &[Tick] { + // SAFETY: `self.changed` is `self.count` long + unsafe { self.changed.cast::().as_slice_unchecked(self.count) } + } + + /// Returns immutable added ticks slice + pub fn get_added_ticks(&self) -> &[Tick] { + // SAFETY: `self.added` is `self.count` long + unsafe { self.added.cast::().as_slice_unchecked(self.count) } + } + + /// Returns the last tick system ran + pub fn last_run(&self) -> Tick { + self.last_run + } + + /// Returns the current tick + pub fn this_run(&self) -> Tick { + self.this_run + } +} + impl<'w, T: ?Sized> From> for Ref<'w, T> { fn from(mut_ref: Mut<'w, T>) -> Self { Self { diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index d672b6ae4d0b7..0310e10812837 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -1,7 +1,9 @@ use crate::{ archetype::{Archetype, Archetypes}, bundle::Bundle, - change_detection::{ComponentTicksMut, ComponentTicksRef, MaybeLocation, Tick}, + change_detection::{ + ComponentTicksMut, ComponentTicksRef, ContiguousComponentTicks, MaybeLocation, Tick, + }, component::{Component, ComponentId, Components, Mutable, StorageType}, entity::{Entities, Entity, EntityLocation}, query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery}, @@ -96,14 +98,16 @@ use variadics_please::all_tuples; /// /// ## Macro expansion /// -/// Expanding the macro will declare one or three additional structs, depending on whether or not the struct is marked as mutable. +/// Expanding the macro will declare one to five additional structs, depending on whether or not the struct is marked as mutable or as contiguous. /// For a struct named `X`, the additional structs will be: /// -/// |Struct name|`mutable` only|Description| -/// |:---:|:---:|---| -/// |`XItem`|---|The type of the query item for `X`| -/// |`XReadOnlyItem`|✓|The type of the query item for `XReadOnly`| -/// |`XReadOnly`|✓|[`ReadOnly`] variant of `X`| +/// |Struct name|`mutable` only|`contiguous` target|Description| +/// |:---:|:---:|:---:|---| +/// |`XItem`|---|---|The type of the query item for `X`| +/// |`XReadOnlyItem`|✓|---|The type of the query item for `XReadOnly`| +/// |`XReadOnly`|✓|---|[`ReadOnly`] variant of `X`| +/// |`XContiguousItem`|---|`mutable` or `all`|The type of the contiguous query item for `X`| +/// |`XReadOnlyContiguousItem`|✓|`immutable` or `all`|The type of the contiguous query item for `XReadOnly`| /// /// ## Adding mutable references /// @@ -139,11 +143,38 @@ use variadics_please::all_tuples; /// } /// ``` /// +/// ## Adding contiguous items +/// +/// To create contiguous items additionally, the struct must be marked with the `#[query_data(contiguous(target))]` attribute, +/// where the target may be `all`, `mutable` or `immutable` (see the table above). +/// +/// For mutable queries it may be done like that: +/// ``` +/// # use bevy_ecs::prelude::*; +/// # use bevy_ecs::query::QueryData; +/// # +/// # #[derive(Component)] +/// # struct ComponentA; +/// # +/// #[derive(QueryData)] +/// /// - contiguous(all) will create contiguous items for both read and mutable versions +/// /// - contiguous(mutable) will only create a contiguous item for the mutable version +/// /// - contiguous(immutable) will only create a contiguous item for the read only version +/// #[query_data(mutable, contiguous(all))] +/// struct CustomQuery { +/// component_a: &'static mut ComponentA, +/// } +/// ``` +/// +/// For immutable queries `contiguous(immutable)` attribute will be **ignored**, meanwhile `contiguous(mutable)` and `contiguous(all)` +/// will only generate a contiguous item for the (original) read only version. +/// /// ## Adding methods to query items /// /// It is possible to add methods to query items in order to write reusable logic about related components. /// This will often make systems more readable because low level logic is moved out from them. -/// It is done by adding `impl` blocks with methods for the `-Item` or `-ReadOnlyItem` generated structs. +/// It is done by adding `impl` blocks with methods for the `-Item`, `-ReadOnlyItem`, `-ContiguousItem` or `ContiguousReadOnlyItem` +/// generated structs. /// /// ``` /// # use bevy_ecs::prelude::*; @@ -208,7 +239,7 @@ use variadics_please::all_tuples; /// # struct ComponentA; /// # /// #[derive(QueryData)] -/// #[query_data(mutable, derive(Debug))] +/// #[query_data(mutable, derive(Debug), contiguous(all))] /// struct CustomQuery { /// component_a: &'static ComponentA, /// } @@ -218,6 +249,8 @@ use variadics_please::all_tuples; /// /// assert_debug::(); /// assert_debug::(); +/// assert_debug::(); +/// assert_debug::(); /// ``` /// /// ## Query composition @@ -344,6 +377,44 @@ pub unsafe trait QueryData: WorldQuery { ) -> Option>; } +/// A [`QueryData`] which allows getting a direct access to contiguous chunks of components' values +/// +// NOTE: The safety rules might not be used to optimize the library, it still might be better to ensure +// that contiguous query data methods match their non-contiguous versions +// NOTE: Even though all component references (&T, &mut T) implement this trait, it won't be executed for +// SparseSet components because in that case the query is not dense. +/// # Safety +/// +/// - The result of [`ContiguousQueryData::fetch_contiguous`] must represent the same result as if +/// [`QueryData::fetch`] was executed for each entity of the set table +#[diagnostic::on_unimplemented( + message = "`{Self}` cannot be iterated contiguously", + label = "invalid contiguous `Query` data", + note = "if `{Self}` is a custom query type, using `QueryData` derive macro, ensure that the `#[query_data(contiguous(target))]` attribute is added" +)] +pub unsafe trait ContiguousQueryData: ArchetypeQueryData { + /// Item returned by [`ContiguousQueryData::fetch_contiguous`]. + /// Represents a contiguous chunk of memory. + type Contiguous<'w, 's>; + + /// Fetch [`ContiguousQueryData::Contiguous`] which represents a contiguous chunk of memory (e.g., an array) in the current [`Table`]. + /// This must always be called after [`WorldQuery::set_table`]. + /// + /// # Safety + /// + /// - Must always be called _after_ [`WorldQuery::set_table`]. + /// - `entities`'s length must match the length of the set table. + /// - `entities` must match the entities of the set table. + /// - `offset` must be less than the length of the set table. + /// - There must not be simultaneous conflicting component access registered in `update_component_access`. + unsafe fn fetch_contiguous<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's>; +} + /// A [`QueryData`] that is read only. /// /// # Safety @@ -463,6 +534,20 @@ impl ReleaseStateQueryData for Entity { impl ArchetypeQueryData for Entity {} +/// SAFETY: matches the [`QueryData::fetch`] implementation +unsafe impl ContiguousQueryData for Entity { + type Contiguous<'w, 's> = &'w [Entity]; + + unsafe fn fetch_contiguous<'w, 's>( + _state: &'s Self::State, + _fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + &entities[offset..] + } +} + /// SAFETY: /// `update_component_access` does nothing. /// This is sound because `fetch` does not access components. @@ -1670,6 +1755,37 @@ unsafe impl QueryData for &T { } } +/// SAFETY: The result represents all values of [`Self`] in the set table. +unsafe impl ContiguousQueryData for &T { + type Contiguous<'w, 's> = &'w [T]; + + unsafe fn fetch_contiguous<'w, 's>( + _state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + fetch.components.extract( + |table| { + // SAFETY: set_table was previously called + let table = unsafe { table.debug_checked_unwrap() }; + // UnsafeCell has the same alignment as T because of transparent representation + // (i.e. repr(transparent)) of UnsafeCell + let table = table.cast::(); + // SAFETY: Caller ensures `rows` is the amount of rows in the table + let item = unsafe { table.as_slice_unchecked(entities.len()) }; + &item[offset..] + }, + |_| { + #[cfg(debug_assertions)] + unreachable!(); + #[cfg(not(debug_assertions))] + core::hint::unreachable_unchecked(); + }, + ) + } +} + /// SAFETY: access is read only unsafe impl ReadOnlyQueryData for &T {} @@ -1894,6 +2010,46 @@ impl ReleaseStateQueryData for Ref<'_, T> { impl ArchetypeQueryData for Ref<'_, T> {} +/// SAFETY: Refer to [`&mut T`]'s implementation +unsafe impl ContiguousQueryData for Ref<'_, T> { + type Contiguous<'w, 's> = (&'w [T], ContiguousComponentTicks<'w, false>); + + unsafe fn fetch_contiguous<'w, 's>( + _state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + fetch.components.extract( + |table| { + // SAFETY: set_table was previously called + let (table_components, added_ticks, changed_ticks, callers) = + unsafe { table.debug_checked_unwrap() }; + + ( + &table_components + .cast::() + .as_slice_unchecked(entities.len())[offset..], + ContiguousComponentTicks::<'w, false>::new( + added_ticks.add_unchecked(offset), + changed_ticks.add_unchecked(offset), + callers.map(|callers| callers.add_unchecked(offset)), + entities.len() - offset, + fetch.last_run, + fetch.this_run, + ), + ) + }, + |_| { + #[cfg(debug_assertions)] + unreachable!(); + #[cfg(not(debug_assertions))] + core::hint::unreachable_unchecked(); + }, + ) + } +} + /// The [`WorldQuery::Fetch`] type for `&mut T`. pub struct WriteFetch<'w, T: Component> { components: StorageSwitch< @@ -2104,6 +2260,46 @@ impl> ReleaseStateQueryData for &mut T { impl> ArchetypeQueryData for &mut T {} +/// SAFETY: +/// - The first element of [`ContiguousQueryData::Contiguous`] tuple represents all components' values in the set table. +/// - The second element of [`ContiguousQueryData::Contiguous`] tuple represents all components' ticks in the set table. +unsafe impl> ContiguousQueryData for &mut T { + type Contiguous<'w, 's> = (&'w mut [T], ContiguousComponentTicks<'w, true>); + + unsafe fn fetch_contiguous<'w, 's>( + _state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + fetch.components.extract( + |table| { + // SAFETY: set_table was previously called + let (table_components, added_ticks, changed_ticks, callers) = + unsafe { table.debug_checked_unwrap() }; + + ( + &mut table_components.as_mut_slice_unchecked(entities.len())[offset..], + ContiguousComponentTicks::<'w, true>::new( + added_ticks.add_unchecked(offset), + changed_ticks.add_unchecked(offset), + callers.map(|callers| callers.add_unchecked(offset)), + entities.len() - offset, + fetch.last_run, + fetch.this_run, + ), + ) + }, + |_| { + #[cfg(debug_assertions)] + unreachable!(); + #[cfg(not(debug_assertions))] + core::hint::unreachable_unchecked(); + }, + ) + } +} + /// When `Mut` is used in a query, it will be converted to `Ref` when transformed into its read-only form, providing access to change detection methods. /// /// By contrast `&mut T` will result in a `Mut` item in mutable form to record mutations, but result in a bare `&T` in read-only form. @@ -2219,6 +2415,20 @@ impl> ReleaseStateQueryData for Mut<'_, T> { impl> ArchetypeQueryData for Mut<'_, T> {} +/// SAFETY: Refer to soundness of `&mut T` implementation +unsafe impl<'__w, T: Component> ContiguousQueryData for Mut<'__w, T> { + type Contiguous<'w, 's> = (&'w mut [T], ContiguousComponentTicks<'w, true>); + + unsafe fn fetch_contiguous<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + <&mut T as ContiguousQueryData>::fetch_contiguous(state, fetch, entities, offset) + } +} + #[doc(hidden)] pub struct OptionFetch<'w, T: WorldQuery> { fetch: T::Fetch<'w>, @@ -2372,6 +2582,23 @@ impl ReleaseStateQueryData for Option { // so it's always an `ArchetypeQueryData`, even for non-archetypal `T`. impl ArchetypeQueryData for Option {} +// SAFETY: matches the [`QueryData::fetch`] impl +unsafe impl ContiguousQueryData for Option { + type Contiguous<'w, 's> = Option>; + + unsafe fn fetch_contiguous<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + fetch + .matches + // SAFETY: The invariants are upheld by the caller + .then(|| unsafe { T::fetch_contiguous(state, &mut fetch.fetch, entities, offset) }) + } +} + /// Returns a bool that describes if an entity has the component `T`. /// /// This can be used in a [`Query`](crate::system::Query) if you want to know whether or not entities @@ -2546,6 +2773,20 @@ impl ReleaseStateQueryData for Has { impl ArchetypeQueryData for Has {} +/// SAFETY: matches [`QueryData::fetch`] +unsafe impl ContiguousQueryData for Has { + type Contiguous<'w, 's> = bool; + + unsafe fn fetch_contiguous<'w, 's>( + _state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + _entities: &'w [Entity], + _offset: usize, + ) -> Self::Contiguous<'w, 's> { + *fetch + } +} + /// The `AnyOf` query parameter fetches entities with any of the component types included in T. /// /// `Query>` is equivalent to `Query<(Option<&A>, Option<&B>, Option<&mut C>), Or<(With, With, With)>>`. @@ -2631,6 +2872,40 @@ macro_rules! impl_tuple_query_data { $(#[$meta])* impl<$($name: ArchetypeQueryData),*> ArchetypeQueryData for ($($name,)*) {} + + #[expect( + clippy::allow_attributes, + reason = "This is a tuple-related macro; as such the lints below may not always apply." + )] + #[allow( + non_snake_case, + reason = "The names of some variables are provided by the macro's caller, not by us." + )] + #[allow( + unused_variables, + reason = "Zero-length tuples won't use any of the parameters." + )] + #[allow( + clippy::unused_unit, + reason = "Zero-length tuples will generate some function bodies equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." + )] + $(#[$meta])* + // SAFETY: The returned result represents the result of individual fetches. + unsafe impl<$($name: ContiguousQueryData),*> ContiguousQueryData for ($($name,)*) { + type Contiguous<'w, 's> = ($($name::Contiguous::<'w, 's>,)*); + + unsafe fn fetch_contiguous<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + let ($($state,)*) = state; + let ($($name,)*) = fetch; + // SAFETY: The invariants are upheld by the caller. + ($(unsafe {$name::fetch_contiguous($state, $name, entities, offset)},)*) + } + } }; } @@ -2827,6 +3102,44 @@ macro_rules! impl_anytuple_fetch { $(#[$meta])* impl<$($name: ArchetypeQueryData),*> ArchetypeQueryData for AnyOf<($($name,)*)> {} + + + #[expect( + clippy::allow_attributes, + reason = "This is a tuple-related macro; as such the lints below may not always apply." + )] + #[allow( + non_snake_case, + reason = "The names of some variables are provided by the macro's caller, not by us." + )] + #[allow( + unused_variables, + reason = "Zero-length tuples won't use any of the parameters." + )] + #[allow( + clippy::unused_unit, + reason = "Zero-length tuples will generate some function bodies equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case." + )] + $(#[$meta])* + // SAFETY: Matches the fetch implementation + unsafe impl<$($name: ContiguousQueryData),*> ContiguousQueryData for AnyOf<($($name,)*)> { + type Contiguous<'w, 's> = ($(Option<$name::Contiguous<'w,'s>>,)*); + + unsafe fn fetch_contiguous<'w, 's>( + state: &'s Self::State, + fetch: &mut Self::Fetch<'w>, + entities: &'w [Entity], + offset: usize, + ) -> Self::Contiguous<'w, 's> { + let ($($name,)*) = fetch; + let ($($state,)*) = state; + // Matches the [`QueryData::fetch`] except it always returns Some + ($( + // SAFETY: The invariants are upheld by the caller + $name.1.then(|| unsafe { $name::fetch_contiguous($state, &mut $name.0, entities, offset) }), + )*) + } + } }; } @@ -3094,6 +3407,7 @@ impl Copy for StorageSwitch {} mod tests { use super::*; use crate::change_detection::DetectChanges; + use crate::query::Without; use crate::system::{assert_is_system, Query}; use bevy_ecs::prelude::Schedule; use bevy_ecs_macros::QueryData; @@ -3338,4 +3652,160 @@ mod tests { // we want EntityRef to use the change ticks of the system schedule.run(&mut world); } + + #[test] + fn test_contiguous_query_data() { + #[derive(Component, PartialEq, Eq, Debug)] + pub struct C(i32); + + #[derive(Component, PartialEq, Eq, Debug)] + pub struct D(bool); + + let mut world = World::new(); + world.spawn((C(0), D(true))); + world.spawn((C(1), D(false))); + world.spawn(C(2)); + + let mut query = world.query::<(&C, &D)>(); + let mut iter = query.iter(&world); + let mut iter = iter.as_contiguous_iter().unwrap(); + let c = iter.next().unwrap(); + assert_eq!(c.0, [C(0), C(1)].as_slice()); + assert_eq!(c.1, [D(true), D(false)].as_slice()); + assert!(iter.next().is_none()); + + let mut query = world.query::<&C>(); + let mut iter = query.iter(&world); + let mut iter = iter.as_contiguous_iter().unwrap(); + let mut present = [false; 3]; + let mut len = 0; + for _ in 0..2 { + let c = iter.next().unwrap(); + for c in c { + present[c.0 as usize] = true; + len += 1; + } + } + assert!(iter.next().is_none()); + assert_eq!(len, 3); + assert_eq!(present, [true; 3]); + + let mut query = world.query::<&mut C>(); + let mut iter = query.iter_mut(&mut world); + let mut iter = iter.as_contiguous_iter().unwrap(); + for _ in 0..2 { + let c = iter.next().unwrap(); + for c in c.0 { + c.0 *= 2; + } + } + assert!(iter.next().is_none()); + let mut iter = query.iter(&world); + let mut iter = iter.as_contiguous_iter().unwrap(); + let mut present = [false; 6]; + let mut len = 0; + for _ in 0..2 { + let c = iter.next().unwrap(); + for c in c { + present[c.0 as usize] = true; + len += 1; + } + } + assert_eq!(present, [true, false, true, false, true, false]); + assert_eq!(len, 3); + + let mut query = world.query_filtered::<&C, Without>(); + let mut iter = query.iter(&world); + let mut iter = iter.as_contiguous_iter().unwrap(); + assert_eq!(iter.next().unwrap(), &[C(4)]); + assert!(iter.next().is_none()); + } + + #[test] + fn sparse_set_contiguous_query() { + #[derive(Component, Debug, PartialEq, Eq)] + #[component(storage = "SparseSet")] + pub struct S(i32); + + let mut world = World::new(); + world.spawn(S(0)); + + let mut query = world.query::<&mut S>(); + let mut iter = query.iter_mut(&mut world); + assert!(iter.as_contiguous_iter().is_none()); + assert_eq!(iter.next().unwrap().as_ref(), &S(0)); + } + + #[test] + fn any_of_contiguous_test() { + #[derive(Component, Debug, Clone, Copy)] + pub struct C(i32); + + #[derive(Component, Debug, Clone, Copy)] + pub struct D(i32); + + let mut world = World::new(); + world.spawn((C(0), D(1))); + world.spawn(C(2)); + world.spawn(D(3)); + world.spawn(()); + + let mut query = world.query::>(); + let mut iter = query.iter(&world); + let mut present = [false; 4]; + + for (c, d) in iter.as_contiguous_iter().unwrap() { + assert!(c.is_some() || d.is_some()); + let c = c.unwrap_or(&[]); + let d = d.unwrap_or(&[]); + for i in 0..c.len().max(d.len()) { + let c = c.get(i).cloned(); + let d = d.get(i).cloned(); + if let Some(C(c)) = c { + assert!(!present[c as usize]); + present[c as usize] = true; + } + if let Some(D(d)) = d { + assert!(!present[d as usize]); + present[d as usize] = true; + } + } + } + + assert_eq!(present, [true; 4]); + } + + #[test] + fn option_contiguous_test() { + #[derive(Component, Clone, Copy)] + struct C(i32); + + #[derive(Component, Clone, Copy)] + struct D(i32); + + let mut world = World::new(); + world.spawn((C(0), D(1))); + world.spawn(D(2)); + world.spawn(C(3)); + + let mut query = world.query::<(Option<&C>, &D)>(); + let mut iter = query.iter(&world); + let mut present = [false; 3]; + + for (c, d) in iter.as_contiguous_iter().unwrap() { + let c = c.unwrap_or(&[]); + for i in 0..d.len() { + let c = c.get(i).cloned(); + let D(d) = d[i]; + if let Some(C(c)) = c { + assert!(!present[c as usize]); + present[c as usize] = true; + } + assert!(!present[d as usize]); + present[d as usize] = true; + } + } + + assert_eq!(present, [true; 3]); + } } diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index f8578d5d21705..86c0941ef589c 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -563,7 +563,6 @@ macro_rules! impl_tuple_query_filter { true $(&& unsafe { $name::filter_fetch($state, $name, entity, table_row) })* } } - }; } @@ -1240,8 +1239,9 @@ unsafe impl QueryFilter for Spawned { /// A marker trait to indicate that the filter works at an archetype level. /// -/// This is needed to implement [`ExactSizeIterator`] for -/// [`QueryIter`](crate::query::QueryIter) that contains archetype-level filters. +/// This is needed to: +/// - implement [`ExactSizeIterator`] for [`QueryIter`](crate::query::QueryIter) that contains archetype-level filters. +/// - ensure table filtering for [`QueryContiguousIter`](crate::query::QueryContiguousIter). /// /// The trait must only be implemented for filters where its corresponding [`QueryFilter::IS_ARCHETYPAL`] /// is [`prim@true`]. As such, only the [`With`] and [`Without`] filters can implement the trait. diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 0bd178ba03727..f39c6222a2ba0 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -4,7 +4,10 @@ use crate::{ bundle::Bundle, change_detection::Tick, entity::{ContainsEntity, Entities, Entity, EntityEquivalent, EntitySet, EntitySetIterator}, - query::{ArchetypeFilter, ArchetypeQueryData, DebugCheckedUnwrap, QueryState, StorageId}, + query::{ + ArchetypeFilter, ArchetypeQueryData, ContiguousQueryData, DebugCheckedUnwrap, QueryState, + StorageId, + }, storage::{Table, TableRow, Tables}, world::{ unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept, @@ -900,6 +903,17 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) } } + + /// Returns a contiguous iter or [`None`] if contiguous access is not supported + pub fn as_contiguous_iter(&mut self) -> Option> + where + D: ContiguousQueryData, + F: ArchetypeFilter, + { + self.cursor + .is_dense + .then_some(QueryContiguousIter { iter: self }) + } } impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> { @@ -992,6 +1006,45 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter> Clone for QueryIter<'w, 's, D } } +/// Iterator for contiguous chunks of memory +pub struct QueryContiguousIter<'a, 'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> { + iter: &'a mut QueryIter<'w, 's, D, F>, +} + +impl<'a, 'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> Iterator + for QueryContiguousIter<'a, 'w, 's, D, F> +{ + type Item = D::Contiguous<'w, 's>; + + #[inline(always)] + fn next(&mut self) -> Option { + // SAFETY: + // `tables` belongs to the same world that the cursor was initialized for. + // `query_state` is the state that was passed to `QueryIterationCursor::init` + unsafe { + self.iter + .cursor + .next_contiguous(self.iter.tables, self.iter.query_state) + } + } + + fn size_hint(&self) -> (usize, Option) { + self.iter.cursor.storage_id_iter.size_hint() + } +} + +// [`QueryIterationCursor::next_contiguous`] always returns None when exhausted +impl<'a, 'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> FusedIterator + for QueryContiguousIter<'a, 'w, 's, D, F> +{ +} + +// self.iter.cursor.storage_id_iter is a slice's iterator hence has an exact size +impl<'a, 'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> ExactSizeIterator + for QueryContiguousIter<'a, 'w, 's, D, F> +{ +} + /// An [`Iterator`] over sorted query results of a [`Query`](crate::system::Query). /// /// This struct is created by the [`QueryIter::sort`], [`QueryIter::sort_unstable`], @@ -2530,6 +2583,54 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { remaining_matched + self.current_len - self.current_row } + /// Returns the next contiguous chunk of memory or [`None`] if it is impossible or there is none + /// + /// # Safety + /// - `tables` must belong to the same world that the [`QueryIterationCursor`] was initialized for. + /// - `query_state` must be the same [`QueryState`] that was passed to `init` or `init_empty`. + /// - Query must be dense + #[inline(always)] + unsafe fn next_contiguous( + &mut self, + tables: &'w Tables, + query_state: &'s QueryState, + ) -> Option> + where + D: ContiguousQueryData, + F: ArchetypeFilter, + { + // SAFETY: Refer to [`Self::next`] + loop { + if self.current_row == self.current_len { + let table_id = self.storage_id_iter.next()?.table_id; + let table = tables.get(table_id).debug_checked_unwrap(); + if table.is_empty() { + continue; + } + D::set_table(&mut self.fetch, &query_state.fetch_state, table); + F::set_table(&mut self.filter, &query_state.filter_state, table); + self.table_entities = table.entities(); + self.current_len = table.entity_count(); + self.current_row = 0; + } + + let offset = self.current_row as usize; + self.current_row = self.current_len; + + // no filtering because `F` implements `ArchetypeFilter` which ensures that `QueryFilter::fetch` + // always returns true + + let item = D::fetch_contiguous( + &query_state.fetch_state, + &mut self.fetch, + self.table_entities, + offset, + ); + + return Some(item); + } + } + // NOTE: If you are changing query iteration code, remember to update the following places, where relevant: // QueryIter, QueryIterationCursor, QuerySortedIter, QueryManyIter, QuerySortedManyIter, QueryCombinationIter, // QueryState::par_fold_init_unchecked_manual, QueryState::par_many_fold_init_unchecked_manual, @@ -2546,6 +2647,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { query_state: &'s QueryState, ) -> Option> { if self.is_dense { + // NOTE: If you are changing this branch's code (the self.is_dense branch), + // don't forget to update [`Self::next_contiguous`] loop { // we are on the beginning of the query, or finished processing a table, so skip to the next if self.current_row == self.current_len { diff --git a/crates/bevy_ptr/src/lib.rs b/crates/bevy_ptr/src/lib.rs index e76fcc9f57250..7a1cf2b67a193 100644 --- a/crates/bevy_ptr/src/lib.rs +++ b/crates/bevy_ptr/src/lib.rs @@ -1104,6 +1104,67 @@ impl<'a, T> ThinSlicePtr<'a, T> { unsafe { &*self.ptr.add(index).as_ptr() } } + /// Returns a slice without performing bounds checks. + /// + /// # Safety + /// + /// `len` must be less or equal to the length of the slice. + pub unsafe fn as_slice_unchecked(&self, len: usize) -> &'a [T] { + #[cfg(debug_assertions)] + assert!(len <= self.len, "tried to create an out-of-bounds slice"); + + // SAFETY: The caller guarantees `len` is not greater than the length of the slice + unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), len) } + } + + /// Casts the slice to another type + pub fn cast(&self) -> ThinSlicePtr<'a, U> { + ThinSlicePtr { + ptr: self.ptr.cast::(), + // self.len is equal the amount of elements of T in the slice, which takes + // size_of:: * self.len bytes, thus the length of the same slice but for U is the amount + // of bytes divided by the size of U. + // + // when the size of U is 0, then the length of the slice may be infinite. + // + // when the size of T is 0 as well, then we can logically assume that the lengths of the both slices (of type T, + // and of type U) are equal. + #[cfg(debug_assertions)] + len: if size_of::() == 0 { + if size_of::() == 0 { + self.len + } else { + isize::MAX as usize + } + } else { + self.len * size_of::() / size_of::() + }, + _marker: PhantomData, + } + } + + /// Offsets the slice beginning by `count` elements + /// + /// # Safety + /// + /// - `count` must be less or equal to the length of the slice + // The result pointer must lie within the same allocation + pub unsafe fn add_unchecked(&self, count: usize) -> ThinSlicePtr<'a, T> { + #[cfg(debug_assertions)] + assert!( + count <= self.len, + "tried to offset the slice by more than the length" + ); + + Self { + // SAFETY: The caller guarantees that count is in-bounds. + ptr: unsafe { self.ptr.add(count) }, + #[cfg(debug_assertions)] + len: self.len - count, + _marker: PhantomData, + } + } + /// Indexes the slice without performing bounds checks. /// /// # Safety @@ -1116,6 +1177,22 @@ impl<'a, T> ThinSlicePtr<'a, T> { } } +impl<'a, T> ThinSlicePtr<'a, UnsafeCell> { + /// Returns a mutable reference of the slice + /// + /// # Safety + /// + /// - There must not be any aliases to the slice + /// - `len` must be less or equal to the length of the slice + pub unsafe fn as_mut_slice_unchecked(&self, len: usize) -> &'a mut [T] { + #[cfg(debug_assertions)] + assert!(len <= self.len, "tried to create an out-of-bounds slice"); + + // SAFETY: The caller ensures no aliases exist and `len` is in-bounds. + unsafe { core::slice::from_raw_parts_mut(UnsafeCell::raw_get(self.ptr.as_ptr()), len) } + } +} + impl<'a, T> Clone for ThinSlicePtr<'a, T> { fn clone(&self) -> Self { *self diff --git a/examples/README.md b/examples/README.md index 440c54d0c2bdb..31dbaff9b97c0 100644 --- a/examples/README.md +++ b/examples/README.md @@ -311,6 +311,7 @@ Example | Description --- | --- [Change Detection](../examples/ecs/change_detection.rs) | Change detection on components and resources [Component Hooks](../examples/ecs/component_hooks.rs) | Define component hooks to manage component lifecycle events +[Contiguous Query](../examples/ecs/contiguous_query.rs) | Demonstrates contiguous queries [Custom Query Parameters](../examples/ecs/custom_query_param.rs) | Groups commonly used compound queries and query filters into a single type [Custom Schedule](../examples/ecs/custom_schedule.rs) | Demonstrates how to add custom schedules [Dynamic ECS](../examples/ecs/dynamic.rs) | Dynamically create components, spawn entities with those components and query those components diff --git a/examples/ecs/contiguous_query.rs b/examples/ecs/contiguous_query.rs new file mode 100644 index 0000000000000..a435967dd36fb --- /dev/null +++ b/examples/ecs/contiguous_query.rs @@ -0,0 +1,55 @@ +//! Demonstrates how contiguous queries work + +use bevy::prelude::*; + +#[derive(Component)] +/// When the value reaches 0.0 the entity dies +pub struct Health(pub f32); + +#[derive(Component)] +/// Each tick an entity will have his health multiplied by the factor, which +/// for a big amount of entities can be accelerated using contiguous queries +pub struct HealthDecay(pub f32); + +fn apply_health_decay(mut query: Query<(&mut Health, &HealthDecay)>) { + // as_contiguous_iter() would return None if query couldn't be iterated contiguously + for ((health, _health_ticks), decay) in query.iter_mut().as_contiguous_iter().unwrap() { + // all slices returned by component queries are the same size + assert!(health.len() == decay.len()); + for i in 0..health.len() { + health[i].0 *= decay[i].0; + } + // we could have updated health's ticks but it is unnecessary hence we can make less work + // _health_ticks.mark_all_as_updated(); + } +} + +fn finish_off_first(mut commands: Commands, mut query: Query<(Entity, &mut Health)>) { + if let Some((entity, mut health)) = query.iter_mut().next() { + health.0 -= 1.0; + if health.0 <= 0.0 { + commands.entity(entity).despawn(); + println!("Finishing off {entity:?}"); + } + } +} + +fn main() { + App::new() + .add_plugins(DefaultPlugins) + .add_systems(Update, (apply_health_decay, finish_off_first).chain()) + .add_systems(Startup, setup) + .run(); +} + +fn setup(mut commands: Commands) { + let mut i = 0; + commands.spawn_batch(std::iter::from_fn(move || { + i += 1; + if i == 10_000 { + None + } else { + Some((Health(i as f32 * 5.0), HealthDecay(0.9))) + } + })); +} diff --git a/examples/ecs/custom_query_param.rs b/examples/ecs/custom_query_param.rs index 8960557ef4e85..28025915bd9a5 100644 --- a/examples/ecs/custom_query_param.rs +++ b/examples/ecs/custom_query_param.rs @@ -28,6 +28,7 @@ fn main() { print_components_iter_mut, print_components_iter, print_components_tuple, + print_components_contiguous_iter, ) .chain(), ) @@ -111,7 +112,7 @@ struct NestedQuery { } #[derive(QueryData)] -#[query_data(derive(Debug))] +#[query_data(derive(Debug), contiguous(mutable))] struct GenericQuery { generic: (&'static T, &'static P), } @@ -193,3 +194,35 @@ fn print_components_tuple( println!("Generic: {generic_c:?} {generic_d:?}"); } } + +/// If you are going to contiguously iterate the data in a query, you must mark it with the `contiguous` attribute, +/// which accepts one of 3 targets (`all`, `immutable` and `mutable`) +/// +/// - `all` will make read only query as well as mutable query both be able to be iterated contiguosly +/// - `mutable` will only make the original query (i.e., in that case [`CustomContiguousQuery`]) be able to be iterated contiguously +/// - `immutable` will only make the read only query (which is only useful when you mark the original query as `mutable`) +/// be able to be iterated contiguously +#[derive(QueryData)] +#[query_data(derive(Debug), contiguous(all))] +struct CustomContiguousQuery { + entity: Entity, + a: &'static ComponentA, + b: Option<&'static ComponentB>, + generic: GenericQuery, +} + +fn print_components_contiguous_iter(query: Query>) { + println!("Print components (contiguous_iter):"); + for e in query.iter().as_contiguous_iter().unwrap() { + let e: CustomContiguousQueryContiguousItem<'_, '_, _, _> = e; + for i in 0..e.entity.len() { + println!("Entity: {:?}", e.entity[i]); + println!("A: {:?}", e.a[i]); + println!("B: {:?}", e.b.map(|b| &b[i])); + println!( + "Generic: {:?} {:?}", + e.generic.generic.0[i], e.generic.generic.1[i] + ); + } + } +}