C++20 Coroutine Task

@Louis-Langholtz This has been a while, but I certainly can do that.

After running a first attempt at fixing @rafoo's code, I've understood the real source of his code's problem. It does a co_await std::suspend_always{}; in two places. This is not the proper usage of std::suspend_always. The problem is that this call returns control to main, but main needs to restart the proper handle, which, without his stack of handles, is lost.

So instead, one needs to have an awaitable that schedules the handle of the awaiting coroutine for future resumption. This is done below by the object suspend{}. The variable execution_handle models a simple executor, main just resumes that handle until it is done. Note however, that the handle will be changed by any coroutine that co_await suspend{};. The rest of the flow is handled by the task continuation handle, allowing to return execution to the coroutine that awaits the task's result.

In a real code, suspend{} could pass the handle to a real multi-threaded executor. Or it could have an asynchronous I/O completion do the re-scheduling once the I/O has completed.

#include <iostream>
#include <coroutine>
#include <memory>
#include <optional>

namespace coroutines = std;

template<typename T = void>
struct [[nodiscard]] task {
    using value_type = T;
    struct task_promise;
    using promise_type = task_promise;

    explicit task(std::coroutine_handle<promise_type> handle) : my_handle(handle) {}

    task(task &&t) noexcept: my_handle(t.my_handle) { t.my_handle = nullptr; }

    /// Disable copy construction/assignment.
    task(const task &) = delete;

    task &operator=(const task &) = delete;

    ~task() {
        if (my_handle)
            my_handle.destroy();
    }

    task &operator=(task &&other) noexcept {
        if (std::addressof(other) != this) {
            if (my_handle) {
                my_handle.destroy();
            }

            my_handle = other.my_handle;
            other.my_handle = nullptr;
        }
        return *this;
    }

    /** \brief Provides for co_awaiting a task. */
    auto operator
    co_await() noexcept
    {
        struct awaiter {
            bool await_ready() const noexcept { return !coro_handle || coro_handle.done(); }

            std::coroutine_handle<>
            await_suspend(std::coroutine_handle<> awaiting_coroutine) noexcept {
                coro_handle.promise().set_continuation(awaiting_coroutine);
                return coro_handle;
            }

            T await_resume() noexcept {
                if constexpr (std::is_same_v<T, void>)
                    return;
                else
                    return std::move(coro_handle.promise().data.value());
            }

            std::coroutine_handle<promise_type> coro_handle;
        };
        return awaiter{my_handle};
    }

    auto get_handle() {
        return my_handle;
    }
private:
    std::coroutine_handle<promise_type> my_handle;
};

/** \brief Component of the promise handling the result.
 *
 * @tparam T Result type
 *
 * \details The C++20 coroutines' promises need to expose either a return_void or
 * return_value function, whether they return an actual value or a void.
 * They cannot expose both. Many implementations specialize the full promise type.
 * Here we only specialize the return value handling.
 */
template<typename T>
struct promise_setter {
    template<typename U>
    void return_value(U &&u) {
        data.template emplace(std::forward<U>(u));
    }

    std::optional<T> data;
};

/** \brief Specialization of promise_setter for void returning tasks. */
template<>
struct promise_setter<void> {
    void return_void() {}
};

/** \brief Promise type associated with a task
 *
 * @tparam T the return type of the task.
 */
template<typename T>
struct task<T>::task_promise : public promise_setter<T> {

    /// \brief Invoked when we first enter a coroutine, build the task object with the coroutine
    /// handle.
    task get_return_object() noexcept {
        return task{std::coroutine_handle<task_promise>::from_promise(*this)};
    };

    auto initial_suspend() const noexcept {
        return std::suspend_always{};
    }

    /// \brief Get an awaiter with the continuation if any.
    /// \details If T is void, co_await final_suspend(); will destroy the coroutine frame.
    auto final_suspend() const noexcept {
        struct awaiter {
            // Return false here to return control to the thread's event loop. Remember that we're
            // running on some async thread at this point.
            /// \brief The current coroutine is done, so it is not ready to restart.
            /// \details await_suspend will be called next allowing to return the handle of the
            /// continuation.
            bool await_ready() const noexcept { return false; }

            void await_resume() const noexcept {}

            /** \brief Suspension method for the ending co-routine.
             * @param h Handle of the current and ending co-routine.
             * @return The handle of the coroutine that continues the work if any.
             */
            std::coroutine_handle<>
            await_suspend(std::coroutine_handle<task::task_promise> h) noexcept {
                return h.promise().continuation;
            }
        };
        return awaiter{};
    }

    void unhandled_exception() noexcept { std::terminate(); }

    void set_continuation(std::coroutine_handle<> handle) { continuation = handle; }

private:
    std::coroutine_handle<> continuation = std::noop_coroutine();
};

/* Simple mock-up of an executor. main() will continuously resume execution_handle.
  A coroutine can be re-scheduled for execution by writing its handle to this variable. */
std::coroutine_handle<> execution_handle;

/// An awaitable for suspending a coroutine and, in this demo, immediately
/// scheduling it to be resumed next.
struct suspend {
    /// \brief The current coroutine wants to pause and is not ready to restart.
    /// \details await_suspend will be called next allowing to schedule for future execution.
    bool await_ready() const noexcept { return false; }

    void await_resume() const noexcept {}

    /** \brief Suspension method for the ending co-routine.
     * @param h Handle of the current and ending co-routine.
     * @return The handle of the coroutine that continues the work if any.
     */
    auto
    await_suspend(std::coroutine_handle<> h) noexcept {
        // Schedule for resumption
        execution_handle = h;
        // Pause the execution
        return std::noop_coroutine();
    }
};


task<double> one()
{
    std::cout << "one takes a coffee" << std::endl;
    // Suspend (and reschedule)
    co_await suspend{};
    co_return 1;
}

task<> noop()
{
    std::cout << "the noop does something" << std::endl;
    co_return;
}

task<short> two()
{
    co_return co_await one() + co_await one();
}

task<int> foo()
{
    int i = 0;
    i += co_await one();
    std::cout << "foo just run one() " << i << std::endl;

    i += co_await two();

    std::cout << "foo just run two() " << i << std::endl;

    std::cout << "foo makes a pause after his work" << std::endl;

    // Suspend (and reschedule)
    co_await suspend{};

    co_await noop();

    co_return i;
}

int main()
{
    // Build the task foo() but does not start any work
    auto tsk = foo();
   execution_handle = tsk.get_handle();

   for(int i = 0; !execution_handle.done(); ++i) {
       std::cout << "Resume # " << i << std::endl;
       execution_handle.resume();
   }

    std::cout << "Result: " << tsk.get_handle().promise().data.value() << std::endl;
    return 0;
}

```