task.h

#ifndef ASYNCIO_TASK_H
#define ASYNCIO_TASK_H
#include <asyncio/handle.h>
#include <asyncio/event_loop.h>
#include <asyncio/result.h>
#include <asyncio/concept/promise.h>
#include <coroutine>
#include <cassert>
#include <variant>
#include <memory>

namespace ASYNCIO_NS {
struct NoWaitAtInitialSuspend {};
inline constexpr NoWaitAtInitialSuspend no_wait_at_initial_suspend;

template<typename R = void> // R为Task的返回类型
struct Task: private NonCopyable {  // Task作为协程函数的返回值（Task必须有一个promise_type，必须重载co_await）
    struct promise_type;
    using coro_handle = std::coroutine_handle<promise_type>;

    template<concepts::Future>
    friend struct ScheduledTask;

    explicit Task(coro_handle h) noexcept: handle_(h) {}

    Task(Task&& t) noexcept
        : handle_(std::exchange(t.handle_, {})) {}

    ~Task() { destroy(); }

    // 获取Task返回值
    decltype(auto) get_result() & {
        return handle_.promise().result();
    }

    decltype(auto) get_result() && {
        return std::move(handle_.promise()).result();
    }

    // co_await返回值的基类
    struct AwaiterBase {
        constexpr bool await_ready() {
            if (self_coro_) [[likely]]  // 有句柄，就检查句柄是否执行完成
            { return self_coro_.done(); }   
            return true;
        }

        template<typename Promise>
        void await_suspend(std::coroutine_handle<Promise> resumer) const noexcept {
            assert(! self_coro_.promise().continuation_);
            // 将调用co_await Task的协程状态设置为SUSPEND，因为需要等待Task执行返回
            resumer.promise().set_state(Handle::SUSPEND);
            // 当前任务执行完或者挂起，就会执行continuation_，将continuation_指定为为调用Task的协程的句柄
            self_coro_.promise().continuation_ = &resumer.promise();    
            self_coro_.promise().schedule();    // 将当前协程加入句柄
        }

        coro_handle self_coro_ {};  // Task所属协程的句柄
    };

    // 重载co_await运算符（左值引用调用）
    auto operator co_await() const & noexcept {
        struct Awaiter: AwaiterBase {
            decltype(auto) await_resume() const {
                if (! AwaiterBase::self_coro_) [[unlikely]]
                { throw InvalidFuture{}; }
                return AwaiterBase::self_coro_.promise().result();  // 返回结果
            }
        };
        return Awaiter {handle_};   // Awaiter的协程句柄就是Task所属协程的句柄
    }

    // 重载co_await运算符（右值引用调用）
    auto operator co_await() const && noexcept {
        struct Awaiter: AwaiterBase {
            decltype(auto) await_resume() const {
                if (! AwaiterBase::self_coro_) [[unlikely]]
                { throw InvalidFuture{}; }
                return std::move(AwaiterBase::self_coro_.promise()).result();   // 返回结果
            }
        };
        return Awaiter {handle_};   // Awaiter的协程句柄就是Task所属协程的句柄
    }

    struct promise_type: CoroHandle, Result<R> {    // promise_type既是协程句柄，也是结果
        promise_type() = default;

        // 
        template<typename... Args> // from free function
        promise_type(NoWaitAtInitialSuspend, Args&&...): wait_at_initial_suspend_{false} { }

        // 
        template<typename Obj, typename... Args> // from member function
        promise_type(Obj&&, NoWaitAtInitialSuspend, Args&&...): wait_at_initial_suspend_{false} { }

        // 
        auto initial_suspend() noexcept {
            struct InitialSuspendAwaiter {
                constexpr bool await_ready() const noexcept { return !wait_at_initial_suspend_; }   // 协程初始需要挂起
                constexpr void await_suspend(std::coroutine_handle<>) const noexcept {}
                constexpr void await_resume() const noexcept {}
                const bool wait_at_initial_suspend_{true};
            };
            return InitialSuspendAwaiter{wait_at_initial_suspend_};
        }

        struct FinalAwaiter {
            constexpr bool await_ready() const noexcept { return false; }   // 协程结束需要挂起
            
            template<typename Promise>
            constexpr void await_suspend(std::coroutine_handle<Promise> h) const noexcept { 
                if (auto cont = h.promise().continuation_) {    // 将调用co_await Task的协程的句柄加入调度
                    get_event_loop().call_soon(*cont);
                }
            }

            constexpr void await_resume() const noexcept {}
        };

        auto final_suspend() noexcept {
            return FinalAwaiter {};
        }

        Task get_return_object() noexcept {
            return Task{coro_handle::from_promise(*this)};
        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////////////////
        // 这些是CoroHandle的函数实现
        void run() final {  // run就是调用当前协程句柄
            coro_handle::from_promise(*this).resume();
        }
        
        const std::source_location& get_frame_info() const final { return frame_info_; }
        
        void dump_backtrace(size_t depth = 0) const final {
            fmt::print("[{}] {}\n", depth, frame_name());
            if (continuation_) { continuation_->dump_backtrace(depth + 1);  // 调用上一次的函数 }
            else { fmt::print("\n"); }
        }
        ///////////////////////////////////////////////////////////////////////////////////////////////////////////////

        const bool wait_at_initial_suspend_ {true}; // 是否在协程开始时挂起
        CoroHandle* continuation_ {};               // 记录调用Task的协程的句柄
        std::source_location frame_info_{};         // 代码所在位置
    };

    bool valid() const { return handle_ != nullptr; }   // 任务是否合法：句柄是否存在

    bool done() const { return handle_.done(); }    // 任务是否完成：句柄是否执行完成

private:
    void destroy() {
        if (auto handle = std::exchange(handle_, nullptr)) {
            handle.promise().cancel();
            handle.destroy();
        }
    }

private:
    coro_handle handle_;    // 协程句柄
};

static_assert(concepts::Promise<Task<>::promise_type>);
static_assert(concepts::Future<Task<>>);
}
#endif // ASYNCIO_TASK_H