rework executor code for performance

crates/bevy_tasks/src/task_pool.rs

@@ -2,13 +2,12 @@ use std::{
     future::Future,
     marker::PhantomData,
     mem,
-    pin::Pin,
     sync::Arc,
     thread::{self, JoinHandle},
 };
 
 use concurrent_queue::ConcurrentQueue;
-use futures_lite::{future, pin};
+use futures_lite::{future, FutureExt};
 
 use crate::Task;
 
@@ -235,40 +234,27 @@ impl TaskPool {
 
         f(scope_ref);
 
-        let fut = async move {
-            let mut results = Vec::with_capacity(spawned.len());
-            while let Ok(task) = spawned.pop() {
-                results.push(task.await);
-            }
+        future::block_on(async move {
+            let get_results = async move {
+                let mut results = Vec::with_capacity(spawned.len());
+                while let Ok(task) = spawned.pop() {
+                    results.push(task.await);
+                }
 
-            results
-        };
+                results
+            };
 
-        // Pin the futures on the stack.
-        pin!(fut);
-
-        // SAFETY: This function blocks until all futures complete, so we do not read/write
-        // the data from futures outside of the 'scope lifetime. However,
-        // rust has no way of knowing this so we must convert to 'static
-        // here to appease the compiler as it is unable to validate safety.
-        let fut: Pin<&mut (dyn Future<Output = Vec<T>> + 'static + Send)> = fut;
-        let fut: Pin<&'static mut (dyn Future<Output = Vec<T>> + 'static + Send)> =
-            unsafe { mem::transmute(fut) };
-
-        // The thread that calls scope() will participate in driving tasks in the pool
-        // forward until the tasks that are spawned by this scope() call
-        // complete. (If the caller of scope() happens to be a thread in
-        // this thread pool, and we only have one thread in the pool, then
-        // simply calling future::block_on(spawned) would deadlock.)
-        let mut spawned = task_scope_executor.spawn(fut);
-        loop {
-            if let Some(result) = future::block_on(future::poll_once(&mut spawned)) {
-                break result;
+            let tick_forever = async move {
+                loop {
+                    self.executor.try_tick();
+                    task_scope_executor.try_tick();
+
+                    future::yield_now().await
+                }
             };
 
-            self.executor.try_tick();
-            task_scope_executor.try_tick();
-        }
+            get_results.or(tick_forever).await
+        })
     }
 
     /// Spawns a static future onto the thread pool. The returned Task is a future. It can also be