Skip to content

Commit

Permalink
Dynamically sized queries
Browse files Browse the repository at this point in the history
  • Loading branch information
Suficio committed Oct 27, 2022
1 parent ba27ee5 commit 2ceac66
Show file tree
Hide file tree
Showing 3 changed files with 230 additions and 164 deletions.
35 changes: 23 additions & 12 deletions src/query/mod.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
mod ptr;
mod vec;

pub use ptr::{ComponentPtr, ComponentPtrDense};
pub use ptr::ComponentPtr;
pub use vec::VecPtr;

#[allow(unreachable_code)]
Expand Down Expand Up @@ -35,13 +35,15 @@ mod tests {
let component_id = world.init_component::<A>();

let mut query = unsafe {
QueryState::<&ComponentPtrDense, ()>::new_with_state(&mut world, component_id, ())
QueryState::<&ComponentPtr, ()>::new_with_state(&mut world, component_id, ())
};
let count = query.iter(&mut world).count();
assert_eq!(count, 0);

// Check dynamic VecPtr variations

let mut query = unsafe {
QueryState::<VecPtr<&ComponentPtrDense>, ()>::new_with_state(
QueryState::<VecPtr<&ComponentPtr>, ()>::new_with_state(
&mut world,
vec![component_id],
(),
Expand All @@ -62,7 +64,7 @@ mod tests {
world.spawn(A);

let mut query =
unsafe { QueryState::<&ComponentPtrDense, ()>::new_with_state(&mut world, a, ()) };
unsafe { QueryState::<&ComponentPtr, ()>::new_with_state(&mut world, a, ()) };
let count = query.iter(&mut world).count();
assert_eq!(count, 2);

Expand All @@ -72,16 +74,25 @@ mod tests {
assert_eq!(count, 1);

let mut query = unsafe {
QueryState::<(&ComponentPtrDense, &ComponentPtr), ()>::new_with_state(
&mut world,
(a, b),
(),
)
QueryState::<(&ComponentPtr, &ComponentPtr), ()>::new_with_state(&mut world, (a, b), ())
};
let count = query.iter(&mut world).count();
assert_eq!(count, 1);

// Check dynamic VecPtr variations

let mut query = unsafe {
QueryState::<VecPtr<&ComponentPtr>, ()>::new_with_state(&mut world, vec![a], ())
};
let count = query.iter(&mut world).count();
assert_eq!(count, 2);

let mut query = unsafe {
QueryState::<VecPtr<&ComponentPtr>, ()>::new_with_state(&mut world, vec![b], ())
};
let count = query.iter(&mut world).count();
assert_eq!(count, 1);

// Pick VecPtr<&ComponentPtr> due to presence of &ComponentPtr in query
let mut query = unsafe {
QueryState::<VecPtr<&ComponentPtr>, ()>::new_with_state(&mut world, vec![a, b], ())
};
Expand All @@ -101,7 +112,7 @@ mod tests {
world.spawn(Data("Hello, World!".to_string()));

let mut query = unsafe {
QueryState::<&ComponentPtrDense, ()>::new_with_state(&mut world, component_id, ())
QueryState::<&ComponentPtr, ()>::new_with_state(&mut world, component_id, ())
};

for data in query.iter(&mut world) {
Expand Down Expand Up @@ -129,7 +140,7 @@ mod tests {
world.spawn(Data("Hello, World!".to_string()));

let mut query = unsafe {
QueryState::<&ComponentPtrDense, ()>::new_with_state(&mut world, component_id, ())
QueryState::<&ComponentPtr, ()>::new_with_state(&mut world, component_id, ())
};

let res = query.iter_mut(&mut world).collect::<Vec<_>>();
Expand Down
209 changes: 57 additions & 152 deletions src/query/ptr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,139 +10,18 @@ use bevy::{
ptr::{Ptr, ThinSlicePtr},
};

/// Type used to query dense components
pub struct ComponentPtrDense;
/// Type used to query non-dense components
pub struct ComponentPtr;

#[doc(hidden)]
pub struct ReadFetchDense<'w> {
pub struct ReadFetchSparse<'w> {
storage_type: StorageType,
component_size: usize,

// T::Storage = TableStorage
table_components: Option<Ptr<'w>>,
entity_table_rows: Option<ThinSlicePtr<'w, usize>>,
}

impl<'w> WorldQueryGats<'w> for &ComponentPtrDense {
type Item = Ptr<'w>;
type Fetch = ReadFetchDense<'w>;
}

unsafe impl ReadOnlyWorldQuery for &ComponentPtrDense {}

unsafe impl WorldQuery for &ComponentPtrDense {
type ReadOnly = Self;
type State = ComponentId;

fn shrink<'wlong: 'wshort, 'wshort>(item: Ptr<'wlong>) -> Ptr<'wshort> {
item
}

const IS_DENSE: bool = true;
const IS_ARCHETYPAL: bool = true;

unsafe fn init_fetch<'w>(
world: &'w World,
&component_id: &ComponentId,
_last_change_tick: u32,
_change_tick: u32,
) -> ReadFetchDense<'w> {
let component_info = world.components().get_info(component_id).unwrap();

debug_assert_eq!(component_info.storage_type(), StorageType::Table);

ReadFetchDense {
component_size: component_info.layout().size(),
table_components: None,
entity_table_rows: None,
}
}

#[inline]
unsafe fn set_archetype<'w>(
fetch: &mut ReadFetchDense<'w>,
&component_id: &ComponentId,
archetype: &'w Archetype,
tables: &'w Tables,
) {
fetch.entity_table_rows = Some(archetype.entity_table_rows().into());
let column = tables[archetype.table_id()]
.get_column(component_id)
.unwrap();
fetch.table_components = Some(column.get_data_ptr());
}

#[inline]
unsafe fn set_table<'w>(
fetch: &mut ReadFetchDense<'w>,
&component_id: &ComponentId,
table: &'w Table,
) {
fetch.table_components = Some(table.get_column(component_id).unwrap().get_data_ptr());
}

#[inline]
unsafe fn archetype_fetch<'w>(
fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
archetype_index: usize,
) -> <Self as WorldQueryGats<'w>>::Item {
let (entity_table_rows, table_components) = fetch
.entity_table_rows
.zip(fetch.table_components)
.unwrap_or_else(|| debug_checked_unreachable());

let table_row = *entity_table_rows.get(archetype_index);
table_components.byte_add(table_row * fetch.component_size)
}

#[inline]
unsafe fn table_fetch<'w>(
fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
table_row: usize,
) -> <Self as WorldQueryGats<'w>>::Item {
let components = fetch
.table_components
.unwrap_or_else(|| debug_checked_unreachable());
components.byte_add(table_row * fetch.component_size)
}

fn update_component_access(
&component_id: &ComponentId,
access: &mut FilteredAccess<ComponentId>,
) {
assert!(
!access.access().has_write(component_id),
"Read access to component with id: {:?} conflicts with a previous access in this query. Shared access cannot coincide with exclusive access.",
component_id
);
access.add_read(component_id);
}

fn update_archetype_component_access(
&component_id: &ComponentId,
archetype: &Archetype,
access: &mut Access<ArchetypeComponentId>,
) {
if let Some(archetype_component_id) = archetype.get_archetype_component_id(component_id) {
access.add_read(archetype_component_id);
}
}

fn init_state(_world: &mut World) -> Self::State {
panic!("Dynamic queries can only be initialized through QueryState::new_with_state");
}

fn matches_component_set(
&component_id: &ComponentId,
set_contains_id: &impl Fn(ComponentId) -> bool,
) -> bool {
set_contains_id(component_id)
}
}

/// Type used to query non-dense components
pub struct ComponentPtr;

#[doc(hidden)]
pub struct ReadFetchSparse<'w> {
// T::Storage = SparseStorage
entities: Option<ThinSlicePtr<'w, Entity>>,
sparse_set: Option<&'w ComponentSparseSet>,
Expand All @@ -169,13 +48,16 @@ unsafe impl WorldQuery for &ComponentPtr {
_last_change_tick: u32,
_change_tick: u32,
) -> <Self as WorldQueryGats<'w>>::Fetch {
let storage_type = world
.components()
.get_info(component_id)
.unwrap()
.storage_type();
let component_info = world.components().get_info(component_id).unwrap();
let storage_type = component_info.storage_type();

ReadFetchSparse {
storage_type,
component_size: component_info.layout().size(),

table_components: None,
entity_table_rows: None,

entities: None,
sparse_set: (storage_type == StorageType::SparseSet)
.then(|| world.storages().sparse_sets.get(component_id).unwrap()),
Expand All @@ -188,46 +70,69 @@ unsafe impl WorldQuery for &ComponentPtr {
#[inline]
unsafe fn set_archetype<'w>(
fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
_component_id: &ComponentId,
&component_id: &ComponentId,
archetype: &'w Archetype,
_tables: &'w Tables,
tables: &'w Tables,
) {
fetch.entities = Some(archetype.entities().into())
match fetch.storage_type {
StorageType::Table => {
fetch.entity_table_rows = Some(archetype.entity_table_rows().into());
let column = tables[archetype.table_id()]
.get_column(component_id)
.unwrap();
fetch.table_components = Some(column.get_data_ptr());
}
StorageType::SparseSet => fetch.entities = Some(archetype.entities().into()),
}
}

#[inline]
unsafe fn set_table<'w>(
_fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
_component_id: &ComponentId,
_table: &'w Table,
fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
&component_id: &ComponentId,
table: &'w Table,
) {
panic!("hi");
debug_checked_unreachable();
fetch.table_components = Some(table.get_column(component_id).unwrap().get_data_ptr());
}

#[inline]
unsafe fn archetype_fetch<'w>(
fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
archetype_index: usize,
) -> <Self as WorldQueryGats<'w>>::Item {
let (entities, sparse_set) = fetch
.entities
.zip(fetch.sparse_set)
.unwrap_or_else(|| debug_checked_unreachable());

let entity = *entities.get(archetype_index);
sparse_set
.get(entity)
.unwrap_or_else(|| debug_checked_unreachable())
match fetch.storage_type {
StorageType::Table => {
let (entity_table_rows, table_components) = fetch
.entity_table_rows
.zip(fetch.table_components)
.unwrap_or_else(|| debug_checked_unreachable());

let table_row = *entity_table_rows.get(archetype_index);
table_components.byte_add(table_row * fetch.component_size)
}
StorageType::SparseSet => {
let (entities, sparse_set) = fetch
.entities
.zip(fetch.sparse_set)
.unwrap_or_else(|| debug_checked_unreachable());

let entity = *entities.get(archetype_index);
sparse_set
.get(entity)
.unwrap_or_else(|| debug_checked_unreachable())
}
}
}

#[inline]
unsafe fn table_fetch<'w>(
_fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
_table_row: usize,
fetch: &mut <Self as WorldQueryGats<'w>>::Fetch,
table_row: usize,
) -> <Self as WorldQueryGats<'w>>::Item {
panic!("hi");
debug_checked_unreachable();
let components = fetch
.table_components
.unwrap_or_else(|| debug_checked_unreachable());
components.byte_add(table_row * fetch.component_size)
}

fn update_component_access(
Expand Down
Loading

0 comments on commit 2ceac66

Please sign in to comment.