Repository: bevyengine/bevy
File: crates/bevy_tasks/src/single_threaded_task_pool.rs (Rust)

Code before:
            let pending_tasks: RefCell<usize> = RefCell::new(0);
            // SAFETY: As above, all futures must complete in this function so we can change the lifetime
            let pending_tasks: &'env RefCell<usize> = unsafe { mem::transmute(&pending_tasks) };

            let mut scope = Scope {
                executor_ref,
                pending_tasks,
                results_ref,
                scope: PhantomData,
                env: PhantomData,
            };

            // SAFETY: As above, all futures must complete in this function so we can change the lifetime
            let scope_ref: &'env mut Scope<'_, 'env, T> = unsafe { mem::transmute(&mut scope) };

            f(scope_ref);

            // Wait until the socpe is complete
            while *pending_tasks.borrow() != 0 {
                block_on(executor.tick());
            }

            results
                .take()
                .into_iter()
                .map(|result| result.unwrap())
                .collect()
        })
    }

    /// Spawns a static future onto the thread pool. The returned Task is a future, which can be polled
    /// to retrieve the output of the original future. Dropping the task will attempt to cancel it.
    /// It can also be "detached", allowing it to continue running without having to be polled by the
    /// end-user.
    ///
    /// If the provided future is non-`Send`, [`TaskPool::spawn_local`] should be used instead.
    pub fn spawn<T>(
        &self,
        future: impl Future<Output = T> + 'static + MaybeSend + MaybeSync,
    ) -> Task<T>
    where
        T: 'static + MaybeSend + MaybeSync,
    {
        crate::cfg::switch! {{
            crate::cfg::web => {
                Task::wrap_future(future)
            }
            crate::cfg::std => {
                LOCAL_EXECUTOR.with(|executor| {
                    let task = executor.spawn(future);
                    // Loop until all tasks are done

Review comment (suggestion):
I think we need a Drop impl on `Scope` to cancel/remove any remaining tasks on executor in case one of them panicks. But I guess we can't cancel tasks on web? Maybe we need to abort instead of just panicking? if the executor gets ticked after the scope returns (maybe the user catch unwinds) it would be UB.

Code after:
            env: PhantomData,
        };

        // SAFETY: As above, all futures must complete in this function so we can change the lifetime
        let scope_ref: &'env mut Scope<'_, 'env, T> = unsafe { mem::transmute(&mut scope) };

        f(scope_ref);

        // Wait until the scope is complete
        block_on(executor.run(async {
            while pending_tasks.get() != 0 {
                futures_lite::future::yield_now().await;
            }
        }));

        results
            .take()
            .into_iter()
            .map(|result| result.unwrap())
            .collect()
    }

    /// Spawns a static future onto the thread pool. The returned Task is a future, which can be polled
    /// to retrieve the output of the original future. Dropping the task will attempt to cancel it.
    /// It can also be "detached", allowing it to continue running without having to be polled by the
    /// end-user.
    ///
    /// If the provided future is non-`Send`, [`TaskPool::spawn_local`] should be used instead.
    pub fn spawn<T>(
        &self,
        future: impl Future<Output = T> + 'static + MaybeSend + MaybeSync,
    ) -> Task<T>
    where
        T: 'static + MaybeSend + MaybeSync,
    {
        crate::cfg::switch! {{
            crate::cfg::web => {
                Task::wrap_future(future)
            }
            crate::cfg::std => {
                LOCAL_EXECUTOR.with(|executor| {
                    let task = executor.spawn(future);
                    // Loop until all tasks are done
                    while executor.try_tick() {}

                    Task::new(task)
                })
            }
            _ => {
                let task = LOCAL_EXECUTOR.spawn(future);
                // Loop until all tasks are done