Add a method iter_combinations on query to iterate over combinations of query results (#1763)

Related to [discussion on discord](https://discord.com/channels/691052431525675048/742569353878437978/824731187724681289)

With const generics, it is now possible to write generic iterator over multiple entities at once.

This enables patterns of query iterations like

```rust
for [e1, e2, e3] in query.iter_combinations() {
   // do something with relation of all three entities
}
```

The compiler is able to infer the correct iterator for given size of array, so either of those work
```rust
for [e1, e2] in query.iter_combinations()  { ... }
for [e1, e2, e3] in query.iter_combinations()  { ... }
```

This feature can be very useful for systems like collision detection.

When you ask for permutations of size K of N entities:
- if K == N, you get one result of all entities
- if K < N, you get all possible subsets of N with size K, without repetition
- if K > N, the result set is empty (no permutation of size K exist)

Co-authored-by: Carter Anderson <mcanders1@gmail.com>
This commit is contained in:
Paweł Grabarz 2021-05-17 23:33:47 +00:00
parent cb98d31b27
commit a81fb7aa7e
7 changed files with 799 additions and 9 deletions

View file

@ -267,6 +267,10 @@ path = "examples/ecs/fixed_timestep.rs"
name = "hierarchy"
path = "examples/ecs/hierarchy.rs"
[[example]]
name = "iter_combinations"
path = "examples/ecs/iter_combinations.rs"
[[example]]
name = "parallel_query"
path = "examples/ecs/parallel_query.rs"

View file

@ -279,6 +279,18 @@ pub struct ReadFetch<T> {
sparse_set: *const ComponentSparseSet,
}
impl<T> Clone for ReadFetch<T> {
fn clone(&self) -> Self {
Self {
storage_type: self.storage_type,
table_components: self.table_components,
entity_table_rows: self.entity_table_rows,
entities: self.entities,
sparse_set: self.sparse_set,
}
}
}
/// SAFETY: access is read only
unsafe impl<T> ReadOnlyFetch for ReadFetch<T> {}
@ -382,6 +394,21 @@ pub struct WriteFetch<T> {
change_tick: u32,
}
impl<T> Clone for WriteFetch<T> {
fn clone(&self) -> Self {
Self {
storage_type: self.storage_type,
table_components: self.table_components,
table_ticks: self.table_ticks,
entities: self.entities,
entity_table_rows: self.entity_table_rows,
sparse_set: self.sparse_set,
last_change_tick: self.last_change_tick,
change_tick: self.change_tick,
}
}
}
/// The [`FetchState`] of `&mut T`.
pub struct WriteState<T> {
component_id: ComponentId,

View file

@ -1,9 +1,10 @@
use crate::{
archetype::{ArchetypeId, Archetypes},
query::{Fetch, FilterFetch, QueryState, WorldQuery},
query::{Fetch, FilterFetch, QueryState, ReadOnlyFetch, WorldQuery},
storage::{TableId, Tables},
world::World,
};
use std::mem::MaybeUninit;
/// An [`Iterator`] over query results of a [`Query`](crate::system::Query).
///
@ -21,15 +22,20 @@ where
archetype_id_iter: std::slice::Iter<'s, ArchetypeId>,
fetch: Q::Fetch,
filter: F::Fetch,
is_dense: bool,
current_len: usize,
current_index: usize,
is_dense: bool,
}
impl<'w, 's, Q: WorldQuery, F: WorldQuery> QueryIter<'w, 's, Q, F>
where
F::Fetch: FilterFetch,
{
/// # Safety
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
/// This does not validate that `world.id()` matches `query_state.world_id`. Calling this on a `world`
/// with a mismatched WorldId is unsound.
pub(crate) unsafe fn new(
world: &'w World,
query_state: &'s QueryState<Q, F>,
@ -48,14 +54,15 @@ where
last_change_tick,
change_tick,
);
QueryIter {
is_dense: fetch.is_dense() && filter.is_dense(),
world,
query_state,
fetch,
filter,
tables: &world.storages().tables,
archetypes: &world.archetypes,
is_dense: fetch.is_dense() && filter.is_dense(),
fetch,
filter,
table_id_iter: query_state.matched_table_ids.iter(),
archetype_id_iter: query_state.matched_archetype_ids.iter(),
current_len: 0,
@ -70,7 +77,10 @@ where
{
type Item = <Q::Fetch as Fetch<'w>>::Item;
#[inline]
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
// We can't currently reuse QueryIterationCursor in QueryIter for performance reasons. See #1763 for context.
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
unsafe {
if self.is_dense {
@ -143,6 +153,186 @@ where
}
}
pub struct QueryCombinationIter<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize>
where
F::Fetch: FilterFetch,
{
tables: &'w Tables,
archetypes: &'w Archetypes,
query_state: &'s QueryState<Q, F>,
world: &'w World,
cursors: [QueryIterationCursor<'s, Q, F>; K],
}
impl<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize> QueryCombinationIter<'w, 's, Q, F, K>
where
F::Fetch: FilterFetch,
{
/// # Safety
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
/// This does not validate that `world.id()` matches `query_state.world_id`. Calling this on a `world`
/// with a mismatched WorldId is unsound.
pub(crate) unsafe fn new(
world: &'w World,
query_state: &'s QueryState<Q, F>,
last_change_tick: u32,
change_tick: u32,
) -> Self {
// Initialize array with cursors.
// There is no FromIterator on arrays, so instead initialize it manually with MaybeUninit
// TODO: use MaybeUninit::uninit_array if it stabilizes
let mut cursors: [MaybeUninit<QueryIterationCursor<'s, Q, F>>; K] =
MaybeUninit::uninit().assume_init();
for (i, cursor) in cursors.iter_mut().enumerate() {
match i {
0 => cursor.as_mut_ptr().write(QueryIterationCursor::init(
world,
query_state,
last_change_tick,
change_tick,
)),
_ => cursor.as_mut_ptr().write(QueryIterationCursor::init_empty(
world,
query_state,
last_change_tick,
change_tick,
)),
}
}
// TODO: use MaybeUninit::array_assume_init if it stabilizes
let cursors: [QueryIterationCursor<'s, Q, F>; K] =
(&cursors as *const _ as *const [QueryIterationCursor<'s, Q, F>; K]).read();
QueryCombinationIter {
world,
query_state,
tables: &world.storages().tables,
archetypes: &world.archetypes,
cursors,
}
}
/// Safety:
/// The lifetime here is not restrictive enough for Fetch with &mut access,
/// as calling `fetch_next_aliased_unchecked` multiple times can produce multiple
/// references to the same component, leading to unique reference aliasing.
///.
/// It is always safe for shared access.
unsafe fn fetch_next_aliased_unchecked<'a>(
&mut self,
) -> Option<[<Q::Fetch as Fetch<'a>>::Item; K]>
where
Q::Fetch: Clone,
F::Fetch: Clone,
{
if K == 0 {
return None;
}
// first, iterate from last to first until next item is found
'outer: for i in (0..K).rev() {
match self.cursors[i].next(&self.tables, &self.archetypes, &self.query_state) {
Some(_) => {
// walk forward up to last element, propagating cursor state forward
for j in (i + 1)..K {
self.cursors[j] = self.cursors[j - 1].clone();
match self.cursors[j].next(
&self.tables,
&self.archetypes,
&self.query_state,
) {
Some(_) => {}
None if i > 0 => continue 'outer,
None => return None,
}
}
break;
}
None if i > 0 => continue,
None => return None,
}
}
// TODO: use MaybeUninit::uninit_array if it stabilizes
let mut values: [MaybeUninit<<Q::Fetch as Fetch<'a>>::Item>; K] =
MaybeUninit::uninit().assume_init();
for (value, cursor) in values.iter_mut().zip(&mut self.cursors) {
value.as_mut_ptr().write(cursor.peek_last().unwrap());
}
// TODO: use MaybeUninit::array_assume_init if it stabilizes
let values: [<Q::Fetch as Fetch<'a>>::Item; K] =
(&values as *const _ as *const [<Q::Fetch as Fetch<'a>>::Item; K]).read();
Some(values)
}
/// Get next combination of queried components
#[inline]
pub fn fetch_next(&mut self) -> Option<[<Q::Fetch as Fetch<'_>>::Item; K]>
where
Q::Fetch: Clone,
F::Fetch: Clone,
{
// safety: we are limiting the returned reference to self,
// making sure this method cannot be called multiple times without getting rid
// of any previously returned unique references first, thus preventing aliasing.
unsafe { self.fetch_next_aliased_unchecked() }
}
}
// Iterator type is intentionally implemented only for read-only access.
// Doing so for mutable references would be unsound, because calling `next`
// multiple times would allow multiple owned references to the same data to exist.
impl<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize> Iterator
for QueryCombinationIter<'w, 's, Q, F, K>
where
Q::Fetch: Clone + ReadOnlyFetch,
F::Fetch: Clone + FilterFetch + ReadOnlyFetch,
{
type Item = [<Q::Fetch as Fetch<'w>>::Item; K];
#[inline]
fn next(&mut self) -> Option<Self::Item> {
// Safety: it is safe to alias for ReadOnlyFetch
unsafe { QueryCombinationIter::fetch_next_aliased_unchecked(self) }
}
// NOTE: For unfiltered Queries this should actually return a exact size hint,
// to fulfil the ExactSizeIterator invariant, but this isn't practical without specialization.
// For more information see Issue #1686.
fn size_hint(&self) -> (usize, Option<usize>) {
if K == 0 {
return (0, Some(0));
}
let max_size: usize = self
.query_state
.matched_archetypes
.ones()
.map(|index| self.world.archetypes[ArchetypeId::new(index)].len())
.sum();
if max_size < K {
return (0, Some(0));
}
// n! / k!(n-k)! = (n*n-1*...*n-k+1) / k!
let max_combinations = (0..K)
.try_fold(1usize, |n, i| n.checked_mul(max_size - i))
.map(|n| {
let k_factorial: usize = (1..=K).product();
n / k_factorial
});
(0, max_combinations)
}
}
// NOTE: We can cheaply implement this for unfiltered Queries because we have:
// (1) pre-computed archetype matches
// (2) each archetype pre-computes length
@ -157,3 +347,150 @@ impl<'w, 's, Q: WorldQuery> ExactSizeIterator for QueryIter<'w, 's, Q, ()> {
.sum()
}
}
struct QueryIterationCursor<'s, Q: WorldQuery, F: WorldQuery> {
table_id_iter: std::slice::Iter<'s, TableId>,
archetype_id_iter: std::slice::Iter<'s, ArchetypeId>,
fetch: Q::Fetch,
filter: F::Fetch,
current_len: usize,
current_index: usize,
is_dense: bool,
}
impl<'s, Q: WorldQuery, F: WorldQuery> Clone for QueryIterationCursor<'s, Q, F>
where
Q::Fetch: Clone,
F::Fetch: Clone,
{
fn clone(&self) -> Self {
Self {
table_id_iter: self.table_id_iter.clone(),
archetype_id_iter: self.archetype_id_iter.clone(),
fetch: self.fetch.clone(),
filter: self.filter.clone(),
current_len: self.current_len,
current_index: self.current_index,
is_dense: self.is_dense,
}
}
}
impl<'s, Q: WorldQuery, F: WorldQuery> QueryIterationCursor<'s, Q, F>
where
F::Fetch: FilterFetch,
{
unsafe fn init_empty(
world: &World,
query_state: &'s QueryState<Q, F>,
last_change_tick: u32,
change_tick: u32,
) -> Self {
QueryIterationCursor {
table_id_iter: [].iter(),
archetype_id_iter: [].iter(),
..Self::init(world, query_state, last_change_tick, change_tick)
}
}
unsafe fn init(
world: &World,
query_state: &'s QueryState<Q, F>,
last_change_tick: u32,
change_tick: u32,
) -> Self {
let fetch = <Q::Fetch as Fetch>::init(
world,
&query_state.fetch_state,
last_change_tick,
change_tick,
);
let filter = <F::Fetch as Fetch>::init(
world,
&query_state.filter_state,
last_change_tick,
change_tick,
);
QueryIterationCursor {
is_dense: fetch.is_dense() && filter.is_dense(),
fetch,
filter,
table_id_iter: query_state.matched_table_ids.iter(),
archetype_id_iter: query_state.matched_archetype_ids.iter(),
current_len: 0,
current_index: 0,
}
}
/// retrieve item returned from most recent `next` call again.
#[inline]
unsafe fn peek_last<'w>(&mut self) -> Option<<Q::Fetch as Fetch<'w>>::Item> {
if self.current_index > 0 {
if self.is_dense {
Some(self.fetch.table_fetch(self.current_index - 1))
} else {
Some(self.fetch.archetype_fetch(self.current_index - 1))
}
} else {
None
}
}
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
// We can't currently reuse QueryIterationCursor in QueryIter for performance reasons. See #1763 for context.
#[inline(always)]
unsafe fn next<'w>(
&mut self,
tables: &'w Tables,
archetypes: &'w Archetypes,
query_state: &'s QueryState<Q, F>,
) -> Option<<Q::Fetch as Fetch<'w>>::Item> {
if self.is_dense {
loop {
if self.current_index == self.current_len {
let table_id = self.table_id_iter.next()?;
let table = &tables[*table_id];
self.fetch.set_table(&query_state.fetch_state, table);
self.filter.set_table(&query_state.filter_state, table);
self.current_len = table.len();
self.current_index = 0;
continue;
}
if !self.filter.table_filter_fetch(self.current_index) {
self.current_index += 1;
continue;
}
let item = self.fetch.table_fetch(self.current_index);
self.current_index += 1;
return Some(item);
}
} else {
loop {
if self.current_index == self.current_len {
let archetype_id = self.archetype_id_iter.next()?;
let archetype = &archetypes[*archetype_id];
self.fetch
.set_archetype(&query_state.fetch_state, archetype, tables);
self.filter
.set_archetype(&query_state.filter_state, archetype, tables);
self.current_len = archetype.len();
self.current_index = 0;
continue;
}
if !self.filter.archetype_filter_fetch(self.current_index) {
self.current_index += 1;
continue;
}
let item = self.fetch.archetype_fetch(self.current_index);
self.current_index += 1;
return Some(item);
}
}
}
}

View file

@ -37,6 +37,135 @@ mod tests {
assert_eq!(values, vec![&B(3)]);
}
#[test]
fn query_iter_combinations() {
let mut world = World::new();
world.spawn().insert_bundle((A(1), B(1)));
world.spawn().insert_bundle((A(2),));
world.spawn().insert_bundle((A(3),));
world.spawn().insert_bundle((A(4),));
let mut a_query = world.query::<&A>();
assert_eq!(a_query.iter_combinations::<0>(&world).count(), 0);
assert_eq!(
a_query.iter_combinations::<0>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query.iter_combinations::<1>(&world).count(), 4);
assert_eq!(
a_query.iter_combinations::<1>(&world).size_hint(),
(0, Some(4))
);
assert_eq!(a_query.iter_combinations::<2>(&world).count(), 6);
assert_eq!(
a_query.iter_combinations::<2>(&world).size_hint(),
(0, Some(6))
);
assert_eq!(a_query.iter_combinations::<3>(&world).count(), 4);
assert_eq!(
a_query.iter_combinations::<3>(&world).size_hint(),
(0, Some(4))
);
assert_eq!(a_query.iter_combinations::<4>(&world).count(), 1);
assert_eq!(
a_query.iter_combinations::<4>(&world).size_hint(),
(0, Some(1))
);
assert_eq!(a_query.iter_combinations::<5>(&world).count(), 0);
assert_eq!(
a_query.iter_combinations::<5>(&world).size_hint(),
(0, Some(0))
);
assert_eq!(a_query.iter_combinations::<1024>(&world).count(), 0);
assert_eq!(
a_query.iter_combinations::<1024>(&world).size_hint(),
(0, Some(0))
);
let values: Vec<[&A; 2]> = world.query::<&A>().iter_combinations(&world).collect();
assert_eq!(
values,
vec![
[&A(1), &A(2)],
[&A(1), &A(3)],
[&A(1), &A(4)],
[&A(2), &A(3)],
[&A(2), &A(4)],
[&A(3), &A(4)],
]
);
let size = a_query.iter_combinations::<3>(&world).size_hint();
assert_eq!(size.1, Some(4));
let values: Vec<[&A; 3]> = a_query.iter_combinations(&world).collect();
assert_eq!(
values,
vec![
[&A(1), &A(2), &A(3)],
[&A(1), &A(2), &A(4)],
[&A(1), &A(3), &A(4)],
[&A(2), &A(3), &A(4)],
]
);
let mut query = world.query::<&mut A>();
let mut combinations = query.iter_combinations_mut(&mut world);
while let Some([mut a, mut b, mut c]) = combinations.fetch_next() {
a.0 += 10;
b.0 += 100;
c.0 += 1000;
}
let values: Vec<[&A; 3]> = a_query.iter_combinations(&world).collect();
assert_eq!(
values,
vec![
[&A(31), &A(212), &A(1203)],
[&A(31), &A(212), &A(3004)],
[&A(31), &A(1203), &A(3004)],
[&A(212), &A(1203), &A(3004)]
]
);
let mut b_query = world.query::<&B>();
assert_eq!(
b_query.iter_combinations::<2>(&world).size_hint(),
(0, Some(0))
);
let values: Vec<[&B; 2]> = b_query.iter_combinations(&world).collect();
assert_eq!(values, Vec::<[&B; 2]>::new());
}
#[test]
fn query_iter_combinations_sparse() {
let mut world = World::new();
world
.register_component(ComponentDescriptor::new::<A>(StorageType::SparseSet))
.unwrap();
world.spawn_batch((1..=4).map(|i| (A(i),)));
let mut query = world.query::<&mut A>();
let mut combinations = query.iter_combinations_mut(&mut world);
while let Some([mut a, mut b, mut c]) = combinations.fetch_next() {
a.0 += 10;
b.0 += 100;
c.0 += 1000;
}
let mut query = world.query::<&A>();
let values: Vec<[&A; 3]> = query.iter_combinations(&world).collect();
assert_eq!(
values,
vec![
[&A(31), &A(212), &A(1203)],
[&A(31), &A(212), &A(3004)],
[&A(31), &A(1203), &A(3004)],
[&A(212), &A(1203), &A(3004)]
]
);
}
#[test]
fn multi_storage_query() {
let mut world = World::new();

View file

@ -3,8 +3,8 @@ use crate::{
component::ComponentId,
entity::Entity,
query::{
Access, Fetch, FetchState, FilterFetch, FilteredAccess, QueryIter, ReadOnlyFetch,
WorldQuery,
Access, Fetch, FetchState, FilterFetch, FilteredAccess, QueryCombinationIter, QueryIter,
ReadOnlyFetch, WorldQuery,
},
storage::TableId,
world::{World, WorldId},
@ -205,6 +205,27 @@ where
unsafe { self.iter_unchecked(world) }
}
#[inline]
pub fn iter_combinations<'w, 's, const K: usize>(
&'s mut self,
world: &'w World,
) -> QueryCombinationIter<'w, 's, Q, F, K>
where
Q::Fetch: ReadOnlyFetch,
{
// SAFE: query is read only
unsafe { self.iter_combinations_unchecked(world) }
}
#[inline]
pub fn iter_combinations_mut<'w, 's, const K: usize>(
&'s mut self,
world: &'w mut World,
) -> QueryCombinationIter<'w, 's, Q, F, K> {
// SAFE: query has unique world access
unsafe { self.iter_combinations_unchecked(world) }
}
/// # Safety
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
@ -222,6 +243,22 @@ where
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
#[inline]
pub unsafe fn iter_combinations_unchecked<'w, 's, const K: usize>(
&'s mut self,
world: &'w World,
) -> QueryCombinationIter<'w, 's, Q, F, K> {
self.validate_world_and_update_archetypes(world);
self.iter_combinations_unchecked_manual(
world,
world.last_change_tick(),
world.read_change_tick(),
)
}
/// # Safety
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
/// This does not validate that `world.id()` matches `self.world_id`. Calling this on a `world`
/// with a mismatched WorldId is unsound.
#[inline]
@ -234,6 +271,21 @@ where
QueryIter::new(world, self, last_change_tick, change_tick)
}
/// # Safety
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
/// This does not validate that `world.id()` matches `self.world_id`. Calling this on a `world`
/// with a mismatched WorldId is unsound.
#[inline]
pub(crate) unsafe fn iter_combinations_unchecked_manual<'w, 's, const K: usize>(
&'s self,
world: &'w World,
last_change_tick: u32,
change_tick: u32,
) -> QueryCombinationIter<'w, 's, Q, F, K> {
QueryCombinationIter::new(world, self, last_change_tick, change_tick)
}
#[inline]
pub fn for_each<'w>(
&mut self,
@ -345,6 +397,8 @@ where
last_change_tick: u32,
change_tick: u32,
) {
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
let mut fetch =
<Q::Fetch as Fetch>::init(world, &self.fetch_state, last_change_tick, change_tick);
let mut filter =
@ -397,6 +451,8 @@ where
last_change_tick: u32,
change_tick: u32,
) {
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
task_pool.scope(|scope| {
let fetch =
<Q::Fetch as Fetch>::init(world, &self.fetch_state, last_change_tick, change_tick);

View file

@ -2,7 +2,8 @@ use crate::{
component::Component,
entity::Entity,
query::{
Fetch, FilterFetch, QueryEntityError, QueryIter, QueryState, ReadOnlyFetch, WorldQuery,
Fetch, FilterFetch, QueryCombinationIter, QueryEntityError, QueryIter, QueryState,
ReadOnlyFetch, WorldQuery,
},
world::{Mut, World},
};
@ -157,6 +158,29 @@ where
}
}
/// Returns an [`Iterator`] over all possible combinations of `K` query results without repetition.
/// This can only be called for read-only queries
///
/// For permutations of size K of query returning N results, you will get:
/// - if K == N: one permutation of all query results
/// - if K < N: all possible K-sized combinations of query results, without repetition
/// - if K > N: empty set (no K-sized combinations exist)
#[inline]
pub fn iter_combinations<const K: usize>(&self) -> QueryCombinationIter<'_, '_, Q, F, K>
where
Q::Fetch: ReadOnlyFetch,
{
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
self.state.iter_combinations_unchecked_manual(
self.world,
self.last_change_tick,
self.change_tick,
)
}
}
/// Returns an [`Iterator`] over the query results.
#[inline]
pub fn iter_mut(&mut self) -> QueryIter<'_, '_, Q, F> {
@ -168,6 +192,42 @@ where
}
}
/// Iterates over all possible combinations of `K` query results without repetition.
///
/// The returned value is not an `Iterator`, because that would lead to aliasing of mutable references.
/// In order to iterate it, use `fetch_next` method with `while let Some(..)` loop pattern.
///
/// ```
/// # struct A;
/// # use bevy_ecs::prelude::*;
/// # fn some_system(mut query: Query<&mut A>) {
/// // iterate using `fetch_next` in while loop
/// let mut combinations = query.iter_combinations_mut();
/// while let Some([mut a, mut b]) = combinations.fetch_next() {
/// // mutably access components data
/// }
/// # }
/// ```
///
/// There is no `for_each` method, because it cannot be safely implemented
/// due to a [compiler bug](https://github.com/rust-lang/rust/issues/62529).
///
/// For immutable access see [`Query::iter_combinations`].
#[inline]
pub fn iter_combinations_mut<const K: usize>(
&mut self,
) -> QueryCombinationIter<'_, '_, Q, F, K> {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
self.state.iter_combinations_unchecked_manual(
self.world,
self.last_change_tick,
self.change_tick,
)
}
}
/// Returns an [`Iterator`] over the query results.
///
/// # Safety
@ -182,6 +242,25 @@ where
.iter_unchecked_manual(self.world, self.last_change_tick, self.change_tick)
}
/// Iterates over all possible combinations of `K` query results without repetition.
/// See [`Query::iter_combinations`].
///
/// # Safety
/// This allows aliased mutability. You must make sure this call does not result in multiple
/// mutable references to the same component
#[inline]
pub unsafe fn iter_combinations_unsafe<const K: usize>(
&self,
) -> QueryCombinationIter<'_, '_, Q, F, K> {
// SEMI-SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
self.state.iter_combinations_unchecked_manual(
self.world,
self.last_change_tick,
self.change_tick,
)
}
/// Runs `f` on each query result. This is faster than the equivalent iter() method, but cannot
/// be chained like a normal [`Iterator`].
///

View file

@ -0,0 +1,158 @@
use bevy::{core::FixedTimestep, prelude::*};
use rand::{thread_rng, Rng};
#[derive(Debug, Hash, PartialEq, Eq, Clone, StageLabel)]
struct FixedUpdateStage;
const DELTA_TIME: f64 = 0.01;
fn main() {
App::build()
.insert_resource(Msaa { samples: 4 })
.add_plugins(DefaultPlugins)
.add_startup_system(generate_bodies.system())
.add_stage_after(
CoreStage::Update,
FixedUpdateStage,
SystemStage::parallel()
.with_run_criteria(FixedTimestep::step(DELTA_TIME))
.with_system(interact_bodies.system())
.with_system(integrate.system()),
)
.run();
}
const GRAVITY_CONSTANT: f32 = 0.001;
const SOFTENING: f32 = 0.01;
const NUM_BODIES: usize = 100;
#[derive(Default)]
struct Mass(f32);
#[derive(Default)]
struct Acceleration(Vec3);
#[derive(Default)]
struct LastPos(Vec3);
#[derive(Bundle, Default)]
struct BodyBundle {
#[bundle]
pbr: PbrBundle,
mass: Mass,
last_pos: LastPos,
acceleration: Acceleration,
}
fn generate_bodies(
mut commands: Commands,
mut meshes: ResMut<Assets<Mesh>>,
mut materials: ResMut<Assets<StandardMaterial>>,
) {
let mesh = meshes.add(Mesh::from(shape::Icosphere {
radius: 1.0,
subdivisions: 3,
}));
let pos_range = 1.0..15.0;
let color_range = 0.5..1.0;
let vel_range = -0.5..0.5;
let mut rng = thread_rng();
for _ in 0..NUM_BODIES {
let mass_value_cube_root: f32 = rng.gen_range(0.5..4.0);
let mass_value: f32 = mass_value_cube_root * mass_value_cube_root * mass_value_cube_root;
let position = Vec3::new(
rng.gen_range(-1.0..1.0),
rng.gen_range(-1.0..1.0),
rng.gen_range(-1.0..1.0),
)
.normalize()
* rng.gen_range(pos_range.clone());
commands.spawn_bundle(BodyBundle {
pbr: PbrBundle {
transform: Transform {
translation: position,
scale: Vec3::splat(mass_value_cube_root * 0.1),
..Default::default()
},
mesh: mesh.clone(),
material: materials.add(
Color::rgb_linear(
rng.gen_range(color_range.clone()),
rng.gen_range(color_range.clone()),
rng.gen_range(color_range.clone()),
)
.into(),
),
..Default::default()
},
mass: Mass(mass_value),
acceleration: Acceleration(Vec3::ZERO),
last_pos: LastPos(
position
- Vec3::new(
rng.gen_range(vel_range.clone()),
rng.gen_range(vel_range.clone()),
rng.gen_range(vel_range.clone()),
) * DELTA_TIME as f32,
),
});
}
// add bigger "star" body in the center
commands
.spawn_bundle(BodyBundle {
pbr: PbrBundle {
transform: Transform {
scale: Vec3::splat(0.5),
..Default::default()
},
mesh: meshes.add(Mesh::from(shape::Icosphere {
radius: 1.0,
subdivisions: 5,
})),
material: materials.add((Color::ORANGE_RED * 10.0).into()),
..Default::default()
},
mass: Mass(1000.0),
..Default::default()
})
.insert(PointLight {
color: Color::ORANGE_RED,
..Default::default()
});
commands.spawn_bundle(PerspectiveCameraBundle {
transform: Transform::from_xyz(0.0, 10.5, -20.0).looking_at(Vec3::ZERO, Vec3::Y),
..Default::default()
});
}
fn interact_bodies(mut query: Query<(&Mass, &GlobalTransform, &mut Acceleration)>) {
let mut iter = query.iter_combinations_mut();
while let Some([(Mass(m1), transform1, mut acc1), (Mass(m2), transform2, mut acc2)]) =
iter.fetch_next()
{
let delta = transform2.translation - transform1.translation;
let distance_sq: f32 = delta.length_squared();
let f = GRAVITY_CONSTANT / (distance_sq * (distance_sq + SOFTENING).sqrt());
let force_unit_mass = delta * f;
acc1.0 += force_unit_mass * *m2;
acc2.0 -= force_unit_mass * *m1;
}
}
fn integrate(mut query: Query<(&mut Acceleration, &mut Transform, &mut LastPos)>) {
let dt_sq = (DELTA_TIME * DELTA_TIME) as f32;
for (mut acceleration, mut transform, mut last_pos) in query.iter_mut() {
// verlet integration
// x(t+dt) = 2x(t) - x(t-dt) + a(t)dt^2 + O(dt^4)
let new_pos =
transform.translation + transform.translation - last_pos.0 + acceleration.0 * dt_sq;
acceleration.0 = Vec3::ZERO;
last_pos.0 = transform.translation;
transform.translation = new_pos;
}
}