launch.h

#pragma once
#include <asyncpp/detail/promise_allocator_base.h>
#include <asyncpp/scope_guard.h>
#include <atomic>
#include <cassert>
#include <exception>
#include <future>
#include <stdexcept>
#include <tuple>
#include <utility>
 
namespace asyncpp {
    namespace detail {
        template<ByteAllocator Allocator = default_allocator_type>
        struct launch_task {
            // Promise type of this task
            struct promise_type : promise_allocator_base<Allocator> {
                constexpr launch_task get_return_object() noexcept { return {}; }
                constexpr suspend_never initial_suspend() noexcept { return {}; }
                constexpr suspend_never final_suspend() noexcept { return {}; }
                constexpr void return_void() noexcept {}
                void unhandled_exception() noexcept { std::terminate(); }
            };
        };
    } // namespace detail
 
    template<typename Awaitable, ByteAllocator Allocator = default_allocator_type>
    void launch(Awaitable&& awaitable, const Allocator& allocator = {}) {
        [](std::decay_t<Awaitable> awaitable, const Allocator&) -> detail::launch_task<Allocator> {
            co_await std::move(awaitable);
        }(std::forward<decltype(awaitable)>(awaitable), allocator);
    }
 
    class async_launch_scope {
        std::atomic<size_t> m_count{0U};
        std::atomic<void*> m_continuation{nullptr};
 
    public:
        constexpr async_launch_scope() noexcept = default;
        async_launch_scope(const async_launch_scope&) = delete;
        async_launch_scope& operator=(const async_launch_scope&) = delete;
        ~async_launch_scope() { assert(m_count.load() == 0); }
 
        template<typename Awaitable, ByteAllocator Allocator = default_allocator_type>
        void launch(Awaitable&& awaitable, const Allocator& allocator = {}) {
            [](async_launch_scope* scope, std::decay_t<Awaitable> awaitable,
               const Allocator&) -> detail::launch_task<Allocator> {
                scope->m_count.fetch_add(1);
                scope_guard guard{[scope]() noexcept {
                    // If this is the last task
                    if (scope->m_count.fetch_sub(1) == 1) {
                        // And we are being awaited
                        auto hdl = scope->m_continuation.exchange(nullptr);
                        // Resume the awaiter
                        if (hdl != nullptr) coroutine_handle<>::from_address(hdl).resume();
                    }
                }};
                co_await awaitable;
            }(this, std::forward<decltype(awaitable)>(awaitable), allocator);
        }
 
        template<typename Callable, typename... Args, ByteAllocator Allocator = default_allocator_type>
            requires(std::is_invocable_v<Callable, Args...>)
        void invoke_tuple(Callable&& callable, std::tuple<Args...>&& args, const Allocator& allocator = {}) {
            [](async_launch_scope* scope, Callable callable, std::tuple<Args...>&& args,
               const Allocator&) -> detail::launch_task<Allocator> {
                scope->m_count.fetch_add(1);
                scope_guard guard{[scope]() noexcept {
                    // If this is the last task
                    if (scope->m_count.fetch_sub(1) == 1) {
                        // And we are being awaited
                        auto hdl = scope->m_continuation.exchange(nullptr);
                        // Resume the awaiter
                        if (hdl != nullptr) coroutine_handle<>::from_address(hdl).resume();
                    }
                }};
                co_await std::apply(callable, std::move(args));
            }(this, std::forward<Callable>(callable), std::move(args), allocator);
        }
 
        template<typename Callable, typename... Args>
            requires(std::is_invocable_v<Callable, Args...>)
        void invoke(Callable&& callable, Args&&... args) {
            [](async_launch_scope* scope, Callable callable, Args&&... args) -> detail::launch_task<> {
                scope->m_count.fetch_add(1);
                scope_guard guard{[scope]() noexcept {
                    // If this is the last task
                    if (scope->m_count.fetch_sub(1) == 1) {
                        // And we are being awaited
                        auto hdl = scope->m_continuation.exchange(nullptr);
                        // Resume the awaiter
                        if (hdl != nullptr) coroutine_handle<>::from_address(hdl).resume();
                    }
                }};
                co_await std::invoke(callable, std::forward<Args>(args)...);
            }(this, std::forward<Callable>(callable), std::forward<Args>(args)...);
        }
 
        [[nodiscard]] auto join() noexcept {
            struct awaiter {
                async_launch_scope* m_scope;
                [[nodiscard]] bool await_ready() const noexcept {
                    // Dont wait if theres nothing to await
                    return m_scope->m_count.load() == 0;
                }
                [[nodiscard]] bool await_suspend(coroutine_handle<> hdl) const {
                    // Set our coroutine if there is noone waiting
                    void* expected = nullptr;
                    if (!m_scope->m_continuation.compare_exchange_strong(expected, hdl.address()))
                        throw std::logic_error("duplicate join");
                    // We might have nothing left to wait on if the last coroutine finished in the meantime
                    return m_scope->m_count.load() != 0;
                }
                constexpr void await_resume() const noexcept {}
            };
            return awaiter{this};
        }
 
        template<ByteAllocator Allocator = default_allocator_type>
        std::future<void> join_future(const Allocator& allocator = {}) noexcept {
            std::promise<void> done;
            auto res = done.get_future();
            [](async_launch_scope* that, std::promise<void> done, const Allocator&) -> detail::launch_task<Allocator> {
                try {
                    co_await that->join();
                    done.set_value();
                } catch (...) { done.set_exception(std::current_exception()); }
            }(this, std::move(done), allocator);
            return res;
        }
 
        [[nodiscard]] size_t inflight_coroutines() const noexcept { return m_count.load(std::memory_order::relaxed); }
 
        [[nodiscard]] bool all_done() const noexcept { return inflight_coroutines() == 0; }
    };
} // namespace asyncpp