task.hpp

Go to the documentation of this file.

#pragma once
 
// C++ coroutine Task<T> (C++20 and C++23): a lazy, single-shot, awaitable that returns T (or void).
//
//   Task<int> compute() { co_return 42; }
//   Task<void> log()    { co_return; }
//
// The task is started when co_awaited (or by Task::start()).
// Exceptions thrown in the coroutine body are stored and rethrown by co_await.
 
#include <coroutine>
#include <exception>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <utility>
 
namespace aria::async {
 
template<typename T = void>
class Task;
 
namespace detail {
 
template<typename T>
struct TaskPromiseBase {
    std::coroutine_handle<> continuation;
    std::exception_ptr exception;
 
    // Detached tasks run through final suspend so the compiler destroys
    // their frames. Owned tasks remain suspended for their Task/awaiter.
    bool detached = false;
 
    auto initial_suspend() noexcept { return std::suspend_always{}; }
 
    struct FinalAwaiter {
        bool detached;
 
        bool await_ready() const noexcept { return detached; }
 
        // Templated to accept any coroutine_handle whose promise inherits TaskPromiseBase<T>
        template<typename Promise>
        std::coroutine_handle<> await_suspend(std::coroutine_handle<Promise> h) noexcept {
            TaskPromiseBase<T>& base = h.promise();
            return base.continuation ? base.continuation : std::noop_coroutine();
        }
 
        void await_resume() const noexcept {}
    };
 
    auto final_suspend() noexcept { return FinalAwaiter{detached}; }
 
    void unhandled_exception() noexcept { exception = std::current_exception(); }
};
 
template<typename T>
struct TaskPromise : TaskPromiseBase<T> {
    using value_type = T;
    std::optional<T> value;
 
    Task<T> get_return_object() noexcept;
 
    template<typename U>
    void return_value(U&& v) noexcept(std::is_nothrow_constructible_v<T, U>) {
        value.emplace(std::forward<U>(v));
    }
};
 
template<>
struct TaskPromise<void> : TaskPromiseBase<void> {
    using value_type = void;
    Task<void> get_return_object() noexcept;
    void return_void() noexcept {}
};
 
}  // namespace detail
 
template<typename T>
class [[nodiscard]] Task {
public:
    using promise_type = detail::TaskPromise<T>;
    using handle_type = std::coroutine_handle<promise_type>;
 
    Task() noexcept = default;
    explicit Task(handle_type h) noexcept : handle_(h) {}
 
    Task(const Task&) = delete;
    Task& operator=(const Task&) = delete;
 
    Task(Task&& o) noexcept : handle_(std::exchange(o.handle_, {})) {}
    Task& operator=(Task&& o) noexcept {
        if (this != &o) {
            if (handle_) handle_.destroy();
            handle_ = std::exchange(o.handle_, {});
        }
        return *this;
    }
 
    ~Task() {
        if (handle_) handle_.destroy();
    }
 
    [[nodiscard]] bool done() const noexcept { return handle_ && handle_.done(); }

    auto operator co_await() && noexcept {
        struct Awaiter {
            handle_type h;
 
            bool await_ready() const noexcept { return !h || h.done(); }
 
            std::coroutine_handle<> await_suspend(std::coroutine_handle<> caller) noexcept {
                h.promise().continuation = caller;
                return h;
            }
 
            T await_resume() {
                if (!h) { throw std::runtime_error("Task: empty handle"); }
                auto& p = h.promise();
                if (p.exception) std::rethrow_exception(p.exception);
                if constexpr (!std::is_void_v<T>) {
                    return std::move(*p.value);
                }
            }
 
            // Awaiter takes ownership: destroy when done.
            ~Awaiter() { if (h) h.destroy(); }
 
            Awaiter(handle_type hh) noexcept : h(hh) {}
            Awaiter(Awaiter&& o) noexcept : h(std::exchange(o.h, {})) {}
            Awaiter& operator=(Awaiter&&) = delete;
            Awaiter(const Awaiter&) = delete;
            Awaiter& operator=(const Awaiter&) = delete;
        };
        // Transfer ownership from the temporary `*this` to the Awaiter.
        return Awaiter{std::exchange(handle_, {})};
    }

    void start() {
        if (handle_ && !handle_.done()) handle_.resume();
    }

    void start_detached() && {
        if (!handle_) return;
        auto h = std::exchange(handle_, {});
        if (h.done()) {
            h.destroy();
            return;
        }
        h.promise().detached = true;
        // Completion may release the frame before resume() returns.
        // If the body suspends, its eventual resumer drives completion.
        // Never inspect h after resuming it.
        h.resume();
    }

    void start_detached_() { std::move(*this).start_detached(); }

    T blocking_get() {
        if (!handle_) throw std::runtime_error("Task: empty handle");
        if (!handle_.done()) handle_.resume();
        if (!handle_.done()) throw std::runtime_error("Task: did not complete synchronously");
        auto& p = handle_.promise();
        if (p.exception) std::rethrow_exception(p.exception);
        if constexpr (!std::is_void_v<T>) {
            return std::move(*p.value);
        }
    }
 
private:
    handle_type handle_{};
};
namespace detail {
 
template<typename T>
Task<T> TaskPromise<T>::get_return_object() noexcept {
    return Task<T>{std::coroutine_handle<TaskPromise<T>>::from_promise(*this)};
}
 
inline Task<void> TaskPromise<void>::get_return_object() noexcept {
    return Task<void>{std::coroutine_handle<TaskPromise<void>>::from_promise(*this)};
}
 
}  // namespace detail
 
}  // namespace aria::async