[Merged by Bors] - Nested spawns on scope

crates/bevy_ecs/src/schedule/executor_parallel.rs

@@ -166,7 +166,7 @@ impl ParallelExecutor {
     /// queues systems with no dependencies to run (or skip) at next opportunity.
     fn prepare_systems<'scope>(
         &mut self,
-        scope: &mut Scope<'scope, ()>,
+        scope: &Scope<'_, 'scope, ()>,
         systems: &'scope mut [SystemContainer],
         world: &'scope World,
     ) {
@@ -236,7 +236,7 @@ impl ParallelExecutor {
                 if system_data.is_send {
                     scope.spawn(task);
                 } else {
-                    scope.spawn_local(task);
+                    scope.spawn_on_scope(task);
                 }
 
                 #[cfg(test)]
@@ -271,7 +271,7 @@ impl ParallelExecutor {
                 if system_data.is_send {
                     scope.spawn(task);
                 } else {
-                    scope.spawn_local(task);
+                    scope.spawn_on_scope(task);
                 }
             }
         }

crates/bevy_tasks/Cargo.toml

@@ -13,6 +13,7 @@ futures-lite = "1.4.0"
 async-executor = "1.3.0"
 async-channel = "1.4.2"
 once_cell = "1.7"
+concurrent-queue = "1.2.2"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 wasm-bindgen-futures = "0.4"

crates/bevy_tasks/src/single_threaded_task_pool.rs

@@ -1,5 +1,6 @@
 use std::{
     future::Future,
+    marker::PhantomData,
     mem,
     sync::{Arc, Mutex},
 };
@@ -61,27 +62,34 @@ impl TaskPool {
     /// to spawn tasks. This function will await the completion of all tasks before returning.
     ///
     /// This is similar to `rayon::scope` and `crossbeam::scope`
-    pub fn scope<'scope, F, T>(&self, f: F) -> Vec<T>
+    pub fn scope<'env, F, T>(&self, f: F) -> Vec<T>
     where
-        F: FnOnce(&mut Scope<'scope, T>) + 'scope + Send,
+        F: for<'scope> FnOnce(&'env mut Scope<'scope, 'env, T>),
         T: Send + 'static,
     {
         let executor = &async_executor::LocalExecutor::new();
-        let executor: &'scope async_executor::LocalExecutor<'scope> =
+        let executor: &'env async_executor::LocalExecutor<'env> =
             unsafe { mem::transmute(executor) };
 
+        let results: Mutex<Vec<Arc<Mutex<Option<T>>>>> = Mutex::new(Vec::new());
+        let results: &'env Mutex<Vec<Arc<Mutex<Option<T>>>>> = unsafe { mem::transmute(&results) };
+
         let mut scope = Scope {
             executor,
-            results: Vec::new(),
+            results,
+            scope: PhantomData,
+            env: PhantomData,
         };
 
-        f(&mut scope);
+        let scope_ref: &'env mut Scope<'_, 'env, T> = unsafe { mem::transmute(&mut scope) };
+
+        f(scope_ref);
 
         // Loop until all tasks are done
         while executor.try_tick() {}
 
-        scope
-            .results
+        let results = scope.results.lock().unwrap();
+        results
             .iter()
             .map(|result| result.lock().unwrap().take().unwrap())
             .collect()
@@ -127,32 +135,36 @@ impl FakeTask {
 ///
 /// For more information, see [`TaskPool::scope`].
 #[derive(Debug)]
-pub struct Scope<'scope, T> {
-    executor: &'scope async_executor::LocalExecutor<'scope>,
+pub struct Scope<'scope, 'env: 'scope, T> {
+    executor: &'env async_executor::LocalExecutor<'env>,
     // Vector to gather results of all futures spawned during scope run
-    results: Vec<Arc<Mutex<Option<T>>>>,
+    results: &'env Mutex<Vec<Arc<Mutex<Option<T>>>>>,
+
+    // make `Scope` invariant over 'scope and 'env
+    scope: PhantomData<&'scope mut &'scope ()>,
+    env: PhantomData<&'env mut &'env ()>,
 }
 
-impl<'scope, T: Send + 'scope> Scope<'scope, T> {
+impl<'scope, 'env, T: Send + 'env> Scope<'scope, 'env, T> {
     /// Spawns a scoped future onto the thread-local executor. The scope *must* outlive
     /// the provided future. The results of the future will be returned as a part of
     /// [`TaskPool::scope`]'s return value.
     ///
     /// On the single threaded task pool, it just calls [`Scope::spawn_local`].
     ///
     /// For more information, see [`TaskPool::scope`].
-    pub fn spawn<Fut: Future<Output = T> + 'scope + Send>(&mut self, f: Fut) {
-        self.spawn_local(f);
+    pub fn spawn<Fut: Future<Output = T> + 'env>(&self, f: Fut) {
+        self.spawn_on_scope(f);
     }
 
-    /// Spawns a scoped future onto the thread-local executor. The scope *must* outlive
-    /// the provided future. The results of the future will be returned as a part of
-    /// [`TaskPool::scope`]'s return value.
+    /// Spawns a scoped future that runs on the thread the scope called from. The
+    /// scope *must* outlive the provided future. The results of the future will be
+    /// returned as a part of [`TaskPool::scope`]'s return value.
     ///
     /// For more information, see [`TaskPool::scope`].
-    pub fn spawn_local<Fut: Future<Output = T> + 'scope>(&mut self, f: Fut) {
+    pub fn spawn_on_scope<Fut: Future<Output = T> + 'env>(&self, f: Fut) {
         let result = Arc::new(Mutex::new(None));
-        self.results.push(result.clone());
+        self.results.lock().unwrap().push(result.clone());
         let f = async move {
             result.lock().unwrap().replace(f.await);
         };