[Merged by Bors] - Remove task_pool parameter from par_for_each(_mut)

benches/benches/bevy_ecs/ecs_bench_suite/heavy_compute.rs

@@ -1,5 +1,5 @@
 use bevy_ecs::prelude::*;
-use bevy_tasks::TaskPool;
+use bevy_tasks::{ComputeTaskPool, TaskPool};
 use glam::*;
 
 #[derive(Component, Copy, Clone)]
@@ -29,8 +29,8 @@ impl Benchmark {
             )
         }));
 
-        fn sys(task_pool: Res<TaskPool>, mut query: Query<(&mut Position, &mut Transform)>) {
-            query.par_for_each_mut(&task_pool, 128, |(mut pos, mut mat)| {
+        fn sys(mut query: Query<(&mut Position, &mut Transform)>) {
+            query.par_for_each_mut(128, |(mut pos, mut mat)| {
                 for _ in 0..100 {
                     mat.0 = mat.0.inverse();
                 }
@@ -39,7 +39,7 @@ impl Benchmark {
             });
         }
 
-        world.insert_resource(TaskPool::default());
+        world.insert_resource(ComputeTaskPool(TaskPool::default()));
         let mut system = IntoSystem::into_system(sys);
         system.initialize(&mut world);
         system.update_archetype_component_access(&world);

crates/bevy_ecs/src/lib.rs

@@ -56,7 +56,7 @@ mod tests {
         query::{Added, ChangeTrackers, Changed, FilteredAccess, With, Without, WorldQuery},
         world::{Mut, World},
     };
-    use bevy_tasks::TaskPool;
+    use bevy_tasks::{ComputeTaskPool, TaskPool};
     use std::{
         any::TypeId,
         sync::{
@@ -373,7 +373,7 @@ mod tests {
     #[test]
     fn par_for_each_dense() {
         let mut world = World::new();
-        let task_pool = TaskPool::default();
+        world.insert_resource(ComputeTaskPool(TaskPool::default()));
         let e1 = world.spawn().insert(A(1)).id();
         let e2 = world.spawn().insert(A(2)).id();
         let e3 = world.spawn().insert(A(3)).id();
@@ -382,7 +382,7 @@ mod tests {
         let results = Arc::new(Mutex::new(Vec::new()));
         world
             .query::<(Entity, &A)>()
-            .par_for_each(&world, &task_pool, 2, |(e, &A(i))| {
+            .par_for_each(&world, 2, |(e, &A(i))| {
                 results.lock().unwrap().push((e, i));
             });
         results.lock().unwrap().sort();
@@ -395,8 +395,7 @@ mod tests {
     #[test]
     fn par_for_each_sparse() {
         let mut world = World::new();
-
-        let task_pool = TaskPool::default();
+        world.insert_resource(ComputeTaskPool(TaskPool::default()));
         let e1 = world.spawn().insert(SparseStored(1)).id();
         let e2 = world.spawn().insert(SparseStored(2)).id();
         let e3 = world.spawn().insert(SparseStored(3)).id();
@@ -405,7 +404,6 @@ mod tests {
         let results = Arc::new(Mutex::new(Vec::new()));
         world.query::<(Entity, &SparseStored)>().par_for_each(
             &world,
-            &task_pool,
             2,
             |(e, &SparseStored(i))| results.lock().unwrap().push((e, i)),
         );

crates/bevy_ecs/src/query/state.rs

@@ -10,15 +10,16 @@ use crate::{
     storage::TableId,
     world::{World, WorldId},
 };
-use bevy_tasks::TaskPool;
+use bevy_tasks::{ComputeTaskPool, TaskPool};
 use fixedbitset::FixedBitSet;
-use std::fmt;
+use std::{fmt, ops::Deref};
 
 use super::{QueryFetch, QueryItem, ROQueryFetch, ROQueryItem};
 
 /// Provides scoped access to a [`World`] state according to a given [`WorldQuery`] and query filter.
 pub struct QueryState<Q: WorldQuery, F: WorldQuery = ()> {
     world_id: WorldId,
+    task_pool: Option<TaskPool>,
     pub(crate) archetype_generation: ArchetypeGeneration,
     pub(crate) matched_tables: FixedBitSet,
     pub(crate) matched_archetypes: FixedBitSet,
@@ -59,6 +60,9 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
 
         let mut state = Self {
             world_id: world.id(),
+            task_pool: world
+                .get_resource::<ComputeTaskPool>()
+                .map(|task_pool| task_pool.deref().clone()),
             archetype_generation: ArchetypeGeneration::initial(),
             matched_table_ids: Vec::new(),
             matched_archetype_ids: Vec::new(),
@@ -683,15 +687,17 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
         );
     }
 
-    /// Runs `func` on each query result in parallel using the given `task_pool`.
+    /// Runs `func` on each query result in parallel.
     ///
     /// This can only be called for read-only queries, see [`Self::par_for_each_mut`] for
     /// write-queries.
+    ///
+    /// # Panics
+    /// [`ComputeTaskPool`] is not stored as a resource in `world`.
     #[inline]
     pub fn par_for_each<'w, FN: Fn(ROQueryItem<'w, Q>) + Send + Sync + Clone>(
         &mut self,
         world: &'w World,
-        task_pool: &TaskPool,
         batch_size: usize,
         func: FN,
     ) {
@@ -700,7 +706,6 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
             self.update_archetypes(world);
             self.par_for_each_unchecked_manual::<ROQueryFetch<Q>, FN>(
                 world,
-                task_pool,
                 batch_size,
                 func,
                 world.last_change_tick(),
@@ -709,12 +714,14 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
         }
     }
 
-    /// Runs `func` on each query result in parallel using the given `task_pool`.
+    /// Runs `func` on each query result in parallel.
+    ///
+    /// # Panics
+    /// [`ComputeTaskPool`] is not stored as a resource in `world`.
     #[inline]
     pub fn par_for_each_mut<'w, FN: Fn(QueryItem<'w, Q>) + Send + Sync + Clone>(
         &mut self,
         world: &'w mut World,
-        task_pool: &TaskPool,
         batch_size: usize,
         func: FN,
     ) {
@@ -723,7 +730,6 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
             self.update_archetypes(world);
             self.par_for_each_unchecked_manual::<QueryFetch<Q>, FN>(
                 world,
-                task_pool,
                 batch_size,
                 func,
                 world.last_change_tick(),
@@ -732,10 +738,13 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
         }
     }
 
-    /// Runs `func` on each query result in parallel using the given `task_pool`.
+    /// Runs `func` on each query result in parallel.
     ///
     /// This can only be called for read-only queries.
     ///
+    /// # Panics
+    /// [`ComputeTaskPool`] is not stored as a resource in `world`.
+    ///
     /// # Safety
     ///
     /// This does not check for mutable query correctness. To be safe, make sure mutable queries
@@ -744,14 +753,12 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
     pub unsafe fn par_for_each_unchecked<'w, FN: Fn(QueryItem<'w, Q>) + Send + Sync + Clone>(
         &mut self,
         world: &'w World,
-        task_pool: &TaskPool,
         batch_size: usize,
         func: FN,
     ) {
         self.update_archetypes(world);
         self.par_for_each_unchecked_manual::<QueryFetch<Q>, FN>(
             world,
-            task_pool,
             batch_size,
             func,
             world.last_change_tick(),
@@ -827,6 +834,9 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
     /// the current change tick are given. This is faster than the equivalent
     /// iter() method, but cannot be chained like a normal [`Iterator`].
     ///
+    /// # Panics
+    /// [`ComputeTaskPool`] is not stored as a resource in `world`.
+    ///
     /// # Safety
     ///
     /// This does not check for mutable query correctness. To be safe, make sure mutable queries
@@ -840,12 +850,15 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
     >(
         &self,
         world: &'w World,
-        task_pool: &TaskPool,
         batch_size: usize,
         func: FN,
         last_change_tick: u32,
         change_tick: u32,
     ) {
+        let task_pool = self
+            .task_pool
+            .clone()
+            .expect("Cannot iterate query in parallel. No ComputeTaskPool initialized.");
         // 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| {

crates/bevy_ecs/src/system/query.rs

@@ -7,7 +7,6 @@ use crate::{
     },
     world::{Mut, World},
 };
-use bevy_tasks::TaskPool;
 use std::{any::TypeId, fmt::Debug};
 
 /// Provides scoped access to components in a [`World`].
@@ -493,7 +492,7 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
         };
     }
 
-    /// Runs `f` on each query result in parallel using the given [`TaskPool`].
+    /// Runs `f` on each query result in parallel using the [`World`]'s [`ComputeTaskPool`].
     ///
     /// This can only be called for immutable data, see [`Self::par_for_each_mut`] for
     /// mutable access.
@@ -502,21 +501,24 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
     ///
     /// The items in the query get sorted into batches.
     /// Internally, this function spawns a group of futures that each take on a `batch_size` sized section of the items (or less if the division is not perfect).
-    /// Then, the tasks in the [`TaskPool`] work through these futures.
+    /// Then, the tasks in the [`ComputeTaskPool`] work through these futures.
     ///
     /// You can use this value to tune between maximum multithreading ability (many small batches) and minimum parallelization overhead (few big batches).
     /// Rule of thumb: If the function body is (mostly) computationally expensive but there are not many items, a small batch size (=more batches) may help to even out the load.
     /// If the body is computationally cheap and you have many items, a large batch size (=fewer batches) avoids spawning additional futures that don't help to even out the load.
     ///
     /// # Arguments
     ///
-    ///* `task_pool` - The [`TaskPool`] to use
     ///* `batch_size` - The number of batches to spawn
     ///* `f` - The function to run on each item in the query
+    ///
+    /// # Panics
+    /// [`ComputeTaskPool`] is not stored as a resource in `world`.
+    ///
+    /// [`ComputeTaskPool`]: bevy_tasks::prelude::ComputeTaskPool
     #[inline]
     pub fn par_for_each<'this>(
         &'this self,
-        task_pool: &TaskPool,
         batch_size: usize,
         f: impl Fn(ROQueryItem<'this, Q>) + Send + Sync + Clone,
     ) {
@@ -526,7 +528,6 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
             self.state
                 .par_for_each_unchecked_manual::<ROQueryFetch<Q>, _>(
                     self.world,
-                    task_pool,
                     batch_size,
                     f,
                     self.last_change_tick,
@@ -535,12 +536,16 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
         };
     }
 
-    /// Runs `f` on each query result in parallel using the given [`TaskPool`].
+    /// Runs `f` on each query result in parallel using the [`World`]'s [`ComputeTaskPool`].
     /// See [`Self::par_for_each`] for more details.
+    ///
+    /// # Panics
+    /// [`ComputeTaskPool`] is not stored as a resource in `world`.
+    ///
+    /// [`ComputeTaskPool`]: bevy_tasks::prelude::ComputeTaskPool
     #[inline]
     pub fn par_for_each_mut<'a, FN: Fn(QueryItem<'a, Q>) + Send + Sync + Clone>(
         &'a mut self,
-        task_pool: &TaskPool,
         batch_size: usize,
         f: FN,
     ) {
@@ -550,7 +555,6 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
             self.state
                 .par_for_each_unchecked_manual::<QueryFetch<Q>, FN>(
                     self.world,
-                    task_pool,
                     batch_size,
                     f,
                     self.last_change_tick,

examples/ecs/parallel_query.rs

@@ -1,4 +1,4 @@
-use bevy::{prelude::*, tasks::prelude::*};
+use bevy::prelude::*;
 use rand::random;
 
 #[derive(Component, Deref)]
@@ -21,26 +21,22 @@ fn spawn_system(mut commands: Commands, asset_server: Res<AssetServer>) {
 }
 
 // Move sprites according to their velocity
-fn move_system(pool: Res<ComputeTaskPool>, mut sprites: Query<(&mut Transform, &Velocity)>) {
+fn move_system(mut sprites: Query<(&mut Transform, &Velocity)>) {
     // Compute the new location of each sprite in parallel on the
     // ComputeTaskPool using batches of 32 sprites
     //
-    // This example is only for demonstrative purposes.  Using a
+    // This example is only for demonstrative purposes. Using a
     // ParallelIterator for an inexpensive operation like addition on only 128
     // elements will not typically be faster than just using a normal Iterator.
     // See the ParallelIterator documentation for more information on when
     // to use or not use ParallelIterator over a normal Iterator.
-    sprites.par_for_each_mut(&pool, 32, |(mut transform, velocity)| {
+    sprites.par_for_each_mut(32, |(mut transform, velocity)| {
         transform.translation += velocity.extend(0.0);
     });
 }
 
 // Bounce sprites outside the window
-fn bounce_system(
-    pool: Res<ComputeTaskPool>,
-    windows: Res<Windows>,
-    mut sprites: Query<(&Transform, &mut Velocity)>,
-) {
+fn bounce_system(windows: Res<Windows>, mut sprites: Query<(&Transform, &mut Velocity)>) {
     let window = windows.primary();
     let width = window.width();
     let height = window.height();
@@ -51,7 +47,7 @@ fn bounce_system(
     sprites
         // Batch size of 32 is chosen to limit the overhead of
         // ParallelIterator, since negating a vector is very inexpensive.
-        .par_for_each_mut(&pool, 32, |(transform, mut v)| {
+        .par_for_each_mut(32, |(transform, mut v)| {
             if !(left < transform.translation.x
                 && transform.translation.x < right
                 && bottom < transform.translation.y