when_all.cppm

export module mcpplibs.cmp:when_all;

import std;
import :task;

namespace mcpplibs::cmp::detail {

template<typename T>
using WhenAllValue = std::conditional_t<
    std::same_as<T, void>,
    std::monostate,
    T>;

struct JoinStartNode final {
    std::coroutine_handle<> coroutine_ {};
    JoinStartNode* next_ {};
};

class JoinStartQueue final {
private:
    static inline thread_local JoinStartNode* head_ {};
    static inline thread_local bool draining_ { false };

public:
    class Scope final {
    private:
        JoinStartNode* savedHead_ { head_ };
        bool savedDraining_ { draining_ };

    public:
        Scope() noexcept {
            head_ = nullptr;
            draining_ = false;
        }

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

        ~Scope() {
            // 同步恢复边界必须排空自己的启动队列，外层待启动节点仍由外层持有
            if (head_ || draining_) {
                std::terminate();
            }

            head_ = savedHead_;
            draining_ = savedDraining_;
        }
    };

    static void prepend(JoinStartNode* first, JoinStartNode* last) noexcept {
        // 嵌套子树排在待启动的兄弟之前，保持深度优先的输入启动顺序
        last->next_ = head_;
        head_ = first;
    }

    static void drain() noexcept {
        if (draining_) {
            return;
        }

        draining_ = true;

        while (head_) {
            auto* const node = head_;
            head_ = node->next_;
            const auto coroutine = node->coroutine_;

            // 恢复后整个聚合可能已完成，节点所属帧也可能已销毁
            coroutine.resume();
        }

        draining_ = false;
    }
};

class JoinCounter final {
private:
    // 多出的 1 是注册哨兵，防止同步子任务过早恢复父协程
    std::atomic<std::size_t> remaining_;
    std::coroutine_handle<> continuation_ {};

public:
    explicit JoinCounter(std::size_t count) noexcept
        : remaining_ { count + 1 } {}

    [[nodiscard]] bool try_await(
        std::coroutine_handle<> continuation) noexcept {
        continuation_ = continuation;
        return remaining_.fetch_sub(1, std::memory_order_acq_rel) > 1;
    }

    [[nodiscard]] std::coroutine_handle<> notify_child_completed() noexcept {
        if (remaining_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
            return continuation_;
        }

        return std::noop_coroutine();
    }
};

template<typename T>
class JoinResultStorage {
private:
    std::optional<T> result_ {};

public:
    template<typename U>
    requires std::constructible_from<T, U&&>
    void return_value(U&& value)
        noexcept(std::is_nothrow_constructible_v<T, U&&>) {
        result_.emplace(std::forward<U>(value));
    }

    T take_result() {
        if (!result_) {
            throw std::logic_error {
                "when_all child completed without a result"
            };
        }

        return std::move(*result_);
    }
};

template<>
class JoinResultStorage<void> {
public:
    constexpr void return_void() const noexcept {}

    [[nodiscard]] constexpr std::monostate take_result() const noexcept {
        return {};
    }
};

template<typename T>
class JoinTask final {
public:
    struct promise_type : JoinResultStorage<T> {
        JoinStartNode startNode_ {};
        JoinCounter* counter_ {};
        std::exception_ptr exception_ {};

        [[nodiscard]] JoinTask get_return_object() noexcept;

        [[nodiscard]] constexpr std::suspend_always initial_suspend() const noexcept {
            return {};
        }

        class FinalAwaiter final {
        public:
            [[nodiscard]] constexpr bool await_ready() const noexcept {
                return false;
            }

            [[nodiscard]] std::coroutine_handle<> await_suspend(
                std::coroutine_handle<promise_type> coroutine) const noexcept {
                return coroutine.promise().counter_->notify_child_completed();
            }

            constexpr void await_resume() const noexcept {}
        };

        [[nodiscard]] constexpr FinalAwaiter final_suspend() const noexcept {
            return {};
        }

        void unhandled_exception() noexcept {
            exception_ = std::current_exception();
        }
    };

private:
    using Handle = std::coroutine_handle<promise_type>;

    Handle coroutine_ {};

    explicit JoinTask(Handle coroutine) noexcept
        : coroutine_ { coroutine } {}

public:
    JoinTask(const JoinTask&) = delete;
    JoinTask& operator=(const JoinTask&) = delete;

    JoinTask(JoinTask&& other) noexcept
        : coroutine_ { std::exchange(other.coroutine_, {}) } {}

    JoinTask& operator=(JoinTask&&) = delete;

    ~JoinTask() {
        if (coroutine_) {
            coroutine_.destroy();
        }
    }

    void prepare_start(
        JoinCounter& counter,
        JoinStartNode*& first,
        JoinStartNode*& last) noexcept {
        coroutine_.promise().counter_ = &counter;
        auto* const node = &coroutine_.promise().startNode_;
        node->coroutine_ = coroutine_;

        if (last) {
            last->next_ = node;
        } else {
            first = node;
        }

        last = node;
    }

    void rethrow_if_exception() const {
        if (coroutine_.promise().exception_) {
            std::rethrow_exception(coroutine_.promise().exception_);
        }
    }

    WhenAllValue<T> take_result() {
        return coroutine_.promise().take_result();
    }
};

template<typename T>
JoinTask<T> JoinTask<T>::promise_type::get_return_object() noexcept {
    return JoinTask {
        std::coroutine_handle<promise_type>::from_promise(*this)
    };
}

template<typename T>
JoinTask<T> make_join_task(Task<T> task) {
    if constexpr (std::same_as<T, void>) {
        co_await std::move(task);
        co_return;
    } else {
        co_return co_await std::move(task);
    }
}

template<typename... Results>
class WhenAllAwaiter final {
private:
    JoinCounter counter_ { sizeof...(Results) };
    std::tuple<JoinTask<Results>...> tasks_;

public:
    explicit WhenAllAwaiter(Task<Results>... tasks)
        : tasks_ { make_join_task(std::move(tasks))... } {}

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

    [[nodiscard]] static constexpr bool await_ready() noexcept {
        return sizeof...(Results) == 0;
    }

    void await_suspend(
        std::coroutine_handle<> continuation) noexcept {
        JoinStartNode* first {};
        JoinStartNode* last {};
        std::apply(
            [this, &first, &last](auto&... tasks) noexcept {
                (tasks.prepare_start(counter_, first, last), ...);
            },
            tasks_);

        JoinStartQueue::prepend(first, last);
        static_cast<void>(counter_.try_await(continuation));

        // 子任务从此处开始运行，后续仅访问线程局部队列
        JoinStartQueue::drain();
    }

    std::tuple<WhenAllValue<Results>...> await_resume() {
        // 先按参数顺序检查异常，避免先移动部分结果再发现子任务失败
        std::apply(
            [](const auto&... tasks) {
                (tasks.rethrow_if_exception(), ...);
            },
            tasks_);

        return std::apply(
            [](auto&... tasks) {
                return std::tuple<WhenAllValue<Results>...> {
                    tasks.take_result()...
                };
            },
            tasks_);
    }
};

template<typename Result>
class RangeWhenAllAwaiter final {
private:
    JoinCounter counter_;
    std::vector<JoinTask<Result>> tasks_ {};

public:
    explicit RangeWhenAllAwaiter(std::vector<Task<Result>> tasks)
        : counter_ { tasks.size() } {
        // 先取得全部容器空间，失败时还没有子任务开始运行
        tasks_.reserve(tasks.size());

        for (auto& task : tasks) {
            tasks_.emplace_back(make_join_task(std::move(task)));
        }
    }

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

    [[nodiscard]] bool await_ready() const noexcept {
        return tasks_.empty();
    }

    void await_suspend(
        std::coroutine_handle<> continuation) noexcept {
        JoinStartNode* first {};
        JoinStartNode* last {};

        for (auto& task : tasks_) {
            task.prepare_start(counter_, first, last);
        }

        JoinStartQueue::prepend(first, last);
        static_cast<void>(counter_.try_await(continuation));

        // 子任务从此处开始运行，后续仅访问线程局部队列
        JoinStartQueue::drain();
    }

    std::vector<WhenAllValue<Result>> await_resume() {
        for (const auto& task : tasks_) {
            task.rethrow_if_exception();
        }

        std::vector<WhenAllValue<Result>> results {};
        results.reserve(tasks_.size());

        for (auto& task : tasks_) {
            results.emplace_back(task.take_result());
        }

        return results;
    }
};

}  // namespace mcpplibs::cmp::detail

export namespace mcpplibs::cmp {

template<typename... Results>
[[nodiscard]] Task<std::tuple<detail::WhenAllValue<Results>...>> when_all(
    Task<Results>... tasks) {
    co_return co_await detail::WhenAllAwaiter<Results...> {
        std::move(tasks)...
    };
}

template<typename Result>
[[nodiscard]] Task<std::vector<detail::WhenAllValue<Result>>> when_all(
    std::vector<Task<Result>> tasks) {
    co_return co_await detail::RangeWhenAllAwaiter<Result> {
        std::move(tasks)
    };
}

}  // namespace mcpplibs::cmp