From e94c9fa6c2544738e562b3f818e3e913c5c35841 Mon Sep 17 00:00:00 2001 From: Charles Bournhonesque Date: Thu, 23 Oct 2025 02:41:23 -0400 Subject: [PATCH 1/7] simple test passes --- crates/bevy_ecs/src/query/mod.rs | 1 + crates/bevy_ecs/src/query/plan.rs | 605 ++++++++++++++++++++++++++++++ 2 files changed, 606 insertions(+) create mode 100644 crates/bevy_ecs/src/query/plan.rs diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index fb8899fd5de87..d724976045ba5 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -9,6 +9,7 @@ mod iter; mod par_iter; mod state; mod world_query; +mod plan; pub use access::*; pub use bevy_ecs_macros::{QueryData, QueryFilter}; diff --git a/crates/bevy_ecs/src/query/plan.rs b/crates/bevy_ecs/src/query/plan.rs new file mode 100644 index 0000000000000..07d4a686fe9f0 --- /dev/null +++ b/crates/bevy_ecs/src/query/plan.rs @@ -0,0 +1,605 @@ +use alloc::borrow::Cow; +use crate::{ + component::ComponentId, + entity::Entity, + query::FilteredAccess, + world::unsafe_world_cell::UnsafeWorldCell, +}; +use alloc::vec::Vec; +use alloc::vec; +use fixedbitset::FixedBitSet; +use bevy_ecs::archetype::Archetype; +use bevy_ecs::component::{Component}; +use bevy_ecs::prelude::{QueryBuilder, World}; +use bevy_ecs::query::{QueryData, QueryFilter}; +use bevy_ecs::relationship::Relationship; +use bevy_ecs::storage::{SparseSetIndex, TableId}; +use bevy_ecs::world::unsafe_world_cell::UnsafeEntityCell; +use crate::relationship::RelationshipAccessor; + +/// Represents a single source in a multi-source query. +/// Each term has its own ComponentAccess requirements. +#[derive(Debug, Clone)] +pub enum QueryOperation { + Relationship(QueryRelationship), + Source(QuerySource), +} + +#[derive(Debug, Clone)] +pub struct QuerySource { + access: FilteredAccess, + variable_idx: u8, +} + + +impl QuerySource { + pub fn matches(&self, archetype: &Archetype) -> bool { + self.matches_component_set(&|id| archetype.contains(id)) + } + + /// Returns `true` if this query matches a set of components. Otherwise, returns `false`. + pub fn matches_component_set(&self, set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { + self.access.filter_sets.iter().any(|set| { + set.with + .ones() + .all(|index| set_contains_id(ComponentId::get_sparse_set_index(index))) + && set + .without + .ones() + .all(|index| !set_contains_id(ComponentId::get_sparse_set_index(index))) + }) + } +} + +/// Describes how two query terms are connected via a relationship. +#[derive(Debug, Clone)] +pub struct QueryRelationship { + /// The source term index within the QueryPlan + pub source_idx: u8, + /// The target term index within the QueryPlan + pub target_idx: u8, + /// The relationship component that links source to target. + pub relationship_component: ComponentId, + /// Accessor to dynamically access the 'target' of the relationship + pub relationship_accessor: RelationshipAccessor, +} + +impl QueryRelationship { + /// Get the 'target' of the source entity + pub unsafe fn get(&self, source: Entity, world: UnsafeWorldCell<'_>) -> Option { + let entity_field_offset = match self.relationship_accessor { + RelationshipAccessor::Relationship { entity_field_offset, .. } => {entity_field_offset} + RelationshipAccessor::RelationshipTarget { .. } => { + unreachable!() + } + }; + let relationship_ptr = world.get_entity(source).ok()?.get_by_id(self.relationship_component)?; + Some(unsafe { *relationship_ptr.byte_add(entity_field_offset).deref() }) + } +} + +pub struct QueryVariable { + index: u8, +} + +impl> From for QueryVariable { + fn from(value: T) -> Self { + Self { + index: value.into() + } + } +} + +pub struct QueryPlanBuilder<'w> { + world: &'w mut World, + plan: QueryPlan, +} + +impl<'w> QueryPlanBuilder<'w> { + pub fn new(world: &'w mut World) -> Self { + Self { + world, + plan: QueryPlan::default(), + } + } + + pub fn add_source_from_builder(&mut self, f: impl Fn(QueryBuilder) -> QueryBuilder) -> &mut Self { + let query_builder = QueryBuilder::new(&mut self.world); + let builder = f(query_builder); + self.plan.add_source(builder.access().clone()); + self + } + + pub fn add_source(&mut self) -> &mut Self { + let fetch_state = D::init_state(&mut self.world); + let filter_state = F::init_state(&mut self.world); + + let mut access = FilteredAccess::default(); + D::update_component_access(&fetch_state, &mut access); + + // Use a temporary empty FilteredAccess for filters. This prevents them from conflicting with the + // main Query's `fetch_state` access. Filters are allowed to conflict with the main query fetch + // because they are evaluated *before* a specific reference is constructed. + let mut filter_access = FilteredAccess::default(); + F::update_component_access(&filter_state, &mut filter_access); + + // Merge the temporary filter access with the main access. This ensures that filter access is + // properly considered in a global "cross-query" context (both within systems and across systems). + access.extend(&filter_access); + + self.plan.add_source(access); + self + } + + + /// Add a relationship between two terms using a typed Relationship component. + pub fn add_relationship( + &mut self, + source_term: u8, + target_term: u8, + ) { + let component_id = self.world.register_component::(); + let accessor = self.world + .components() + .get_info(component_id) + .unwrap() + .relationship_accessor() + .unwrap().clone(); + + self.plan.add_relationship( + source_term, + target_term, + component_id, + accessor, + ); + } + + pub fn compile(self) -> QueryPlan { + self.plan.compile() + } +} + +#[derive(thiserror::Error, Debug)] +pub enum QueryPlanError { + /// The source does not exist + #[error("The source with index {0} does not exist")] + QuerySourceNotFound(u8), +} + +/// A dynamic query plan that describes how to match multiple entities +/// connected through relationships. +#[derive(Debug, Default, Clone)] +pub struct QueryPlan { + /// All operations in this query. + pub operations: Vec, + pub num_variables: u8, +} + + +impl QueryPlan { + pub fn query_iter<'w, 's>(&'s self, world: UnsafeWorldCell<'w>) -> Iter<'w, 's> { + Iter { + world, + query_state: Cow::Borrowed(self), + iter_state: IterState::new(self), + } + } +} + + +#[derive(Debug, Copy, Clone)] +pub enum VariableState { + // we are searching through all the entities of a table + Search { + table: Option, + // offset in the current table + offset: u32, + // length of the current table + current_len: u32, + // index of the table inside the StorageState + storage_idx: u32, + }, + // An entity has been found by following a relationship + FixedByRelationship(Entity) +} + +impl Default for VariableState { + fn default() -> Self { + Self::Search { + table: None, + offset: 0, + current_len: 0, + storage_idx: 0, + } + } +} + + +#[derive(Default, Clone)] +pub struct StorageState<'w> { + tables: Cow<'w, [TableId]>, +} + +pub struct IterState<'w> { + // are we backtracking through the plan? + pub backtrack: bool, + // index of the operation in the plan we are currently executing + pub curr_op: u8, + /// Index of the current entity for each variable + pub variable_state: Vec, + /// List of matching tables/archetypes to iterate through for each variable + pub storages_state: Vec>, + /// Whether we have already found an Entity for the source when we are about to process the next operation + written: Vec, +} + +impl<'w> IterState<'w> { + fn new(plan: &QueryPlan) -> Self { + let num_operations = plan.operations.len(); + let num_variables = plan.num_variables as usize; + let mut variable_state = vec![VariableState::default(); num_variables]; + let mut storages_state = vec![StorageState::default(); num_variables]; + let mut written = vec![FixedBitSet::with_capacity(num_variables); num_operations]; + Self { + backtrack: false, + curr_op: 0, + variable_state, + storages_state, + written, + } + } + +} + + +/// Iterator that iterates through a dynamic query plan +pub struct Iter<'w, 's> { + world: UnsafeWorldCell<'w>, + query_state: Cow<'s, QueryPlan>, + iter_state: IterState<'s>, +} + + +impl<'w, 's> Iter<'w, 's> { + + /// Returns true if the `variable` was currently assigned an Entity + fn written(&self, variable: u8) -> bool { + self.iter_state.written[self.iter_state.curr_op as usize].contains(variable as usize) + } + + /// Get the entity currently written to the variable indexed by `variable`, if it has been written + fn written_entity(&self, variable: u8) -> Option { + let variable_idx = variable as usize; + Some(match self.iter_state.variable_state[variable_idx] { + VariableState::Search { table, offset, .. } => { + let table_id = table?; + // Safety: + let table = unsafe { self.world.storages() }.tables.get(table_id)?; + unsafe{ *table.entities().get_unchecked(offset as usize) } + } + VariableState::FixedByRelationship(entity) => {entity} + }) + } + + // we found entity1, QueryOp:Rel -> set entity 2 to one of the entities via the RelationShip + fn dispatch(&mut self) -> bool { + let op = self.current_op(); + match op { + QueryOperation::Relationship(_) => self.op_relationship(), + QueryOperation::Source(_) => self.op_query(), + } + } + + /// Try to find an entity via the relationship + fn op_relationship(&mut self) -> bool { + if self.iter_state.backtrack { + return false; + } + let QueryOperation::Relationship(rel) = self.current_op() else { + unreachable!() + }; + // TODO: do we always find the source first? what if the target was written but not the source? + debug_assert!(self.written(rel.source_idx), "we only support queries where the source has been found before we are querying the relationship"); + // we already found the target term! + if self.written(rel.target_idx) { + let target_state = self.iter_state.variable_state[rel.target_idx as usize]; + match target_state { + VariableState::Search { .. } => { + todo!("Check if target_entity is equal to the relationship.get() value") + } + VariableState::FixedByRelationship(_) => { + true + } + } + } else { + // need to find the target term! We do this by simply following the relationship.get() + let source_entity = self.written_entity(rel.source_idx).unwrap(); + // it is guaranteed that the target exists, since the Relationship component is present + // on the source entity + // SAFETY: TODO + let target_entity = unsafe { rel.get(source_entity, self.world).unwrap_unchecked() }; + self.set_variable_state(rel.target_idx, VariableState::FixedByRelationship(target_entity)); + true + } + } + + fn op_query(&mut self) -> bool { + let QueryOperation::Source(source) = self.current_op() else { + unreachable!() + }; + let variable_idx = source.variable_idx as usize; + + // we already have a potential candidate: check if it matches the query + if let VariableState::FixedByRelationship(entity) = self.iter_state.variable_state[variable_idx] { + // in this case keep backtracking + if self.iter_state.backtrack { + return false + } + let archetype = unsafe { self.world.get_entity(entity).unwrap_unchecked().archetype() }; + return source.matches(archetype) + } + + // we haven't found the entity yet. + // - TODO: either we already had some constraints on the entity (i.e. a potential list of archetypes + // that this entity can be part of), in which case we can further filter down these archetypes + // -> this can only happen if we allow multiple queries for a single variable + // + // we need to use the component index to find a list of potential archetypes + // that could match the query + + // TODO: what about caching? + + // the first time we evaluate the query, we set the list of potential tables + if !self.iter_state.backtrack { + // if there are required components, we can optimize by only iterating through archetypes + // that contain at least one of the required components + let potential_archetypes = source + .access + .required + .ones() + .filter_map(|idx| { + let component_id = ComponentId::get_sparse_set_index(idx); + self.world + .archetypes() + .component_index() + .get(&component_id) + .map(|index| index.keys()) + }) + // select the component with the fewest archetypes + .min_by_key(ExactSizeIterator::len); + let mut matching_tables = Vec::new(); + if let Some(archetypes) = potential_archetypes { + for archetype_id in archetypes { + // SAFETY: get_potential_archetypes only returns archetype ids that are valid for the world + let archetype = &self.world.archetypes()[*archetype_id]; + if source.matches(archetype) { + matching_tables.push(archetype.table_id()) + } + } + } + let Some(&table) = matching_tables.first() else { + return false; + }; + self.iter_state.storages_state[variable_idx] = StorageState { + tables: Cow::Owned(matching_tables) + }; + self.set_variable_state(variable_idx as u8, VariableState::Search { + table: Some(table), + offset: 0, + storage_idx: 0, + current_len: 0, + }); + return true; + } + + // else we backtracked! we need to advance in the current table, or go to the next table + let storage_state = &self.iter_state.storages_state[variable_idx]; + let target_state = &mut self.iter_state.variable_state[variable_idx]; + let VariableState::Search { table: Some(table_id), offset, current_len, storage_idx, } = target_state else { + unreachable!(); + }; + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + let iteration_start = *current_len == 0; + // loop to skip empty tables + loop { + // either beginning of the iteration, or finished processing a table, so skip to the next + if offset == current_len { + // go to next table + if !iteration_start { + *storage_idx += 1; + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + *table_id = storage_state.tables[*storage_idx as usize]; + *offset = 0; + } + let table = unsafe { self.world.world().storages().tables.get(*table_id).unwrap_unchecked() }; + *current_len = table.entity_count(); + let table = unsafe { self.world.world().storages().tables.get(*table_id).unwrap_unchecked() }; + if table.is_empty() { + // skip table + continue; + } + } + + // TODO: store `table.entities()` somewhere so we don't have to fetch it again every time + // let table = unsafe { self.world.world().storages().tables.get(*table_id).unwrap_unchecked() }; + // let entity = unsafe{ table.entities().get_unchecked(*offset as usize) }; + *offset += 1; + + // this entity is our current candidate + return true; + } + } + + /// Assign the variable state (and also set the entity as written for the next op) + fn set_variable_state(&mut self, variable: u8, state: VariableState) { + // the operation writes the entity as written for the **next** operation + let curr_op = self.iter_state.curr_op as usize; + let variable_idx = variable as usize; + if curr_op + 1 < self.iter_state.written.len() { + self.iter_state.written[curr_op + 1].grow_and_insert(variable_idx); + } + self.iter_state.variable_state[variable_idx] = state; + } + + fn current_op(&self) -> &QueryOperation { + &self.query_state.operations[self.iter_state.curr_op as usize] + } + + /// Check if we can still find a matching item + fn next_match(&mut self) -> bool { + let num_operations = self.query_state.operations.len() as u8; + // Past the end of operations so must have matched, return to previous op + if self.iter_state.curr_op == num_operations { + self.iter_state.backtrack = true; + self.iter_state.curr_op -= 1; + } + + while self.iter_state.curr_op < num_operations { + let op_index = self.iter_state.curr_op as usize; + let matched = self.dispatch(); + + if !matched { + // Operation did not match, return to previous op + self.iter_state.backtrack = true; + // Returned from all operations, no more matches + if self.iter_state.curr_op == 0 { + return false; + } + self.iter_state.curr_op -= 1; + } else { + // Operation did match, move to next op + self.iter_state.backtrack = false; + self.iter_state.curr_op += 1; + + if self.iter_state.curr_op < num_operations { + // Setup written state for next operation. The ops themselves have already updated the written state for the next + // op, but we also need to propagate the existing written from the current op + + // (we have a written bitset for each operation so that on backtracking we can retrieve the previous `written` state) + let (written_next_op, written_op) = self.iter_state.written.get_mut(op_index..op_index + 2).unwrap().split_at_mut(1); + // written[op_index + 1].union_with(written[op_index]) + written_next_op[0].union_with(&written_op[0]); + } + } + } + true + } +} + +pub struct DynamicItem<'w> { + // TODO: return FilteredEntityRef/FilteredEntityMut by borrowing from the State, which contains the access. + entities: Vec>, +} + +impl<'w, 's> Iterator for Iter<'w, 's> { + type Item = DynamicItem<'w>; + + fn next(&mut self) -> Option { + if !self.next_match() { + return None; + } + // retrieve the Item + let num_variables = self.query_state.num_variables as usize; + let mut matching_entities = Vec::with_capacity(num_variables); + for idx in 0..num_variables { + let entity = self.written_entity(idx as u8).unwrap(); + let entity_cell = self.world.get_entity(entity).unwrap(); + matching_entities.push(entity_cell); + } + Some(DynamicItem { + entities: matching_entities + }) + } +} + + +impl QueryPlan { + /// Create a new empty query plan. + pub fn new() -> Self { + Self { + operations: Vec::new(), + num_variables: 0, + } + } + + /// Add a term to the query plan. + pub(crate) fn add_source(&mut self, access: FilteredAccess) -> usize { + let op_index = self.operations.len(); + self.operations.push(QueryOperation::Source(QuerySource { + access, + variable_idx: self.num_variables, + })); + self.num_variables += 1; + op_index + } + + /// Add a relationship between two terms. + pub fn add_relationship( + &mut self, + source: impl Into, + target: impl Into, + relationship_component: ComponentId, + accessor: RelationshipAccessor, + ) -> &mut Self { + self.operations.push(QueryOperation::Relationship(QueryRelationship { + source_idx: source.into().index, + target_idx: target.into().index, + relationship_component, + relationship_accessor: accessor, + })); + self + } + + /// Do some optimizations and return the compiled plan + pub fn compile(self) -> QueryPlan { + // TODO: add the relationship/relationshipTarget accesses to each source term + self + } +} + + +#[cfg(test)] +mod tests { + use bevy_ecs::prelude::*; + use super::*; + use crate::{ + component::Component, + hierarchy::ChildOf, + prelude::World, + }; + + #[derive(Component)] + struct Marker; + + #[test] + fn test_query_plan_basic() { + let mut world = World::new(); + + // Create parent-child relationship + let parent = world.spawn_empty().id(); + let child = world.spawn((Marker, ChildOf(parent))).id(); + world.flush(); // Ensure Children component is added to parent + + let marker_id = world.register_component::(); + let child_of_id = world.register_component::(); + + // Build a simple plan using the builder API + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::<&Marker, ()>(); + builder.add_relationship::(0, 1); + builder.add_source::(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].entities[0].entity(), child); + assert_eq!(results[0].entities[1].entity(), parent); + } +} + From d2aeb5d547ad6b19fe6ff0f9fb290a5506d509a1 Mon Sep 17 00:00:00 2001 From: Charles Bournhonesque Date: Thu, 23 Oct 2025 10:15:00 -0400 Subject: [PATCH 2/7] add more tests --- crates/bevy_ecs/src/query/plan.rs | 123 ++++++++++++++++++++---------- 1 file changed, 82 insertions(+), 41 deletions(-) diff --git a/crates/bevy_ecs/src/query/plan.rs b/crates/bevy_ecs/src/query/plan.rs index 07d4a686fe9f0..5f99423bf1c51 100644 --- a/crates/bevy_ecs/src/query/plan.rs +++ b/crates/bevy_ecs/src/query/plan.rs @@ -351,45 +351,49 @@ impl<'w, 's> Iter<'w, 's> { // the first time we evaluate the query, we set the list of potential tables if !self.iter_state.backtrack { - // if there are required components, we can optimize by only iterating through archetypes - // that contain at least one of the required components - let potential_archetypes = source - .access - .required - .ones() - .filter_map(|idx| { - let component_id = ComponentId::get_sparse_set_index(idx); - self.world - .archetypes() - .component_index() - .get(&component_id) - .map(|index| index.keys()) - }) - // select the component with the fewest archetypes - .min_by_key(ExactSizeIterator::len); - let mut matching_tables = Vec::new(); - if let Some(archetypes) = potential_archetypes { - for archetype_id in archetypes { - // SAFETY: get_potential_archetypes only returns archetype ids that are valid for the world - let archetype = &self.world.archetypes()[*archetype_id]; - if source.matches(archetype) { - matching_tables.push(archetype.table_id()) + let tables = &self.iter_state.storages_state[variable_idx].tables; + // only set the list of potential tables if we didn't do so before + if tables.len() == 0 { + // if there are required components, we can optimize by only iterating through archetypes + // that contain at least one of the required components + let potential_archetypes = source + .access + .required + .ones() + .filter_map(|idx| { + let component_id = ComponentId::get_sparse_set_index(idx); + self.world + .archetypes() + .component_index() + .get(&component_id) + .map(|index| index.keys()) + }) + // select the component with the fewest archetypes + .min_by_key(ExactSizeIterator::len); + let mut matching_tables = Vec::new(); + if let Some(archetypes) = potential_archetypes { + for archetype_id in archetypes { + // SAFETY: get_potential_archetypes only returns archetype ids that are valid for the world + let archetype = &self.world.archetypes()[*archetype_id]; + if source.matches(archetype) { + matching_tables.push(archetype.table_id()) + } } } + let Some(&table) = matching_tables.first() else { + return false; + }; + self.iter_state.storages_state[variable_idx] = StorageState { + tables: Cow::Owned(matching_tables) + }; + self.set_variable_state(variable_idx as u8, VariableState::Search { + table: Some(table), + offset: 0, + storage_idx: 0, + current_len: 0, + }); + return true; } - let Some(&table) = matching_tables.first() else { - return false; - }; - self.iter_state.storages_state[variable_idx] = StorageState { - tables: Cow::Owned(matching_tables) - }; - self.set_variable_state(variable_idx as u8, VariableState::Search { - table: Some(table), - offset: 0, - storage_idx: 0, - current_len: 0, - }); - return true; } // else we backtracked! we need to advance in the current table, or go to the next table @@ -579,13 +583,10 @@ mod tests { fn test_query_plan_basic() { let mut world = World::new(); - // Create parent-child relationship + // Correct pair let parent = world.spawn_empty().id(); let child = world.spawn((Marker, ChildOf(parent))).id(); - world.flush(); // Ensure Children component is added to parent - - let marker_id = world.register_component::(); - let child_of_id = world.register_component::(); + world.flush(); // Build a simple plan using the builder API let mut builder = QueryPlanBuilder::new(&mut world); @@ -601,5 +602,45 @@ mod tests { assert_eq!(results[0].entities[0].entity(), child); assert_eq!(results[0].entities[1].entity(), parent); } + + /// Checks that filters in the source or target of the relationship are respected + #[test] + fn test_query_plan_single_relationship() { + let mut world = World::new(); + + // Parent does not have the marker + let parent1 = world.spawn_empty().id(); + let child1 = world.spawn((Marker, ChildOf(parent1))).id(); + world.flush(); + + // Child does not have the marker + let parent2 = world.spawn(Marker).id(); + let child2 = world.spawn(ChildOf(parent2)).id(); + + // Both have markers but there is no relationship + let parent3 = world.spawn(Marker).id(); + let child3 = world.spawn(Marker).id(); + + // Two correct pairs, (Child, Parent) and (Parent, Grandparent) + let grandparent4 = world.spawn(Marker).id(); + let parent4 = world.spawn((Marker, ChildOf(grandparent4))).id(); + let child4 = world.spawn((Marker, ChildOf(parent4))).id(); + + // Both sources must have the Marker + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::<&Marker, ()>(); + builder.add_relationship::(0, 1); + builder.add_source::>(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].entities[0].entity(), child4); + assert_eq!(results[0].entities[1].entity(), parent4); + assert_eq!(results[1].entities[0].entity(), parent4); + assert_eq!(results[2].entities[1].entity(), grandparent4); + } } From 89c201e1acbd67330b83e95b5d9b2a88de1bcea0 Mon Sep 17 00:00:00 2001 From: Charles Bournhonesque Date: Sat, 25 Oct 2025 21:58:12 -0400 Subject: [PATCH 3/7] handle RelationshipTarget and fix backtracking --- crates/bevy_ecs/src/query/plan.rs | 379 +++++++++++++++++++----------- 1 file changed, 247 insertions(+), 132 deletions(-) diff --git a/crates/bevy_ecs/src/query/plan.rs b/crates/bevy_ecs/src/query/plan.rs index 5f99423bf1c51..4e76e9688322e 100644 --- a/crates/bevy_ecs/src/query/plan.rs +++ b/crates/bevy_ecs/src/query/plan.rs @@ -7,9 +7,7 @@ use crate::{ }; use alloc::vec::Vec; use alloc::vec; -use fixedbitset::FixedBitSet; use bevy_ecs::archetype::Archetype; -use bevy_ecs::component::{Component}; use bevy_ecs::prelude::{QueryBuilder, World}; use bevy_ecs::query::{QueryData, QueryFilter}; use bevy_ecs::relationship::Relationship; @@ -25,6 +23,11 @@ pub enum QueryOperation { Source(QuerySource), } +// - we want to track when we change the VariableState, so that we can backtrack and restore to the previous state +// - we cannot just cancel a VariableState::Fixed, because the variable might have been Fixed in a previous operation. +// - we want to handle queries like Q1, R21 (i.e. find all entities $2 where R($1, $2), which means the possibilities for $2 are within a list of entities: RelationshipTarget) + + #[derive(Debug, Clone)] pub struct QuerySource { access: FilteredAccess, @@ -58,15 +61,18 @@ pub struct QueryRelationship { pub source_idx: u8, /// The target term index within the QueryPlan pub target_idx: u8, - /// The relationship component that links source to target. + /// The relationship component that links source to target. (e.g. ChildOf) pub relationship_component: ComponentId, + /// The relationship target component that links target to source. (e.g. Children) + pub relationship_target_component: ComponentId, /// Accessor to dynamically access the 'target' of the relationship pub relationship_accessor: RelationshipAccessor, + pub relationship_target_accessor: RelationshipAccessor, } impl QueryRelationship { /// Get the 'target' of the source entity - pub unsafe fn get(&self, source: Entity, world: UnsafeWorldCell<'_>) -> Option { + pub unsafe fn get_target(&self, source: Entity, world: UnsafeWorldCell<'_>) -> Option { let entity_field_offset = match self.relationship_accessor { RelationshipAccessor::Relationship { entity_field_offset, .. } => {entity_field_offset} RelationshipAccessor::RelationshipTarget { .. } => { @@ -76,6 +82,18 @@ impl QueryRelationship { let relationship_ptr = world.get_entity(source).ok()?.get_by_id(self.relationship_component)?; Some(unsafe { *relationship_ptr.byte_add(entity_field_offset).deref() }) } + + pub unsafe fn get_sources(&self, target: Entity, world: UnsafeWorldCell<'_>) -> Option> { + let iter = match self.relationship_target_accessor { + RelationshipAccessor::Relationship { .. } => { + unreachable!() + } + RelationshipAccessor::RelationshipTarget { iter, .. } => {iter} + }; + let relationship_ptr = world.get_entity(target).ok()?.get_by_id(self.relationship_target_component)?; + let sources: Vec<_> = unsafe { iter(relationship_ptr).collect() }; + Some(sources) + } } pub struct QueryVariable { @@ -139,18 +157,27 @@ impl<'w> QueryPlanBuilder<'w> { target_term: u8, ) { let component_id = self.world.register_component::(); + let target_component_id = self.world.register_component::<::RelationshipTarget>(); let accessor = self.world .components() .get_info(component_id) .unwrap() .relationship_accessor() .unwrap().clone(); + let target_accessor = self.world + .components() + .get_info(target_component_id) + .unwrap() + .relationship_accessor() + .unwrap().clone(); self.plan.add_relationship( source_term, target_term, component_id, + target_component_id, accessor, + target_accessor, ); } @@ -187,11 +214,12 @@ impl QueryPlan { } -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Clone)] pub enum VariableState { + // Default state: we haven't applied any constraints yet + Unknown, // we are searching through all the entities of a table Search { - table: Option, // offset in the current table offset: u32, // length of the current table @@ -200,17 +228,17 @@ pub enum VariableState { storage_idx: u32, }, // An entity has been found by following a relationship - FixedByRelationship(Entity) + FixedByRelationship(Entity), + // The entity is among the RelationshipTargets + FixedByRelationshipTarget { + sources: Vec, + index: usize, + } } impl Default for VariableState { fn default() -> Self { - Self::Search { - table: None, - offset: 0, - current_len: 0, - storage_idx: 0, - } + Self::Unknown } } @@ -227,10 +255,22 @@ pub struct IterState<'w> { pub curr_op: u8, /// Index of the current entity for each variable pub variable_state: Vec, + /// For each operation, store the previous variable state to restore on backtracking + /// in case this operation modified it + variable_state_diff: Vec>, /// List of matching tables/archetypes to iterate through for each variable pub storages_state: Vec>, - /// Whether we have already found an Entity for the source when we are about to process the next operation - written: Vec, +} + +pub struct PreviousVariableState { + variable_idx: u8, + previous_state: PreviousVariableDiff, +} + +pub enum PreviousVariableDiff { + VariableState(VariableState), + // the previous variable state is obtained by decrementing the index in the list of candidates + Decrement, } impl<'w> IterState<'w> { @@ -238,14 +278,15 @@ impl<'w> IterState<'w> { let num_operations = plan.operations.len(); let num_variables = plan.num_variables as usize; let mut variable_state = vec![VariableState::default(); num_variables]; + let mut variable_state_diff = Vec::new(); + variable_state_diff.resize_with(num_operations, Default::default); let mut storages_state = vec![StorageState::default(); num_variables]; - let mut written = vec![FixedBitSet::with_capacity(num_variables); num_operations]; Self { backtrack: false, curr_op: 0, variable_state, + variable_state_diff, storages_state, - written, } } @@ -264,21 +305,27 @@ impl<'w, 's> Iter<'w, 's> { /// Returns true if the `variable` was currently assigned an Entity fn written(&self, variable: u8) -> bool { - self.iter_state.written[self.iter_state.curr_op as usize].contains(variable as usize) + !matches!(self.iter_state.variable_state[variable as usize], VariableState::Unknown) } /// Get the entity currently written to the variable indexed by `variable`, if it has been written fn written_entity(&self, variable: u8) -> Option { let variable_idx = variable as usize; - Some(match self.iter_state.variable_state[variable_idx] { - VariableState::Search { table, offset, .. } => { - let table_id = table?; + match &self.iter_state.variable_state[variable_idx] { + VariableState::Unknown => { + None + } + VariableState::Search { storage_idx, offset, .. } => { // Safety: - let table = unsafe { self.world.storages() }.tables.get(table_id)?; - unsafe{ *table.entities().get_unchecked(offset as usize) } + let table_id = &self.iter_state.storages_state[variable_idx].tables[*storage_idx as usize]; + let table = unsafe { self.world.storages() }.tables.get(*table_id)?; + Some(unsafe{ *table.entities().get_unchecked(*offset as usize) }) } - VariableState::FixedByRelationship(entity) => {entity} - }) + VariableState::FixedByRelationship(entity) => {Some(*entity)} + VariableState::FixedByRelationshipTarget { sources, index } => { + Some(sources[*index]) + } + } } // we found entity1, QueryOp:Rel -> set entity 2 to one of the entities via the RelationShip @@ -292,68 +339,86 @@ impl<'w, 's> Iter<'w, 's> { /// Try to find an entity via the relationship fn op_relationship(&mut self) -> bool { - if self.iter_state.backtrack { - return false; - } let QueryOperation::Relationship(rel) = self.current_op() else { unreachable!() }; - // TODO: do we always find the source first? what if the target was written but not the source? - debug_assert!(self.written(rel.source_idx), "we only support queries where the source has been found before we are querying the relationship"); - // we already found the target term! - if self.written(rel.target_idx) { - let target_state = self.iter_state.variable_state[rel.target_idx as usize]; - match target_state { - VariableState::Search { .. } => { - todo!("Check if target_entity is equal to the relationship.get() value") + let source_idx = rel.source_idx as usize; + let target_idx = rel.target_idx as usize; + if self.iter_state.backtrack { + // possibly check the next source in the list + if let VariableState::FixedByRelationshipTarget { sources, index } = &mut self.iter_state.variable_state[source_idx] { + // TODO: this needs to be done via set_variable_state to track the changes for backtracking + *index += 1; + if *index >= sources.len() { + return false; + } + self.set_variable_state_increment(source_idx as u8); + return true; + } + // possibly check the next source in the list + if let VariableState::FixedByRelationshipTarget { sources, index } = &mut self.iter_state.variable_state[target_idx] { + // TODO: this needs to be done via set_variable_state to track the changes for backtracking + *index += 1; + if *index >= sources.len() { + return false; + } + self.set_variable_state_increment(target_idx as u8); + return true; + } + return false; + } + + match (self.written(rel.source_idx), self.written(rel.target_idx)) { + (false, false) => { + unreachable!("we only support queries where the source has been found before we are querying the relationship"); + } + (true, false) => { + // we found the source, need to find the target + let source_entity = self.written_entity(rel.source_idx).unwrap(); + match unsafe { rel.get_target(source_entity, self.world) } { + None => false, + Some(target_entity) => { + self.set_variable_state(rel.target_idx, VariableState::FixedByRelationship(target_entity)); + true + } } - VariableState::FixedByRelationship(_) => { - true + } + (false, true) => { + // we found the target, need to find the source + let target_entity = self.written_entity(rel.target_idx).unwrap(); + match unsafe { rel.get_sources(target_entity, self.world) } { + None => false, + Some(sources) => { + self.set_variable_state(rel.source_idx, VariableState::FixedByRelationshipTarget { + sources, + index: 0, + }); + true + } } } - } else { - // need to find the target term! We do this by simply following the relationship.get() - let source_entity = self.written_entity(rel.source_idx).unwrap(); - // it is guaranteed that the target exists, since the Relationship component is present - // on the source entity - // SAFETY: TODO - let target_entity = unsafe { rel.get(source_entity, self.world).unwrap_unchecked() }; - self.set_variable_state(rel.target_idx, VariableState::FixedByRelationship(target_entity)); - true + (true, true) => { + // we found both, need to check if they are linked by the relationship + let source_entity = self.written_entity(rel.source_idx).unwrap(); + let target_entity = self.written_entity(rel.target_idx).unwrap(); + let expected_target_entity = unsafe { rel.get_target(source_entity, self.world).unwrap_unchecked() }; + target_entity == expected_target_entity + } } + } fn op_query(&mut self) -> bool { - let QueryOperation::Source(source) = self.current_op() else { + let QueryOperation::Source(source) = &self.query_state.operations[self.iter_state.curr_op as usize] else { unreachable!() }; let variable_idx = source.variable_idx as usize; - - // we already have a potential candidate: check if it matches the query - if let VariableState::FixedByRelationship(entity) = self.iter_state.variable_state[variable_idx] { - // in this case keep backtracking - if self.iter_state.backtrack { - return false - } - let archetype = unsafe { self.world.get_entity(entity).unwrap_unchecked().archetype() }; - return source.matches(archetype) - } - - // we haven't found the entity yet. - // - TODO: either we already had some constraints on the entity (i.e. a potential list of archetypes - // that this entity can be part of), in which case we can further filter down these archetypes - // -> this can only happen if we allow multiple queries for a single variable - // - // we need to use the component index to find a list of potential archetypes - // that could match the query - - // TODO: what about caching? - - // the first time we evaluate the query, we set the list of potential tables - if !self.iter_state.backtrack { - let tables = &self.iter_state.storages_state[variable_idx].tables; - // only set the list of potential tables if we didn't do so before - if tables.len() == 0 { + let storage_state = &self.iter_state.storages_state[variable_idx]; + match &mut self.iter_state.variable_state[variable_idx] { + VariableState::Unknown => { + // the first time we evaluate the query, we set the list of potential tables + let tables = &storage_state.tables; + assert_eq!(tables.len(), 0); // if there are required components, we can optimize by only iterating through archetypes // that contain at least one of the required components let potential_archetypes = source @@ -380,75 +445,105 @@ impl<'w, 's> Iter<'w, 's> { } } } - let Some(&table) = matching_tables.first() else { + if matching_tables.first().is_none() { return false; - }; + } self.iter_state.storages_state[variable_idx] = StorageState { tables: Cow::Owned(matching_tables) }; self.set_variable_state(variable_idx as u8, VariableState::Search { - table: Some(table), offset: 0, storage_idx: 0, current_len: 0, }); - return true; + true } - } + VariableState::Search { + offset, + current_len, + storage_idx, + } => { + // we are already searching through a list of tables, we need to increment the index + assert!(self.iter_state.backtrack); + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + let iteration_start = *current_len == 0; + // loop to skip empty tables + loop { + // either beginning of the iteration, or finished processing a table, so skip to the next + if offset == current_len { + // go to next table + if !iteration_start { + *storage_idx += 1; + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + *offset = 0; + } + let table_id = storage_state.tables[*storage_idx as usize]; + let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; - // else we backtracked! we need to advance in the current table, or go to the next table - let storage_state = &self.iter_state.storages_state[variable_idx]; - let target_state = &mut self.iter_state.variable_state[variable_idx]; - let VariableState::Search { table: Some(table_id), offset, current_len, storage_idx, } = target_state else { - unreachable!(); - }; - if *storage_idx >= storage_state.tables.len() as u32 { - return false - } - let iteration_start = *current_len == 0; - // loop to skip empty tables - loop { - // either beginning of the iteration, or finished processing a table, so skip to the next - if offset == current_len { - // go to next table - if !iteration_start { - *storage_idx += 1; - if *storage_idx >= storage_state.tables.len() as u32 { - return false + *current_len = table.entity_count(); + if table.is_empty() { + // skip table + continue; + } + let storage_idx = *storage_idx; + let current_len = *current_len; + self.set_variable_state(variable_idx as u8, VariableState::Search { + offset: 0, + storage_idx, + current_len, + }); + } else { + *offset += 1; + self.set_variable_state_increment(variable_idx as u8); } - *table_id = storage_state.tables[*storage_idx as usize]; - *offset = 0; + // this entity is our current candidate + return true; } - let table = unsafe { self.world.world().storages().tables.get(*table_id).unwrap_unchecked() }; - *current_len = table.entity_count(); - let table = unsafe { self.world.world().storages().tables.get(*table_id).unwrap_unchecked() }; - if table.is_empty() { - // skip table - continue; + } + _ => { + // the entity has been fixed by a relationship, we need to check if it matches the query + // (unless we are backtracking) + if self.iter_state.backtrack { + return false } + let candidate = unsafe { self.written_entity(source.variable_idx).unwrap_unchecked() }; + let archetype = unsafe { self.world.get_entity(candidate).unwrap_unchecked().archetype() }; + source.matches(archetype) } - - // TODO: store `table.entities()` somewhere so we don't have to fetch it again every time - // let table = unsafe { self.world.world().storages().tables.get(*table_id).unwrap_unchecked() }; - // let entity = unsafe{ table.entities().get_unchecked(*offset as usize) }; - *offset += 1; - - // this entity is our current candidate - return true; } } - /// Assign the variable state (and also set the entity as written for the next op) + /// Assign the variable state (and keeps track of the previous state for backtracking) fn set_variable_state(&mut self, variable: u8, state: VariableState) { - // the operation writes the entity as written for the **next** operation let curr_op = self.iter_state.curr_op as usize; let variable_idx = variable as usize; - if curr_op + 1 < self.iter_state.written.len() { - self.iter_state.written[curr_op + 1].grow_and_insert(variable_idx); + if curr_op + 1 < self.iter_state.variable_state_diff.len() { + // store the previous state to restore on backtracking + self.iter_state.variable_state_diff[curr_op + 1] = Some(PreviousVariableState { + variable_idx: variable, + previous_state: PreviousVariableDiff::VariableState(self.iter_state.variable_state[variable_idx].clone()), + }); } self.iter_state.variable_state[variable_idx] = state; } + /// Assign the variable state (and keeps track of the previous state for backtracking) + fn set_variable_state_increment(&mut self, variable: u8) { + // the variable state has already been incremented + let curr_op = self.iter_state.curr_op as usize; + if curr_op + 1 < self.iter_state.variable_state_diff.len() { + // store the previous state to restore on backtracking + self.iter_state.variable_state_diff[curr_op + 1] = Some(PreviousVariableState { + variable_idx: variable, + previous_state: PreviousVariableDiff::Decrement, + }); + } + } + fn current_op(&self) -> &QueryOperation { &self.query_state.operations[self.iter_state.curr_op as usize] } @@ -469,6 +564,29 @@ impl<'w, 's> Iter<'w, 's> { if !matched { // Operation did not match, return to previous op self.iter_state.backtrack = true; + // restore the previous variable state if it was modified by this operation + if let Some(previous_state) = self.iter_state.variable_state_diff[op_index].take() { + let variable_idx = previous_state.variable_idx as usize; + match previous_state.previous_state { + PreviousVariableDiff::VariableState(state) => { + self.iter_state.variable_state[previous_state.variable_idx as usize] = state; + } + PreviousVariableDiff::Decrement => { + // decrement the index in the list of candidates + match &mut self.iter_state.variable_state[variable_idx] { + VariableState::Search { offset, .. } => { + *offset -= 1; + } + VariableState::FixedByRelationshipTarget { index, .. } => { + *index -= 1; + } + _ => { + unreachable!() + } + }; + } + } + } // Returned from all operations, no more matches if self.iter_state.curr_op == 0 { return false; @@ -478,16 +596,6 @@ impl<'w, 's> Iter<'w, 's> { // Operation did match, move to next op self.iter_state.backtrack = false; self.iter_state.curr_op += 1; - - if self.iter_state.curr_op < num_operations { - // Setup written state for next operation. The ops themselves have already updated the written state for the next - // op, but we also need to propagate the existing written from the current op - - // (we have a written bitset for each operation so that on backtracking we can retrieve the previous `written` state) - let (written_next_op, written_op) = self.iter_state.written.get_mut(op_index..op_index + 2).unwrap().split_at_mut(1); - // written[op_index + 1].union_with(written[op_index]) - written_next_op[0].union_with(&written_op[0]); - } } } true @@ -547,13 +655,17 @@ impl QueryPlan { source: impl Into, target: impl Into, relationship_component: ComponentId, + relationship_target_component: ComponentId, accessor: RelationshipAccessor, + target_accessor: RelationshipAccessor, ) -> &mut Self { self.operations.push(QueryOperation::Relationship(QueryRelationship { source_idx: source.into().index, target_idx: target.into().index, relationship_component, + relationship_target_component, relationship_accessor: accessor, + relationship_target_accessor: target_accessor, })); self } @@ -610,16 +722,16 @@ mod tests { // Parent does not have the marker let parent1 = world.spawn_empty().id(); - let child1 = world.spawn((Marker, ChildOf(parent1))).id(); + let _ = world.spawn((Marker, ChildOf(parent1))).id(); world.flush(); // Child does not have the marker let parent2 = world.spawn(Marker).id(); - let child2 = world.spawn(ChildOf(parent2)).id(); + let _ = world.spawn(ChildOf(parent2)).id(); // Both have markers but there is no relationship - let parent3 = world.spawn(Marker).id(); - let child3 = world.spawn(Marker).id(); + let _ = world.spawn(Marker).id(); + let _ = world.spawn(Marker).id(); // Two correct pairs, (Child, Parent) and (Parent, Grandparent) let grandparent4 = world.spawn(Marker).id(); @@ -644,3 +756,6 @@ mod tests { } } + +// Add test-case: Q1, R21, Q2 (find sources) +// Add test-case: Q1, R12, Q2, R13, Q3, R23 ($2 and $3 have been fixed, and then we check R23) \ No newline at end of file From d39f853d0c449fdf4dd707e54361ad33ebbf65a6 Mon Sep 17 00:00:00 2001 From: Charles Bournhonesque Date: Sun, 26 Oct 2025 13:57:19 -0400 Subject: [PATCH 4/7] Fixed the backtracking logic and simplified the code. Added support for relationships in both directions. Added more complicated backtracking logic and tests. --- crates/bevy_ecs/src/query/plan.rs | 550 +++++++++++++++++++++--------- 1 file changed, 389 insertions(+), 161 deletions(-) diff --git a/crates/bevy_ecs/src/query/plan.rs b/crates/bevy_ecs/src/query/plan.rs index 4e76e9688322e..7bae67e1425b0 100644 --- a/crates/bevy_ecs/src/query/plan.rs +++ b/crates/bevy_ecs/src/query/plan.rs @@ -7,10 +7,11 @@ use crate::{ }; use alloc::vec::Vec; use alloc::vec; +use fixedbitset::FixedBitSet; use bevy_ecs::archetype::Archetype; -use bevy_ecs::prelude::{QueryBuilder, World}; +use bevy_ecs::prelude::{ContainsEntity, QueryBuilder, World}; use bevy_ecs::query::{QueryData, QueryFilter}; -use bevy_ecs::relationship::Relationship; +use bevy_ecs::relationship::{Relationship, RelationshipTarget}; use bevy_ecs::storage::{SparseSetIndex, TableId}; use bevy_ecs::world::unsafe_world_cell::UnsafeEntityCell; use crate::relationship::RelationshipAccessor; @@ -57,7 +58,9 @@ impl QuerySource { /// Describes how two query terms are connected via a relationship. #[derive(Debug, Clone)] pub struct QueryRelationship { - /// The source term index within the QueryPlan + /// The source term index within the QueryPlan. The source is always from the point of view of the Relationship component, not the RElationshipTarget. + /// ChildOf(0, 1) adds 0 as source, 1 as target. + /// Children(0, 1) adds 1 as source, 0 as target. pub source_idx: u8, /// The target term index within the QueryPlan pub target_idx: u8, @@ -181,6 +184,37 @@ impl<'w> QueryPlanBuilder<'w> { ); } + /// Add a relationship between two terms using a typed Relationship component. + pub fn add_relationship_target( + &mut self, + source_term: u8, + target_term: u8, + ) { + let component_id = self.world.register_component::<::Relationship>(); + let target_component_id = self.world.register_component::(); + let accessor = self.world + .components() + .get_info(component_id) + .unwrap() + .relationship_accessor() + .unwrap().clone(); + let target_accessor = self.world + .components() + .get_info(target_component_id) + .unwrap() + .relationship_accessor() + .unwrap().clone(); + + self.plan.add_relationship( + target_term, + source_term, + component_id, + target_component_id, + accessor, + target_accessor, + ); + } + pub fn compile(self) -> QueryPlan { self.plan.compile() } @@ -255,11 +289,10 @@ pub struct IterState<'w> { pub curr_op: u8, /// Index of the current entity for each variable pub variable_state: Vec, - /// For each operation, store the previous variable state to restore on backtracking - /// in case this operation modified it - variable_state_diff: Vec>, /// List of matching tables/archetypes to iterate through for each variable pub storages_state: Vec>, + /// For each operation, which variables were written by the operation + written: Vec, } pub struct PreviousVariableState { @@ -277,16 +310,15 @@ impl<'w> IterState<'w> { fn new(plan: &QueryPlan) -> Self { let num_operations = plan.operations.len(); let num_variables = plan.num_variables as usize; - let mut variable_state = vec![VariableState::default(); num_variables]; - let mut variable_state_diff = Vec::new(); - variable_state_diff.resize_with(num_operations, Default::default); - let mut storages_state = vec![StorageState::default(); num_variables]; + let variable_state = vec![VariableState::default(); num_variables]; + let storages_state = vec![StorageState::default(); num_variables]; + let written = vec![FixedBitSet::with_capacity(num_variables); num_operations]; Self { backtrack: false, curr_op: 0, variable_state, - variable_state_diff, storages_state, + written, } } @@ -303,9 +335,9 @@ pub struct Iter<'w, 's> { impl<'w, 's> Iter<'w, 's> { - /// Returns true if the `variable` was currently assigned an Entity + /// Returns true if the `variable` was assigned an Entity in the current operation fn written(&self, variable: u8) -> bool { - !matches!(self.iter_state.variable_state[variable as usize], VariableState::Unknown) + self.iter_state.written[self.iter_state.curr_op as usize].contains(variable as usize) } /// Get the entity currently written to the variable indexed by `variable`, if it has been written @@ -342,39 +374,21 @@ impl<'w, 's> Iter<'w, 's> { let QueryOperation::Relationship(rel) = self.current_op() else { unreachable!() }; - let source_idx = rel.source_idx as usize; - let target_idx = rel.target_idx as usize; if self.iter_state.backtrack { - // possibly check the next source in the list - if let VariableState::FixedByRelationshipTarget { sources, index } = &mut self.iter_state.variable_state[source_idx] { - // TODO: this needs to be done via set_variable_state to track the changes for backtracking - *index += 1; - if *index >= sources.len() { - return false; - } - self.set_variable_state_increment(source_idx as u8); - return true; + if self.written(rel.source_idx) { + return self.backtrack_variable(rel.source_idx); } - // possibly check the next source in the list - if let VariableState::FixedByRelationshipTarget { sources, index } = &mut self.iter_state.variable_state[target_idx] { - // TODO: this needs to be done via set_variable_state to track the changes for backtracking - *index += 1; - if *index >= sources.len() { - return false; - } - self.set_variable_state_increment(target_idx as u8); - return true; + if self.written(rel.target_idx) { + return self.backtrack_variable(rel.target_idx); } return false; } - - match (self.written(rel.source_idx), self.written(rel.target_idx)) { - (false, false) => { + match (self.written_entity(rel.source_idx), self.written_entity(rel.target_idx)) { + (None, None) => { unreachable!("we only support queries where the source has been found before we are querying the relationship"); } - (true, false) => { + (Some(source_entity), None) => { // we found the source, need to find the target - let source_entity = self.written_entity(rel.source_idx).unwrap(); match unsafe { rel.get_target(source_entity, self.world) } { None => false, Some(target_entity) => { @@ -383,9 +397,8 @@ impl<'w, 's> Iter<'w, 's> { } } } - (false, true) => { + (None, Some(target_entity)) => { // we found the target, need to find the source - let target_entity = self.written_entity(rel.target_idx).unwrap(); match unsafe { rel.get_sources(target_entity, self.world) } { None => false, Some(sources) => { @@ -397,12 +410,9 @@ impl<'w, 's> Iter<'w, 's> { } } } - (true, true) => { + (Some(source_entity), Some(target_entity)) => { // we found both, need to check if they are linked by the relationship - let source_entity = self.written_entity(rel.source_idx).unwrap(); - let target_entity = self.written_entity(rel.target_idx).unwrap(); - let expected_target_entity = unsafe { rel.get_target(source_entity, self.world).unwrap_unchecked() }; - target_entity == expected_target_entity + unsafe { rel.get_target(source_entity, self.world) }.is_some_and(|expected_target_entity| target_entity == expected_target_entity) } } @@ -416,6 +426,9 @@ impl<'w, 's> Iter<'w, 's> { let storage_state = &self.iter_state.storages_state[variable_idx]; match &mut self.iter_state.variable_state[variable_idx] { VariableState::Unknown => { + if self.iter_state.backtrack { + return false + } // the first time we evaluate the query, we set the list of potential tables let tables = &storage_state.tables; assert_eq!(tables.len(), 0); @@ -423,10 +436,8 @@ impl<'w, 's> Iter<'w, 's> { // that contain at least one of the required components let potential_archetypes = source .access - .required - .ones() - .filter_map(|idx| { - let component_id = ComponentId::get_sparse_set_index(idx); + .with_filters() + .filter_map(|component_id| { self.world .archetypes() .component_index() @@ -436,11 +447,20 @@ impl<'w, 's> Iter<'w, 's> { // select the component with the fewest archetypes .min_by_key(ExactSizeIterator::len); let mut matching_tables = Vec::new(); + let mut current_len = 0; if let Some(archetypes) = potential_archetypes { for archetype_id in archetypes { // SAFETY: get_potential_archetypes only returns archetype ids that are valid for the world let archetype = &self.world.archetypes()[*archetype_id]; - if source.matches(archetype) { + // NOTE: you could have empty archetypes even if the table is not empty + // and you could have non-empty archetypes with empty tables (e.g. when the archetype only has sparse sets) + let table_id = archetype.table_id(); + let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; + // skip empty tables + if !table.is_empty() && source.matches(archetype) { + if current_len == 0 { + current_len = table.entity_count(); + } matching_tables.push(archetype.table_id()) } } @@ -448,61 +468,22 @@ impl<'w, 's> Iter<'w, 's> { if matching_tables.first().is_none() { return false; } + self.iter_state.storages_state[variable_idx] = StorageState { tables: Cow::Owned(matching_tables) }; self.set_variable_state(variable_idx as u8, VariableState::Search { offset: 0, storage_idx: 0, - current_len: 0, + current_len, }); true } - VariableState::Search { - offset, - current_len, - storage_idx, - } => { + VariableState::Search { .. } => { // we are already searching through a list of tables, we need to increment the index assert!(self.iter_state.backtrack); - if *storage_idx >= storage_state.tables.len() as u32 { - return false - } - let iteration_start = *current_len == 0; - // loop to skip empty tables - loop { - // either beginning of the iteration, or finished processing a table, so skip to the next - if offset == current_len { - // go to next table - if !iteration_start { - *storage_idx += 1; - if *storage_idx >= storage_state.tables.len() as u32 { - return false - } - *offset = 0; - } - let table_id = storage_state.tables[*storage_idx as usize]; - let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; - - *current_len = table.entity_count(); - if table.is_empty() { - // skip table - continue; - } - let storage_idx = *storage_idx; - let current_len = *current_len; - self.set_variable_state(variable_idx as u8, VariableState::Search { - offset: 0, - storage_idx, - current_len, - }); - } else { - *offset += 1; - self.set_variable_state_increment(variable_idx as u8); - } - // this entity is our current candidate - return true; - } + assert!(self.written(source.variable_idx)); + self.backtrack_variable(source.variable_idx) } _ => { // the entity has been fixed by a relationship, we need to check if it matches the query @@ -517,36 +498,83 @@ impl<'w, 's> Iter<'w, 's> { } } - /// Assign the variable state (and keeps track of the previous state for backtracking) + /// Assign the variable state and keep track that the variable was written in the current operation fn set_variable_state(&mut self, variable: u8, state: VariableState) { - let curr_op = self.iter_state.curr_op as usize; - let variable_idx = variable as usize; - if curr_op + 1 < self.iter_state.variable_state_diff.len() { - // store the previous state to restore on backtracking - self.iter_state.variable_state_diff[curr_op + 1] = Some(PreviousVariableState { - variable_idx: variable, - previous_state: PreviousVariableDiff::VariableState(self.iter_state.variable_state[variable_idx].clone()), - }); - } - self.iter_state.variable_state[variable_idx] = state; + self.iter_state.written[self.iter_state.curr_op as usize].insert(variable as usize); + self.iter_state.variable_state[variable as usize] = state; + } + + fn current_op(&self) -> &QueryOperation { + &self.query_state.operations[self.iter_state.curr_op as usize] } - /// Assign the variable state (and keeps track of the previous state for backtracking) - fn set_variable_state_increment(&mut self, variable: u8) { - // the variable state has already been incremented + /// Backtrack the variable if it had been written in the current operation + fn backtrack_variable(&mut self, variable: u8) -> bool { let curr_op = self.iter_state.curr_op as usize; - if curr_op + 1 < self.iter_state.variable_state_diff.len() { - // store the previous state to restore on backtracking - self.iter_state.variable_state_diff[curr_op + 1] = Some(PreviousVariableState { - variable_idx: variable, - previous_state: PreviousVariableDiff::Decrement, - }); + let variable_idx = variable as usize; + // wrap in closure to allow early return + let res = (|| { + match &mut self.iter_state.variable_state[variable_idx] { + VariableState::Unknown => { + unreachable!() + } + VariableState::Search { offset, current_len, storage_idx } => { + let storage_state = &self.iter_state.storages_state[variable_idx]; + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + // either beginning of the iteration, or finished processing a table, so skip to the next + if (*offset + 1) == *current_len { + // go to next table + *storage_idx += 1; + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + *offset = 0; + let table_id = storage_state.tables[*storage_idx as usize]; + let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; + *current_len = table.entity_count(); + } else { + *offset += 1; + } + return true + } + VariableState::FixedByRelationship(_) => { + false + } + VariableState::FixedByRelationshipTarget { sources, index } => { + *index += 1; + if *index >= sources.len() { + return false + } + true + } + } + })(); + if !res { + // if we don't find a next candidate, reset the variable state to Unknown and remove from written + self.iter_state.written[curr_op].remove(variable_idx); + self.iter_state.variable_state[variable_idx] = VariableState::Unknown; } + res } - fn current_op(&self) -> &QueryOperation { - &self.query_state.operations[self.iter_state.curr_op as usize] - } + // fn backtrack(&mut self) { + // self.iter_state.backtrack = true; + // self.iter_state.curr_op -= 1; + // self.iter_state.writt + // if let Some(previous_state) = self.iter_state.variable_state_diff[self.iter_state.curr_op as usize].take() { + // let variable_idx = previous_state.variable_idx as usize; + // match previous_state.previous_state { + // PreviousVariableDiff::VariableState(state) => { + // self.iter_state.variable_state[variable_idx] = state; + // } + // PreviousVariableDiff::Decrement => { + // // note: we do nothing! because we need to keep incrementing the index in the list of candidates + // } + // } + // } + // } /// Check if we can still find a matching item fn next_match(&mut self) -> bool { @@ -558,39 +586,13 @@ impl<'w, 's> Iter<'w, 's> { } while self.iter_state.curr_op < num_operations { - let op_index = self.iter_state.curr_op as usize; let matched = self.dispatch(); - if !matched { - // Operation did not match, return to previous op - self.iter_state.backtrack = true; - // restore the previous variable state if it was modified by this operation - if let Some(previous_state) = self.iter_state.variable_state_diff[op_index].take() { - let variable_idx = previous_state.variable_idx as usize; - match previous_state.previous_state { - PreviousVariableDiff::VariableState(state) => { - self.iter_state.variable_state[previous_state.variable_idx as usize] = state; - } - PreviousVariableDiff::Decrement => { - // decrement the index in the list of candidates - match &mut self.iter_state.variable_state[variable_idx] { - VariableState::Search { offset, .. } => { - *offset -= 1; - } - VariableState::FixedByRelationshipTarget { index, .. } => { - *index -= 1; - } - _ => { - unreachable!() - } - }; - } - } - } // Returned from all operations, no more matches if self.iter_state.curr_op == 0 { return false; } + self.iter_state.backtrack = true; self.iter_state.curr_op -= 1; } else { // Operation did match, move to next op @@ -607,6 +609,14 @@ pub struct DynamicItem<'w> { entities: Vec>, } +impl core::fmt::Debug for DynamicItem<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_list() + .entries(self.entities.iter().map(|e| e.entity())) + .finish() + } +} + impl<'w, 's> Iterator for Iter<'w, 's> { type Item = DynamicItem<'w>; @@ -680,6 +690,8 @@ impl QueryPlan { #[cfg(test)] mod tests { + use std::dbg; + use std::process::Child; use bevy_ecs::prelude::*; use super::*; use crate::{ @@ -689,20 +701,43 @@ mod tests { }; #[derive(Component)] - struct Marker; + struct A; + + #[derive(Component)] + struct B; + + #[derive(Component)] + struct C; + + #[derive(Component)] + #[relationship(relationship_target = R1Target)] + struct R1(#[entities] Entity); + #[derive(Component)] + #[relationship_target(relationship = R1)] + struct R1Target(#[entities] Vec); + + #[derive(Component)] + #[relationship(relationship_target = R2Target)] + struct R2(#[entities] Entity); + + #[derive(Component)] + #[relationship_target(relationship = R2)] + struct R2Target(#[entities] Vec); + + /// Q1 -> R12 -> Q2 #[test] fn test_query_plan_basic() { let mut world = World::new(); // Correct pair let parent = world.spawn_empty().id(); - let child = world.spawn((Marker, ChildOf(parent))).id(); + let child = world.spawn((A, ChildOf(parent))).id(); world.flush(); // Build a simple plan using the builder API let mut builder = QueryPlanBuilder::new(&mut world); - builder.add_source::<&Marker, ()>(); + builder.add_source::<&A, ()>(); builder.add_relationship::(0, 1); builder.add_source::(); let plan = builder.compile(); @@ -716,46 +751,239 @@ mod tests { } /// Checks that filters in the source or target of the relationship are respected + /// Q1 -> R12 -> Q2 #[test] fn test_query_plan_single_relationship() { let mut world = World::new(); // Parent does not have the marker let parent1 = world.spawn_empty().id(); - let _ = world.spawn((Marker, ChildOf(parent1))).id(); + let _ = world.spawn((A, ChildOf(parent1))).id(); world.flush(); // Child does not have the marker - let parent2 = world.spawn(Marker).id(); + let parent2 = world.spawn(A).id(); let _ = world.spawn(ChildOf(parent2)).id(); // Both have markers but there is no relationship - let _ = world.spawn(Marker).id(); - let _ = world.spawn(Marker).id(); + let _ = world.spawn(A).id(); + let _ = world.spawn(A).id(); // Two correct pairs, (Child, Parent) and (Parent, Grandparent) - let grandparent4 = world.spawn(Marker).id(); - let parent4 = world.spawn((Marker, ChildOf(grandparent4))).id(); - let child4 = world.spawn((Marker, ChildOf(parent4))).id(); + let grandparent4 = world.spawn(A).id(); + let parent4 = world.spawn((A, ChildOf(grandparent4))).id(); + let child4 = world.spawn((A, ChildOf(parent4))).id(); // Both sources must have the Marker let mut builder = QueryPlanBuilder::new(&mut world); - builder.add_source::<&Marker, ()>(); + builder.add_source::>(); builder.add_relationship::(0, 1); - builder.add_source::>(); + builder.add_source::>(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].entities[0].entity(), parent4); + assert_eq!(results[0].entities[1].entity(), grandparent4); + assert_eq!(results[1].entities[0].entity(), child4); + assert_eq!(results[1].entities[1].entity(), parent4); + } + + /// Checks that filters in the source or target of the relationship are respected + /// Q1 -> R12 -> Q2 (but expressed using relationship target) + #[test] + fn test_query_plan_single_relationship_target() { + let mut world = World::new(); + + // Parent does not have the marker + let parent1 = world.spawn_empty().id(); + let _ = world.spawn((A, ChildOf(parent1))).id(); + world.flush(); + + // Child does not have the marker + let parent2 = world.spawn(A).id(); + let _ = world.spawn(ChildOf(parent2)).id(); + + // Both have markers but there is no relationship + let _ = world.spawn(A).id(); + let _ = world.spawn(A).id(); + + // Two correct pairs, (Child, Parent) and (Parent, Grandparent) + let grandparent4 = world.spawn(A).id(); + let parent4 = world.spawn((A, ChildOf(grandparent4))).id(); + let child4 = world.spawn((A, ChildOf(parent4))).id(); + + // Both sources must have the Marker + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::>(); + builder.add_relationship_target::(1, 0); + builder.add_source::>(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].entities[0].entity(), parent4); + assert_eq!(results[0].entities[1].entity(), grandparent4); + assert_eq!(results[1].entities[0].entity(), child4); + assert_eq!(results[1].entities[1].entity(), parent4); + } + + /// The first variable $1 is the target of the relationship + /// Q1 -> R21 -> Q2 + #[test] + fn test_query_plan_single_relationship_reverse() { + let mut world = World::new(); + + // Parent does not have the marker + let parent1 = world.spawn_empty().id(); + let _ = world.spawn((A, ChildOf(parent1))).id(); + world.flush(); + + // Child does not have the marker + let parent2 = world.spawn(A).id(); + let _ = world.spawn(ChildOf(parent2)).id(); + + // Both have markers but there is no relationship + let _ = world.spawn(A).id(); + let _ = world.spawn(A).id(); + + // Two correct pairs, (Child, Parent) and (Parent, Grandparent) + let grandparent4 = world.spawn(A).id(); + let parent4 = world.spawn((A, ChildOf(grandparent4))).id(); + let child4 = world.spawn((A, ChildOf(parent4))).id(); + + // Both sources must have the Marker + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::>(); + builder.add_relationship::(1, 0); + builder.add_source::>(); let plan = builder.compile(); let iter = plan.query_iter(world.as_unsafe_world_cell()); let results: Vec = iter.collect(); assert_eq!(results.len(), 2); - assert_eq!(results[0].entities[0].entity(), child4); - assert_eq!(results[0].entities[1].entity(), parent4); - assert_eq!(results[1].entities[0].entity(), parent4); - assert_eq!(results[2].entities[1].entity(), grandparent4); + assert_eq!(results[0].entities[0].entity(), parent4); + assert_eq!(results[0].entities[1].entity(), child4); + assert_eq!(results[1].entities[0].entity(), grandparent4); + assert_eq!(results[1].entities[1].entity(), parent4); } + + /// Q1 -> R12 -> Q2 -> R23 -> Q3 + #[test] + fn test_query_plan_multi_relationship() { + let mut world = World::new(); + + // Valid triplet + let valid_3 = world.spawn(C).id(); + let valid_2 = world.spawn((B, R2(valid_3))).id(); + let valid_1 = world.spawn((A, R1(valid_2))).id(); + + // Invalid triplet: the constraint Q3 is not satisfied + let invalid_3_a = world.spawn_empty().id(); + let invalid_2_a = world.spawn((B, R2(invalid_3_a))).id(); + let _ = world.spawn((A, R1(invalid_2_a))).id(); + + // Invalid triplet: the constraint Q3 is not satisfied + let invalid_3_b = world.spawn(C).id(); + let invalid_2_b = world.spawn(R2(invalid_3_b)).id(); + let _ = world.spawn((A, R1(invalid_2_b))).id(); + + // Invalid triplet: the constraint Q1 is not satisfied + let invalid_3_c = world.spawn(C).id(); + let invalid_2_c = world.spawn((B, R2(invalid_3_c))).id(); + let _ = world.spawn(R1(invalid_2_c)).id(); + + // Invalid triplet: the constraint R23 is not satisfied + let invalid_3_d = world.spawn(C).id(); + let invalid_2_d = world.spawn(B).id(); + let _ = world.spawn((A, R1(invalid_2_d))).id(); + + // Invalid triplet: the constraint R12 is not satisfied + let invalid_3_e = world.spawn(C).id(); + let invalid_2_e = world.spawn((B, R2(invalid_3_e))).id(); + let _ = world.spawn(A).id(); + + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::<(), With>(); + builder.add_relationship::(0, 1); + builder.add_source::<(), With>(); + builder.add_relationship::(1, 2); + builder.add_source::<(), With>(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].entities[0].entity(), valid_1); + assert_eq!(results[0].entities[1].entity(), valid_2); + assert_eq!(results[0].entities[2].entity(), valid_3); + } + + + /// Q1 -> R12 -> Q2 -> R23 -> Q3 -> R13 + /// At the end, we backtrack through R13 but we don't modify the fixed state of $1 and $3 since + /// they were written in previous ops + #[test] + fn test_query_plan_multi_relationship_fixed_backtrack() { + let mut world = World::new(); + + // Valid triplet + let valid_3 = world.spawn(C).id(); + let valid_2 = world.spawn((B, R1(valid_3))).id(); + let valid_1 = world.spawn((A, ChildOf(valid_2), R2(valid_3))).id(); + + // Invalid triplet: the last constraint (R2) is not satisfied + let invalid_3 = world.spawn(C).id(); + let invalid_2 = world.spawn((B, R1(valid_3))).id(); + let invalid_1 = world.spawn((A, ChildOf(valid_2))).id(); + + world.flush(); + + // Both sources must have the Marker + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::<(), With>(); + builder.add_relationship::(0, 1); + builder.add_source::<(), With>(); + builder.add_relationship::(1, 2); + builder.add_source::<(), With>(); + builder.add_relationship::(0, 2); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].entities[0].entity(), valid_1); + assert_eq!(results[0].entities[1].entity(), valid_2); + assert_eq!(results[0].entities[2].entity(), valid_3); + } + } +// Vec to track which variables were written. +// In each op, keep track of which variables were written by the op. +// The op only writes when it modifies the variable state from Unknown to something else. (Relationship when (true, false) or (false, true)) +// For query, when (Unknown -> Search) +// On backtracking, we reset the op's 'written' variables to zero. + +// - when we backtrack a source that needs to be iterated, we should go to the next entity in the list of candidates +// instead of resetting it to Unknown + + +// Q1, R12, Q2 +// - on backtrack of Q2: everything was fixed so we don't modify written +// - on backtrack of R12: $2 was written here, we reset written + // Add test-case: Q1, R21, Q2 (find sources) -// Add test-case: Q1, R12, Q2, R13, Q3, R23 ($2 and $3 have been fixed, and then we check R23) \ No newline at end of file +// - on backtrack of R21: we should increment the list of candidates for $2. $2 is still marked as written. + +// Add test-case: Q1, R12, Q2, R13, Q3, R23, R31 ($2 and $3 have been fixed, and then we check R23) +// - when we backtrack R31, we shouldn't modify the search state for $1 from Q1. +// - when we backtrack R23, we shouldn't remove the fixed state of $2 and $3. From 6211d05140249c3e636a98a3c589760566850377 Mon Sep 17 00:00:00 2001 From: Charles Bournhonesque Date: Sun, 26 Oct 2025 13:57:19 -0400 Subject: [PATCH 5/7] Fixed the backtracking logic and simplified the code. Added support for relationships in both directions. Added more complicated backtracking logic and tests. --- crates/bevy_ecs/src/query/plan.rs | 535 ++++++++++++++++++++---------- 1 file changed, 354 insertions(+), 181 deletions(-) diff --git a/crates/bevy_ecs/src/query/plan.rs b/crates/bevy_ecs/src/query/plan.rs index 4e76e9688322e..d10acdc3ce112 100644 --- a/crates/bevy_ecs/src/query/plan.rs +++ b/crates/bevy_ecs/src/query/plan.rs @@ -7,10 +7,11 @@ use crate::{ }; use alloc::vec::Vec; use alloc::vec; +use fixedbitset::FixedBitSet; use bevy_ecs::archetype::Archetype; -use bevy_ecs::prelude::{QueryBuilder, World}; +use bevy_ecs::prelude::{ContainsEntity, QueryBuilder, World}; use bevy_ecs::query::{QueryData, QueryFilter}; -use bevy_ecs::relationship::Relationship; +use bevy_ecs::relationship::{Relationship, RelationshipTarget}; use bevy_ecs::storage::{SparseSetIndex, TableId}; use bevy_ecs::world::unsafe_world_cell::UnsafeEntityCell; use crate::relationship::RelationshipAccessor; @@ -23,11 +24,6 @@ pub enum QueryOperation { Source(QuerySource), } -// - we want to track when we change the VariableState, so that we can backtrack and restore to the previous state -// - we cannot just cancel a VariableState::Fixed, because the variable might have been Fixed in a previous operation. -// - we want to handle queries like Q1, R21 (i.e. find all entities $2 where R($1, $2), which means the possibilities for $2 are within a list of entities: RelationshipTarget) - - #[derive(Debug, Clone)] pub struct QuerySource { access: FilteredAccess, @@ -57,7 +53,9 @@ impl QuerySource { /// Describes how two query terms are connected via a relationship. #[derive(Debug, Clone)] pub struct QueryRelationship { - /// The source term index within the QueryPlan + /// The source term index within the QueryPlan. The source is always from the point of view of the Relationship component, not the RElationshipTarget. + /// ChildOf(0, 1) adds 0 as source, 1 as target. + /// Children(0, 1) adds 1 as source, 0 as target. pub source_idx: u8, /// The target term index within the QueryPlan pub target_idx: u8, @@ -181,6 +179,37 @@ impl<'w> QueryPlanBuilder<'w> { ); } + /// Add a relationship between two terms using a typed Relationship component. + pub fn add_relationship_target( + &mut self, + source_term: u8, + target_term: u8, + ) { + let component_id = self.world.register_component::<::Relationship>(); + let target_component_id = self.world.register_component::(); + let accessor = self.world + .components() + .get_info(component_id) + .unwrap() + .relationship_accessor() + .unwrap().clone(); + let target_accessor = self.world + .components() + .get_info(target_component_id) + .unwrap() + .relationship_accessor() + .unwrap().clone(); + + self.plan.add_relationship( + target_term, + source_term, + component_id, + target_component_id, + accessor, + target_accessor, + ); + } + pub fn compile(self) -> QueryPlan { self.plan.compile() } @@ -245,6 +274,7 @@ impl Default for VariableState { #[derive(Default, Clone)] pub struct StorageState<'w> { + /// List of tables that are being iterated for this variable. (maybe merge with VariableState::Search?) tables: Cow<'w, [TableId]>, } @@ -255,38 +285,26 @@ pub struct IterState<'w> { pub curr_op: u8, /// Index of the current entity for each variable pub variable_state: Vec, - /// For each operation, store the previous variable state to restore on backtracking - /// in case this operation modified it - variable_state_diff: Vec>, /// List of matching tables/archetypes to iterate through for each variable pub storages_state: Vec>, -} - -pub struct PreviousVariableState { - variable_idx: u8, - previous_state: PreviousVariableDiff, -} - -pub enum PreviousVariableDiff { - VariableState(VariableState), - // the previous variable state is obtained by decrementing the index in the list of candidates - Decrement, + /// For each operation, which variables were written by the operation + /// This is useful to know if we need to backtrack a variable when backtracking the operation. + written: Vec, } impl<'w> IterState<'w> { fn new(plan: &QueryPlan) -> Self { let num_operations = plan.operations.len(); let num_variables = plan.num_variables as usize; - let mut variable_state = vec![VariableState::default(); num_variables]; - let mut variable_state_diff = Vec::new(); - variable_state_diff.resize_with(num_operations, Default::default); - let mut storages_state = vec![StorageState::default(); num_variables]; + let variable_state = vec![VariableState::default(); num_variables]; + let storages_state = vec![StorageState::default(); num_variables]; + let written = vec![FixedBitSet::with_capacity(num_variables); num_operations]; Self { backtrack: false, curr_op: 0, variable_state, - variable_state_diff, storages_state, + written, } } @@ -303,9 +321,9 @@ pub struct Iter<'w, 's> { impl<'w, 's> Iter<'w, 's> { - /// Returns true if the `variable` was currently assigned an Entity + /// Returns true if the `variable` was assigned an Entity in the current operation fn written(&self, variable: u8) -> bool { - !matches!(self.iter_state.variable_state[variable as usize], VariableState::Unknown) + self.iter_state.written[self.iter_state.curr_op as usize].contains(variable as usize) } /// Get the entity currently written to the variable indexed by `variable`, if it has been written @@ -328,7 +346,6 @@ impl<'w, 's> Iter<'w, 's> { } } - // we found entity1, QueryOp:Rel -> set entity 2 to one of the entities via the RelationShip fn dispatch(&mut self) -> bool { let op = self.current_op(); match op { @@ -342,39 +359,21 @@ impl<'w, 's> Iter<'w, 's> { let QueryOperation::Relationship(rel) = self.current_op() else { unreachable!() }; - let source_idx = rel.source_idx as usize; - let target_idx = rel.target_idx as usize; if self.iter_state.backtrack { - // possibly check the next source in the list - if let VariableState::FixedByRelationshipTarget { sources, index } = &mut self.iter_state.variable_state[source_idx] { - // TODO: this needs to be done via set_variable_state to track the changes for backtracking - *index += 1; - if *index >= sources.len() { - return false; - } - self.set_variable_state_increment(source_idx as u8); - return true; + if self.written(rel.source_idx) { + return self.backtrack_variable(rel.source_idx); } - // possibly check the next source in the list - if let VariableState::FixedByRelationshipTarget { sources, index } = &mut self.iter_state.variable_state[target_idx] { - // TODO: this needs to be done via set_variable_state to track the changes for backtracking - *index += 1; - if *index >= sources.len() { - return false; - } - self.set_variable_state_increment(target_idx as u8); - return true; + if self.written(rel.target_idx) { + return self.backtrack_variable(rel.target_idx); } return false; } - - match (self.written(rel.source_idx), self.written(rel.target_idx)) { - (false, false) => { + match (self.written_entity(rel.source_idx), self.written_entity(rel.target_idx)) { + (None, None) => { unreachable!("we only support queries where the source has been found before we are querying the relationship"); } - (true, false) => { + (Some(source_entity), None) => { // we found the source, need to find the target - let source_entity = self.written_entity(rel.source_idx).unwrap(); match unsafe { rel.get_target(source_entity, self.world) } { None => false, Some(target_entity) => { @@ -383,9 +382,8 @@ impl<'w, 's> Iter<'w, 's> { } } } - (false, true) => { + (None, Some(target_entity)) => { // we found the target, need to find the source - let target_entity = self.written_entity(rel.target_idx).unwrap(); match unsafe { rel.get_sources(target_entity, self.world) } { None => false, Some(sources) => { @@ -397,12 +395,9 @@ impl<'w, 's> Iter<'w, 's> { } } } - (true, true) => { + (Some(source_entity), Some(target_entity)) => { // we found both, need to check if they are linked by the relationship - let source_entity = self.written_entity(rel.source_idx).unwrap(); - let target_entity = self.written_entity(rel.target_idx).unwrap(); - let expected_target_entity = unsafe { rel.get_target(source_entity, self.world).unwrap_unchecked() }; - target_entity == expected_target_entity + unsafe { rel.get_target(source_entity, self.world) }.is_some_and(|expected_target_entity| target_entity == expected_target_entity) } } @@ -416,6 +411,9 @@ impl<'w, 's> Iter<'w, 's> { let storage_state = &self.iter_state.storages_state[variable_idx]; match &mut self.iter_state.variable_state[variable_idx] { VariableState::Unknown => { + if self.iter_state.backtrack { + return false + } // the first time we evaluate the query, we set the list of potential tables let tables = &storage_state.tables; assert_eq!(tables.len(), 0); @@ -423,10 +421,8 @@ impl<'w, 's> Iter<'w, 's> { // that contain at least one of the required components let potential_archetypes = source .access - .required - .ones() - .filter_map(|idx| { - let component_id = ComponentId::get_sparse_set_index(idx); + .with_filters() + .filter_map(|component_id| { self.world .archetypes() .component_index() @@ -436,11 +432,20 @@ impl<'w, 's> Iter<'w, 's> { // select the component with the fewest archetypes .min_by_key(ExactSizeIterator::len); let mut matching_tables = Vec::new(); + let mut current_len = 0; if let Some(archetypes) = potential_archetypes { for archetype_id in archetypes { // SAFETY: get_potential_archetypes only returns archetype ids that are valid for the world let archetype = &self.world.archetypes()[*archetype_id]; - if source.matches(archetype) { + // NOTE: you could have empty archetypes even if the table is not empty + // and you could have non-empty archetypes with empty tables (e.g. when the archetype only has sparse sets) + let table_id = archetype.table_id(); + let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; + // skip empty tables + if !table.is_empty() && source.matches(archetype) { + if current_len == 0 { + current_len = table.entity_count(); + } matching_tables.push(archetype.table_id()) } } @@ -448,61 +453,22 @@ impl<'w, 's> Iter<'w, 's> { if matching_tables.first().is_none() { return false; } + self.iter_state.storages_state[variable_idx] = StorageState { tables: Cow::Owned(matching_tables) }; self.set_variable_state(variable_idx as u8, VariableState::Search { offset: 0, storage_idx: 0, - current_len: 0, + current_len, }); true } - VariableState::Search { - offset, - current_len, - storage_idx, - } => { + VariableState::Search { .. } => { // we are already searching through a list of tables, we need to increment the index assert!(self.iter_state.backtrack); - if *storage_idx >= storage_state.tables.len() as u32 { - return false - } - let iteration_start = *current_len == 0; - // loop to skip empty tables - loop { - // either beginning of the iteration, or finished processing a table, so skip to the next - if offset == current_len { - // go to next table - if !iteration_start { - *storage_idx += 1; - if *storage_idx >= storage_state.tables.len() as u32 { - return false - } - *offset = 0; - } - let table_id = storage_state.tables[*storage_idx as usize]; - let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; - - *current_len = table.entity_count(); - if table.is_empty() { - // skip table - continue; - } - let storage_idx = *storage_idx; - let current_len = *current_len; - self.set_variable_state(variable_idx as u8, VariableState::Search { - offset: 0, - storage_idx, - current_len, - }); - } else { - *offset += 1; - self.set_variable_state_increment(variable_idx as u8); - } - // this entity is our current candidate - return true; - } + assert!(self.written(source.variable_idx)); + self.backtrack_variable(source.variable_idx) } _ => { // the entity has been fixed by a relationship, we need to check if it matches the query @@ -517,37 +483,67 @@ impl<'w, 's> Iter<'w, 's> { } } - /// Assign the variable state (and keeps track of the previous state for backtracking) + /// Assign the variable state and keep track that the variable was written in the current operation fn set_variable_state(&mut self, variable: u8, state: VariableState) { - let curr_op = self.iter_state.curr_op as usize; - let variable_idx = variable as usize; - if curr_op + 1 < self.iter_state.variable_state_diff.len() { - // store the previous state to restore on backtracking - self.iter_state.variable_state_diff[curr_op + 1] = Some(PreviousVariableState { - variable_idx: variable, - previous_state: PreviousVariableDiff::VariableState(self.iter_state.variable_state[variable_idx].clone()), - }); - } - self.iter_state.variable_state[variable_idx] = state; - } - - /// Assign the variable state (and keeps track of the previous state for backtracking) - fn set_variable_state_increment(&mut self, variable: u8) { - // the variable state has already been incremented - let curr_op = self.iter_state.curr_op as usize; - if curr_op + 1 < self.iter_state.variable_state_diff.len() { - // store the previous state to restore on backtracking - self.iter_state.variable_state_diff[curr_op + 1] = Some(PreviousVariableState { - variable_idx: variable, - previous_state: PreviousVariableDiff::Decrement, - }); - } + self.iter_state.written[self.iter_state.curr_op as usize].insert(variable as usize); + self.iter_state.variable_state[variable as usize] = state; } fn current_op(&self) -> &QueryOperation { &self.query_state.operations[self.iter_state.curr_op as usize] } + /// Backtrack the variable if it had been written in the current operation + fn backtrack_variable(&mut self, variable: u8) -> bool { + let curr_op = self.iter_state.curr_op as usize; + let variable_idx = variable as usize; + // wrap in closure to allow early return + let res = (|| { + match &mut self.iter_state.variable_state[variable_idx] { + VariableState::Unknown => { + unreachable!() + } + VariableState::Search { offset, current_len, storage_idx } => { + let storage_state = &self.iter_state.storages_state[variable_idx]; + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + // either beginning of the iteration, or finished processing a table, so skip to the next + if (*offset + 1) == *current_len { + // go to next table + *storage_idx += 1; + if *storage_idx >= storage_state.tables.len() as u32 { + return false + } + *offset = 0; + let table_id = storage_state.tables[*storage_idx as usize]; + let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; + *current_len = table.entity_count(); + } else { + *offset += 1; + } + return true + } + VariableState::FixedByRelationship(_) => { + false + } + VariableState::FixedByRelationshipTarget { sources, index } => { + *index += 1; + if *index >= sources.len() { + return false + } + true + } + } + })(); + if !res { + // if we don't find a next candidate, reset the variable state to Unknown and remove from written + self.iter_state.written[curr_op].remove(variable_idx); + self.iter_state.variable_state[variable_idx] = VariableState::Unknown; + } + res + } + /// Check if we can still find a matching item fn next_match(&mut self) -> bool { let num_operations = self.query_state.operations.len() as u8; @@ -558,39 +554,13 @@ impl<'w, 's> Iter<'w, 's> { } while self.iter_state.curr_op < num_operations { - let op_index = self.iter_state.curr_op as usize; let matched = self.dispatch(); - if !matched { - // Operation did not match, return to previous op - self.iter_state.backtrack = true; - // restore the previous variable state if it was modified by this operation - if let Some(previous_state) = self.iter_state.variable_state_diff[op_index].take() { - let variable_idx = previous_state.variable_idx as usize; - match previous_state.previous_state { - PreviousVariableDiff::VariableState(state) => { - self.iter_state.variable_state[previous_state.variable_idx as usize] = state; - } - PreviousVariableDiff::Decrement => { - // decrement the index in the list of candidates - match &mut self.iter_state.variable_state[variable_idx] { - VariableState::Search { offset, .. } => { - *offset -= 1; - } - VariableState::FixedByRelationshipTarget { index, .. } => { - *index -= 1; - } - _ => { - unreachable!() - } - }; - } - } - } // Returned from all operations, no more matches if self.iter_state.curr_op == 0 { return false; } + self.iter_state.backtrack = true; self.iter_state.curr_op -= 1; } else { // Operation did match, move to next op @@ -607,6 +577,14 @@ pub struct DynamicItem<'w> { entities: Vec>, } +impl core::fmt::Debug for DynamicItem<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_list() + .entries(self.entities.iter().map(|e| e.entity())) + .finish() + } +} + impl<'w, 's> Iterator for Iter<'w, 's> { type Item = DynamicItem<'w>; @@ -680,6 +658,8 @@ impl QueryPlan { #[cfg(test)] mod tests { + use std::dbg; + use std::process::Child; use bevy_ecs::prelude::*; use super::*; use crate::{ @@ -689,20 +669,43 @@ mod tests { }; #[derive(Component)] - struct Marker; + struct A; + + #[derive(Component)] + struct B; + #[derive(Component)] + struct C; + + #[derive(Component)] + #[relationship(relationship_target = R1Target)] + struct R1(#[entities] Entity); + + #[derive(Component)] + #[relationship_target(relationship = R1)] + struct R1Target(#[entities] Vec); + + #[derive(Component)] + #[relationship(relationship_target = R2Target)] + struct R2(#[entities] Entity); + + #[derive(Component)] + #[relationship_target(relationship = R2)] + struct R2Target(#[entities] Vec); + + /// Q1 -> R12 -> Q2 #[test] fn test_query_plan_basic() { let mut world = World::new(); // Correct pair let parent = world.spawn_empty().id(); - let child = world.spawn((Marker, ChildOf(parent))).id(); + let child = world.spawn((A, ChildOf(parent))).id(); world.flush(); // Build a simple plan using the builder API let mut builder = QueryPlanBuilder::new(&mut world); - builder.add_source::<&Marker, ()>(); + builder.add_source::<&A, ()>(); builder.add_relationship::(0, 1); builder.add_source::(); let plan = builder.compile(); @@ -716,46 +719,216 @@ mod tests { } /// Checks that filters in the source or target of the relationship are respected + /// Q1 -> R12 -> Q2 #[test] fn test_query_plan_single_relationship() { let mut world = World::new(); // Parent does not have the marker let parent1 = world.spawn_empty().id(); - let _ = world.spawn((Marker, ChildOf(parent1))).id(); + let _ = world.spawn((A, ChildOf(parent1))).id(); world.flush(); // Child does not have the marker - let parent2 = world.spawn(Marker).id(); + let parent2 = world.spawn(A).id(); let _ = world.spawn(ChildOf(parent2)).id(); // Both have markers but there is no relationship - let _ = world.spawn(Marker).id(); - let _ = world.spawn(Marker).id(); + let _ = world.spawn(A).id(); + let _ = world.spawn(A).id(); // Two correct pairs, (Child, Parent) and (Parent, Grandparent) - let grandparent4 = world.spawn(Marker).id(); - let parent4 = world.spawn((Marker, ChildOf(grandparent4))).id(); - let child4 = world.spawn((Marker, ChildOf(parent4))).id(); + let grandparent4 = world.spawn(A).id(); + let parent4 = world.spawn((A, ChildOf(grandparent4))).id(); + let child4 = world.spawn((A, ChildOf(parent4))).id(); // Both sources must have the Marker let mut builder = QueryPlanBuilder::new(&mut world); - builder.add_source::<&Marker, ()>(); + builder.add_source::>(); builder.add_relationship::(0, 1); - builder.add_source::>(); + builder.add_source::>(); let plan = builder.compile(); let iter = plan.query_iter(world.as_unsafe_world_cell()); let results: Vec = iter.collect(); assert_eq!(results.len(), 2); - assert_eq!(results[0].entities[0].entity(), child4); - assert_eq!(results[0].entities[1].entity(), parent4); - assert_eq!(results[1].entities[0].entity(), parent4); - assert_eq!(results[2].entities[1].entity(), grandparent4); + assert_eq!(results[0].entities[0].entity(), parent4); + assert_eq!(results[0].entities[1].entity(), grandparent4); + assert_eq!(results[1].entities[0].entity(), child4); + assert_eq!(results[1].entities[1].entity(), parent4); } -} + /// Checks that filters in the source or target of the relationship are respected + /// Q1 -> R12 -> Q2 (but expressed using relationship target) + #[test] + fn test_query_plan_single_relationship_target() { + let mut world = World::new(); + + // Parent does not have the marker + let parent1 = world.spawn_empty().id(); + let _ = world.spawn((A, ChildOf(parent1))).id(); + world.flush(); + + // Child does not have the marker + let parent2 = world.spawn(A).id(); + let _ = world.spawn(ChildOf(parent2)).id(); + + // Both have markers but there is no relationship + let _ = world.spawn(A).id(); + let _ = world.spawn(A).id(); + + // Two correct pairs, (Child, Parent) and (Parent, Grandparent) + let grandparent4 = world.spawn(A).id(); + let parent4 = world.spawn((A, ChildOf(grandparent4))).id(); + let child4 = world.spawn((A, ChildOf(parent4))).id(); -// Add test-case: Q1, R21, Q2 (find sources) -// Add test-case: Q1, R12, Q2, R13, Q3, R23 ($2 and $3 have been fixed, and then we check R23) \ No newline at end of file + // Both sources must have the Marker + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::>(); + builder.add_relationship_target::(1, 0); + builder.add_source::>(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].entities[0].entity(), parent4); + assert_eq!(results[0].entities[1].entity(), grandparent4); + assert_eq!(results[1].entities[0].entity(), child4); + assert_eq!(results[1].entities[1].entity(), parent4); + } + + /// The first variable $1 is the target of the relationship + /// Q1 -> R21 -> Q2 + #[test] + fn test_query_plan_single_relationship_reverse() { + let mut world = World::new(); + + // Parent does not have the marker + let parent1 = world.spawn_empty().id(); + let _ = world.spawn((A, ChildOf(parent1))).id(); + world.flush(); + + // Child does not have the marker + let parent2 = world.spawn(A).id(); + let _ = world.spawn(ChildOf(parent2)).id(); + + // Both have markers but there is no relationship + let _ = world.spawn(A).id(); + let _ = world.spawn(A).id(); + + // Two correct pairs, (Child, Parent) and (Parent, Grandparent) + let grandparent4 = world.spawn(A).id(); + let parent4 = world.spawn((A, ChildOf(grandparent4))).id(); + let child4 = world.spawn((A, ChildOf(parent4))).id(); + + // Both sources must have the Marker + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::>(); + builder.add_relationship::(1, 0); + builder.add_source::>(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].entities[0].entity(), parent4); + assert_eq!(results[0].entities[1].entity(), child4); + assert_eq!(results[1].entities[0].entity(), grandparent4); + assert_eq!(results[1].entities[1].entity(), parent4); + } + + /// Q1 -> R12 -> Q2 -> R23 -> Q3 + #[test] + fn test_query_plan_multi_relationship() { + let mut world = World::new(); + + // Valid triplet + let valid_3 = world.spawn(C).id(); + let valid_2 = world.spawn((B, R2(valid_3))).id(); + let valid_1 = world.spawn((A, R1(valid_2))).id(); + + // Invalid triplet: the constraint Q3 is not satisfied + let invalid_3_a = world.spawn_empty().id(); + let invalid_2_a = world.spawn((B, R2(invalid_3_a))).id(); + let _ = world.spawn((A, R1(invalid_2_a))).id(); + + // Invalid triplet: the constraint Q3 is not satisfied + let invalid_3_b = world.spawn(C).id(); + let invalid_2_b = world.spawn(R2(invalid_3_b)).id(); + let _ = world.spawn((A, R1(invalid_2_b))).id(); + + // Invalid triplet: the constraint Q1 is not satisfied + let invalid_3_c = world.spawn(C).id(); + let invalid_2_c = world.spawn((B, R2(invalid_3_c))).id(); + let _ = world.spawn(R1(invalid_2_c)).id(); + + // Invalid triplet: the constraint R23 is not satisfied + let invalid_3_d = world.spawn(C).id(); + let invalid_2_d = world.spawn(B).id(); + let _ = world.spawn((A, R1(invalid_2_d))).id(); + + // Invalid triplet: the constraint R12 is not satisfied + let invalid_3_e = world.spawn(C).id(); + let invalid_2_e = world.spawn((B, R2(invalid_3_e))).id(); + let _ = world.spawn(A).id(); + + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::<(), With>(); + builder.add_relationship::(0, 1); + builder.add_source::<(), With>(); + builder.add_relationship::(1, 2); + builder.add_source::<(), With>(); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].entities[0].entity(), valid_1); + assert_eq!(results[0].entities[1].entity(), valid_2); + assert_eq!(results[0].entities[2].entity(), valid_3); + } + + + /// Q1 -> R12 -> Q2 -> R23 -> Q3 -> R13 + /// At the end, we backtrack through R13 but we don't modify the fixed state of $1 and $3 since + /// they were written in previous ops + #[test] + fn test_query_plan_multi_relationship_fixed_backtrack() { + let mut world = World::new(); + + // Valid triplet + let valid_3 = world.spawn(C).id(); + let valid_2 = world.spawn((B, R1(valid_3))).id(); + let valid_1 = world.spawn((A, ChildOf(valid_2), R2(valid_3))).id(); + + // Invalid triplet: the last constraint (R2) is not satisfied + let invalid_3 = world.spawn(C).id(); + let invalid_2 = world.spawn((B, R1(valid_3))).id(); + let invalid_1 = world.spawn((A, ChildOf(valid_2))).id(); + + world.flush(); + + // Both sources must have the Marker + let mut builder = QueryPlanBuilder::new(&mut world); + builder.add_source::<(), With>(); + builder.add_relationship::(0, 1); + builder.add_source::<(), With>(); + builder.add_relationship::(1, 2); + builder.add_source::<(), With>(); + builder.add_relationship::(0, 2); + let plan = builder.compile(); + + let iter = plan.query_iter(world.as_unsafe_world_cell()); + let results: Vec = iter.collect(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].entities[0].entity(), valid_1); + assert_eq!(results[0].entities[1].entity(), valid_2); + assert_eq!(results[0].entities[2].entity(), valid_3); + } +} \ No newline at end of file From 8b8f73ba73da40d58091e1ecd32052f1858fea68 Mon Sep 17 00:00:00 2001 From: Charles Bournhonesque Date: Sun, 26 Oct 2025 18:54:22 -0400 Subject: [PATCH 6/7] start type-erasing query functions --- crates/bevy_ecs/src/query/plan.rs | 143 +++++++++++++++++++++++------- 1 file changed, 112 insertions(+), 31 deletions(-) diff --git a/crates/bevy_ecs/src/query/plan.rs b/crates/bevy_ecs/src/query/plan.rs index d10acdc3ce112..916c1d7410b81 100644 --- a/crates/bevy_ecs/src/query/plan.rs +++ b/crates/bevy_ecs/src/query/plan.rs @@ -7,13 +7,15 @@ use crate::{ }; use alloc::vec::Vec; use alloc::vec; +use std::ptr::NonNull; use fixedbitset::FixedBitSet; use bevy_ecs::archetype::Archetype; use bevy_ecs::prelude::{ContainsEntity, QueryBuilder, World}; use bevy_ecs::query::{QueryData, QueryFilter}; use bevy_ecs::relationship::{Relationship, RelationshipTarget}; -use bevy_ecs::storage::{SparseSetIndex, TableId}; +use bevy_ecs::storage::{SparseSetIndex, TableId, TableRow}; use bevy_ecs::world::unsafe_world_cell::UnsafeEntityCell; +use bevy_ptr::{OwningPtr, Ptr, PtrMut}; use crate::relationship::RelationshipAccessor; /// Represents a single source in a multi-source query. @@ -24,21 +26,53 @@ pub enum QueryOperation { Source(QuerySource), } +type MatchesArchetypeFn = fn(archetype: &Archetype, component_access: &FilteredAccess, fetch_state: Ptr, filter_state: Ptr) -> bool; + +/// Check if the archetype should be considered +pub unsafe fn matches_archetype(archetype: &Archetype, component_access: &FilteredAccess, fetch_state: Ptr, filter_state: Ptr) -> bool { + let fetch_state = unsafe { fetch_state.deref::() }; + let filter_state = unsafe { filter_state.deref::() }; + // let fetch_state = unsafe { fetch_state.as_ptr().cast::().as_ref().unwrap_unchecked() }; + // let filter_state = unsafe { filter_state.as_ptr().cast::().as_ref().unwrap_unchecked() }; + D::matches_component_set(fetch_state, &|id| archetype.contains(id)) + && F::matches_component_set(filter_state, &|id| archetype.contains(id)) + && QuerySource::matches_component_set(component_access, &|id| archetype.contains(id)) + } + +type FilterFetchFn = fn(state: Ptr, fetch: PtrMut, entity: Entity, table_row: TableRow) -> bool; + +/// Returns true if the provided [`Entity`] and [`TableRow`] should be included in the query results. +/// If false, the entity will be skipped. +pub unsafe fn filter_fetch(state: Ptr, fetch: PtrMut, entity: Entity, table_row: TableRow) -> bool { + let state = unsafe { state.deref::() }; + let fetch = unsafe { fetch.deref_mut::>() }; + unsafe { F::filter_fetch(state, fetch, entity, table_row) } +} + #[derive(Debug, Clone)] pub struct QuerySource { access: FilteredAccess, + fetch_state: NonNull, + filter_state: NonNull, + matches_archetype_fn: MatchesArchetypeFn, + filter_fetch_fn: FilterFetchFn, variable_idx: u8, } impl QuerySource { pub fn matches(&self, archetype: &Archetype) -> bool { - self.matches_component_set(&|id| archetype.contains(id)) + (self.matches_archetype_fn)( + archetype, + &self.access, + unsafe { Ptr::new(self.fetch_state) }, + unsafe { Ptr::new(self.filter_state) }, + ) } - /// Returns `true` if this query matches a set of components. Otherwise, returns `false`. - pub fn matches_component_set(&self, set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { - self.access.filter_sets.iter().any(|set| { + /// Returns `true` if this access matches a set of components. Otherwise, returns `false`. + pub fn matches_component_set(access: &FilteredAccess, set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { + access.filter_sets.iter().any(|set| { set.with .ones() .all(|index| set_contains_id(ComponentId::get_sparse_set_index(index))) @@ -119,10 +153,29 @@ impl<'w> QueryPlanBuilder<'w> { } } - pub fn add_source_from_builder(&mut self, f: impl Fn(QueryBuilder) -> QueryBuilder) -> &mut Self { + pub fn add_source_from_builder(&mut self, f: impl Fn(QueryBuilder) -> QueryBuilder) -> &mut Self { let query_builder = QueryBuilder::new(&mut self.world); - let builder = f(query_builder); - self.plan.add_source(builder.access().clone()); + let mut builder = f(query_builder); + let builder = builder.transmute_filtered::(); + + let mut fetch_state = D::init_state(builder.world_mut()); + let filter_state = F::init_state(builder.world_mut()); + + let mut component_access = FilteredAccess::default(); + D::update_component_access(&fetch_state, &mut component_access); + D::provide_extra_access( + &mut fetch_state, + component_access.access_mut(), + builder.access().access(), + ); + + let access = builder.access().clone(); + let fetch_state = NonNull::from(&fetch_state).cast::(); + let filter_state = NonNull::from(&filter_state).cast::(); + + let matches_archetype_fn: MatchesArchetypeFn = unsafe { core::mem::transmute(matches_archetype::) }; + let filter_fetch_fn: FilterFetchFn = unsafe { core::mem::transmute(filter_fetch::) }; + self.plan.add_source(access, fetch_state, filter_state, matches_archetype_fn, filter_fetch_fn); self } @@ -143,7 +196,13 @@ impl<'w> QueryPlanBuilder<'w> { // properly considered in a global "cross-query" context (both within systems and across systems). access.extend(&filter_access); - self.plan.add_source(access); + let fetch_state = NonNull::from(&fetch_state).cast::(); + let filter_state = NonNull::from(&filter_state).cast::(); + + let matches_archetype_fn: MatchesArchetypeFn = unsafe { core::mem::transmute(matches_archetype::) }; + let filter_fetch_fn: FilterFetchFn = unsafe { core::mem::transmute(filter_fetch::) }; + + self.plan.add_source(access, fetch_state, filter_state, matches_archetype_fn, filter_fetch_fn); self } @@ -417,26 +476,11 @@ impl<'w, 's> Iter<'w, 's> { // the first time we evaluate the query, we set the list of potential tables let tables = &storage_state.tables; assert_eq!(tables.len(), 0); - // if there are required components, we can optimize by only iterating through archetypes - // that contain at least one of the required components - let potential_archetypes = source - .access - .with_filters() - .filter_map(|component_id| { - self.world - .archetypes() - .component_index() - .get(&component_id) - .map(|index| index.keys()) - }) - // select the component with the fewest archetypes - .min_by_key(ExactSizeIterator::len); + let mut matching_tables = Vec::new(); let mut current_len = 0; - if let Some(archetypes) = potential_archetypes { - for archetype_id in archetypes { - // SAFETY: get_potential_archetypes only returns archetype ids that are valid for the world - let archetype = &self.world.archetypes()[*archetype_id]; + if source.access.required.is_empty() { + for archetype in self.world.archetypes().iter() { // NOTE: you could have empty archetypes even if the table is not empty // and you could have non-empty archetypes with empty tables (e.g. when the archetype only has sparse sets) let table_id = archetype.table_id(); @@ -449,14 +493,48 @@ impl<'w, 's> Iter<'w, 's> { matching_tables.push(archetype.table_id()) } } + } else { + // if there are required components, we can optimize by only iterating through archetypes + // that contain at least one of the required components + let potential_archetypes = source + .access + .required + .ones() + .filter_map(|idx| { + let component_id = ComponentId::get_sparse_set_index(idx); + self.world + .archetypes() + .component_index() + .get(&component_id) + .map(|index| index.keys()) + }) + // select the component with the fewest archetypes + .min_by_key(ExactSizeIterator::len); + if let Some(archetypes) = potential_archetypes { + for archetype_id in archetypes { + // SAFETY: get_potential_archetypes only returns archetype ids that are valid for the world + let archetype = &self.world.archetypes()[*archetype_id]; + // NOTE: you could have empty archetypes even if the table is not empty + // and you could have non-empty archetypes with empty tables (e.g. when the archetype only has sparse sets) + let table_id = archetype.table_id(); + let table = unsafe { self.world.world().storages().tables.get(table_id).unwrap_unchecked() }; + // skip empty tables + if !table.is_empty() && source.matches(archetype) { + if current_len == 0 { + current_len = table.entity_count(); + } + matching_tables.push(archetype.table_id()) + } + } + } } if matching_tables.first().is_none() { return false; } - self.iter_state.storages_state[variable_idx] = StorageState { tables: Cow::Owned(matching_tables) }; + // TODO: need to iterate here until we find a good candidate! check with filter_fetch self.set_variable_state(variable_idx as u8, VariableState::Search { offset: 0, storage_idx: 0, @@ -617,10 +695,14 @@ impl QueryPlan { } /// Add a term to the query plan. - pub(crate) fn add_source(&mut self, access: FilteredAccess) -> usize { + pub(crate) fn add_source(&mut self, access: FilteredAccess, fetch_state: NonNull, filter_state: NonNull, matches_archetype_fn: MatchesArchetypeFn, filter_fetch_fn: FilterFetchFn) -> usize { let op_index = self.operations.len(); self.operations.push(QueryOperation::Source(QuerySource { access, + fetch_state, + filter_state, + matches_archetype_fn, + filter_fetch_fn, variable_idx: self.num_variables, })); self.num_variables += 1; @@ -658,8 +740,7 @@ impl QueryPlan { #[cfg(test)] mod tests { - use std::dbg; - use std::process::Child; + #![allow(unused_variables)] use bevy_ecs::prelude::*; use super::*; use crate::{ From 323c0c98eaa328e0ceec8a3d1aa1e492edc983ef Mon Sep 17 00:00:00 2001 From: Charles Bournhonesque Date: Sun, 26 Oct 2025 23:03:40 -0400 Subject: [PATCH 7/7] Use type-erased Queries --- crates/bevy_ecs/src/query/plan.rs | 108 ++++++++++++++++-------------- 1 file changed, 56 insertions(+), 52 deletions(-) diff --git a/crates/bevy_ecs/src/query/plan.rs b/crates/bevy_ecs/src/query/plan.rs index 916c1d7410b81..30b43849f3ce2 100644 --- a/crates/bevy_ecs/src/query/plan.rs +++ b/crates/bevy_ecs/src/query/plan.rs @@ -1,3 +1,4 @@ +#![allow(dead_code)] use alloc::borrow::Cow; use crate::{ component::ComponentId, @@ -5,41 +6,41 @@ use crate::{ query::FilteredAccess, world::unsafe_world_cell::UnsafeWorldCell, }; +use alloc::boxed::Box; use alloc::vec::Vec; use alloc::vec; +use std::any::Any; use std::ptr::NonNull; use fixedbitset::FixedBitSet; use bevy_ecs::archetype::Archetype; use bevy_ecs::prelude::{ContainsEntity, QueryBuilder, World}; -use bevy_ecs::query::{QueryData, QueryFilter}; +use bevy_ecs::query::{QueryData, QueryFilter, WorldQuery}; use bevy_ecs::relationship::{Relationship, RelationshipTarget}; use bevy_ecs::storage::{SparseSetIndex, TableId, TableRow}; use bevy_ecs::world::unsafe_world_cell::UnsafeEntityCell; -use bevy_ptr::{OwningPtr, Ptr, PtrMut}; +use bevy_ptr::{Ptr, PtrMut}; use crate::relationship::RelationshipAccessor; /// Represents a single source in a multi-source query. /// Each term has its own ComponentAccess requirements. -#[derive(Debug, Clone)] +#[derive(Debug)] pub enum QueryOperation { Relationship(QueryRelationship), Source(QuerySource), } -type MatchesArchetypeFn = fn(archetype: &Archetype, component_access: &FilteredAccess, fetch_state: Ptr, filter_state: Ptr) -> bool; +type MatchesArchetypeFn = unsafe fn(archetype: &Archetype, component_access: &FilteredAccess, fetch_state: Ptr, filter_state: Ptr) -> bool; /// Check if the archetype should be considered pub unsafe fn matches_archetype(archetype: &Archetype, component_access: &FilteredAccess, fetch_state: Ptr, filter_state: Ptr) -> bool { let fetch_state = unsafe { fetch_state.deref::() }; let filter_state = unsafe { filter_state.deref::() }; - // let fetch_state = unsafe { fetch_state.as_ptr().cast::().as_ref().unwrap_unchecked() }; - // let filter_state = unsafe { filter_state.as_ptr().cast::().as_ref().unwrap_unchecked() }; D::matches_component_set(fetch_state, &|id| archetype.contains(id)) - && F::matches_component_set(filter_state, &|id| archetype.contains(id)) - && QuerySource::matches_component_set(component_access, &|id| archetype.contains(id)) - } + && F::matches_component_set(filter_state, &|id| archetype.contains(id)) + && QuerySource::matches_component_set(component_access, &|id| archetype.contains(id)) +} -type FilterFetchFn = fn(state: Ptr, fetch: PtrMut, entity: Entity, table_row: TableRow) -> bool; +type FilterFetchFn = unsafe fn(state: Ptr, fetch: PtrMut, entity: Entity, table_row: TableRow) -> bool; /// Returns true if the provided [`Entity`] and [`TableRow`] should be included in the query results. /// If false, the entity will be skipped. @@ -49,11 +50,11 @@ pub unsafe fn filter_fetch(state: Ptr, fetch: PtrMut, entity: En unsafe { F::filter_fetch(state, fetch, entity, table_row) } } -#[derive(Debug, Clone)] +#[derive(Debug)] pub struct QuerySource { access: FilteredAccess, - fetch_state: NonNull, - filter_state: NonNull, + fetch_state: Box, + filter_state: Box, matches_archetype_fn: MatchesArchetypeFn, filter_fetch_fn: FilterFetchFn, variable_idx: u8, @@ -62,12 +63,14 @@ pub struct QuerySource { impl QuerySource { pub fn matches(&self, archetype: &Archetype) -> bool { - (self.matches_archetype_fn)( + let fetch_state = unsafe { NonNull::new_unchecked(self.fetch_state.as_ref() as *const dyn Any as *mut dyn Any) }; + let filter_state = unsafe { NonNull::new_unchecked(self.filter_state.as_ref() as *const dyn Any as *mut dyn Any) }; + unsafe { (self.matches_archetype_fn)( archetype, &self.access, - unsafe { Ptr::new(self.fetch_state) }, - unsafe { Ptr::new(self.filter_state) }, - ) + Ptr::new(fetch_state.cast::()), + Ptr::new(filter_state.cast::()), + ) } } /// Returns `true` if this access matches a set of components. Otherwise, returns `false`. @@ -153,7 +156,9 @@ impl<'w> QueryPlanBuilder<'w> { } } - pub fn add_source_from_builder(&mut self, f: impl Fn(QueryBuilder) -> QueryBuilder) -> &mut Self { + pub fn add_source_from_builder(&mut self, f: impl Fn(QueryBuilder) -> QueryBuilder) -> &mut Self + where ::State: 'static, + ::State: 'static, { let query_builder = QueryBuilder::new(&mut self.world); let mut builder = f(query_builder); let builder = builder.transmute_filtered::(); @@ -170,16 +175,15 @@ impl<'w> QueryPlanBuilder<'w> { ); let access = builder.access().clone(); - let fetch_state = NonNull::from(&fetch_state).cast::(); - let filter_state = NonNull::from(&filter_state).cast::(); - - let matches_archetype_fn: MatchesArchetypeFn = unsafe { core::mem::transmute(matches_archetype::) }; - let filter_fetch_fn: FilterFetchFn = unsafe { core::mem::transmute(filter_fetch::) }; - self.plan.add_source(access, fetch_state, filter_state, matches_archetype_fn, filter_fetch_fn); + let matches_archetype_fn: MatchesArchetypeFn = matches_archetype::; + let filter_fetch_fn: FilterFetchFn = filter_fetch::; + self.plan.add_source(access, Box::new(fetch_state), Box::new(filter_state), matches_archetype_fn, filter_fetch_fn); self } - pub fn add_source(&mut self) -> &mut Self { + pub fn add_source(&mut self) -> &mut Self + where ::State: 'static, + ::State: 'static, { let fetch_state = D::init_state(&mut self.world); let filter_state = F::init_state(&mut self.world); @@ -196,13 +200,9 @@ impl<'w> QueryPlanBuilder<'w> { // properly considered in a global "cross-query" context (both within systems and across systems). access.extend(&filter_access); - let fetch_state = NonNull::from(&fetch_state).cast::(); - let filter_state = NonNull::from(&filter_state).cast::(); - - let matches_archetype_fn: MatchesArchetypeFn = unsafe { core::mem::transmute(matches_archetype::) }; - let filter_fetch_fn: FilterFetchFn = unsafe { core::mem::transmute(filter_fetch::) }; - - self.plan.add_source(access, fetch_state, filter_state, matches_archetype_fn, filter_fetch_fn); + let matches_archetype_fn: MatchesArchetypeFn = matches_archetype::; + let filter_fetch_fn: FilterFetchFn = filter_fetch::; + self.plan.add_source(access, Box::new(fetch_state), Box::new(filter_state), matches_archetype_fn, filter_fetch_fn); self } @@ -283,7 +283,7 @@ pub enum QueryPlanError { /// A dynamic query plan that describes how to match multiple entities /// connected through relationships. -#[derive(Debug, Default, Clone)] +#[derive(Debug, Default)] pub struct QueryPlan { /// All operations in this query. pub operations: Vec, @@ -295,7 +295,7 @@ impl QueryPlan { pub fn query_iter<'w, 's>(&'s self, world: UnsafeWorldCell<'w>) -> Iter<'w, 's> { Iter { world, - query_state: Cow::Borrowed(self), + query_state: &self, iter_state: IterState::new(self), } } @@ -373,7 +373,7 @@ impl<'w> IterState<'w> { /// Iterator that iterates through a dynamic query plan pub struct Iter<'w, 's> { world: UnsafeWorldCell<'w>, - query_state: Cow<'s, QueryPlan>, + query_state: &'s QueryPlan, iter_state: IterState<'s>, } @@ -586,6 +586,7 @@ impl<'w, 's> Iter<'w, 's> { if *storage_idx >= storage_state.tables.len() as u32 { return false } + // TODO: apply filter_fetch // either beginning of the iteration, or finished processing a table, so skip to the next if (*offset + 1) == *current_len { // go to next table @@ -695,7 +696,8 @@ impl QueryPlan { } /// Add a term to the query plan. - pub(crate) fn add_source(&mut self, access: FilteredAccess, fetch_state: NonNull, filter_state: NonNull, matches_archetype_fn: MatchesArchetypeFn, filter_fetch_fn: FilterFetchFn) -> usize { + pub(crate) fn add_source(&mut self, access: FilteredAccess, fetch_state: Box, filter_state: Box, matches_archetype_fn: MatchesArchetypeFn, filter_fetch_fn: FilterFetchFn) -> + usize { let op_index = self.operations.len(); self.operations.push(QueryOperation::Source(QuerySource { access, @@ -784,11 +786,13 @@ mod tests { let child = world.spawn((A, ChildOf(parent))).id(); world.flush(); + let ptr_a = world.register_component::(); + // Build a simple plan using the builder API let mut builder = QueryPlanBuilder::new(&mut world); - builder.add_source::<&A, ()>(); + builder.add_source::<(), With>(); builder.add_relationship::(0, 1); - builder.add_source::(); + builder.add_source::<(), ()>(); let plan = builder.compile(); let iter = plan.query_iter(world.as_unsafe_world_cell()); @@ -825,19 +829,19 @@ mod tests { // Both sources must have the Marker let mut builder = QueryPlanBuilder::new(&mut world); - builder.add_source::>(); + builder.add_source::<(), With>(); builder.add_relationship::(0, 1); - builder.add_source::>(); + builder.add_source::<(), With>(); let plan = builder.compile(); let iter = plan.query_iter(world.as_unsafe_world_cell()); let results: Vec = iter.collect(); assert_eq!(results.len(), 2); - assert_eq!(results[0].entities[0].entity(), parent4); - assert_eq!(results[0].entities[1].entity(), grandparent4); - assert_eq!(results[1].entities[0].entity(), child4); - assert_eq!(results[1].entities[1].entity(), parent4); + assert_eq!(results[1].entities[0].entity(), parent4); + assert_eq!(results[1].entities[1].entity(), grandparent4); + assert_eq!(results[0].entities[0].entity(), child4); + assert_eq!(results[0].entities[1].entity(), parent4); } /// Checks that filters in the source or target of the relationship are respected @@ -875,10 +879,10 @@ mod tests { let results: Vec = iter.collect(); assert_eq!(results.len(), 2); - assert_eq!(results[0].entities[0].entity(), parent4); - assert_eq!(results[0].entities[1].entity(), grandparent4); - assert_eq!(results[1].entities[0].entity(), child4); - assert_eq!(results[1].entities[1].entity(), parent4); + assert_eq!(results[1].entities[0].entity(), parent4); + assert_eq!(results[1].entities[1].entity(), grandparent4); + assert_eq!(results[0].entities[0].entity(), child4); + assert_eq!(results[0].entities[1].entity(), parent4); } /// The first variable $1 is the target of the relationship @@ -916,10 +920,10 @@ mod tests { let results: Vec = iter.collect(); assert_eq!(results.len(), 2); - assert_eq!(results[0].entities[0].entity(), parent4); - assert_eq!(results[0].entities[1].entity(), child4); - assert_eq!(results[1].entities[0].entity(), grandparent4); - assert_eq!(results[1].entities[1].entity(), parent4); + assert_eq!(results[1].entities[0].entity(), parent4); + assert_eq!(results[1].entities[1].entity(), child4); + assert_eq!(results[0].entities[0].entity(), grandparent4); + assert_eq!(results[0].entities[1].entity(), parent4); } /// Q1 -> R12 -> Q2 -> R23 -> Q3