c++20制作简易协程库4

auto set() noexcept -> void; // 普通调用
auto wait() noexcept -> awaiter; // 协程调用，awaiter 的 await_resume 返回 void

template<typename value_type>
auto set(value_type&& value) noexcept -> void; // 普通调用
auto wait() noexcept -> awaiter; // 协程调用，awaiter 的 await_resume 返回 set 设置的值

event<int> ev;
task<> set_func() {
  // codes...
  ev.set(number);
}
task<> wait_func() {
  auto number = co_await ev.wait();
  // codes ...
}

// expected: 期望的原子变量的状态
// desired: 替换为的状态
// 语义: 如果当前原子变量的状态为expected, 就尝试替换为desired, 替换成功返回true, 否则返回false
// 如果返回false, expected会被改为当前原子变量的新值
bool compare_exchange_strong(T& expected, T desired,
                             memory_order success_order,
                             memory_order failure_order) noexcept;

bool compare_exchange_weak(T& expected, T desired,
                           memory_order success_order,
                           memory_order failure_order) noexcept;

// 内存序默认memory_order_seq_cst
// 原子读
T load(memory_order order);
// 原子写
void store(T val, memory_order order);
// 原子交换, 返回原子变量中原先存的值
T exchange(T val, memory_order order);

class event_base {
public:
    struct base_awaiter {
        friend class event_base;
        base_awaiter(event_base& event) noexcept : m_event(event), m_ctx(local_context()) {}
        auto await_ready() noexcept -> bool {
            m_ctx.register_wait();
            return m_event.is_set();
        }
        auto await_suspend(std::coroutine_handle<> handle) noexcept -> bool {
            m_handle = handle;
            void *old_state = m_event.m_state.load(std::memory_order_acquire);
            while (true) {
                // 竞争情况: 准备队列时, event被set了
                if (old_state == reinterpret_cast<void*>(kSet)) {
                    return false;
                }
                // 正常情况: 把自己加入等待者链表
                m_next = (old_state == kUnsetNowaiter) ? nullptr : static_cast<base_awaiter*>(old_state);
                if (m_event.m_state.compare_exchange_strong(
                    old_state, 
                    this,
                    std::memory_order_acq_rel,
                    std::memory_order_relaxed
                )) {
                    return true;
                }
            }
        }
        auto await_resume() noexcept -> void{
            m_ctx.unregister_wait();
        }
    protected:
        event_base& m_event;
        context& m_ctx;
        std::coroutine_handle<> m_handle{nullptr};
        base_awaiter* m_next{nullptr};
    };
    friend struct base_awaiter;
public:
    event_base() noexcept : m_state(kUnsetNowaiter) {};
    ~event_base() noexcept {}

    auto is_set() const noexcept -> bool {
        return m_state.load(std::memory_order_acquire) == reinterpret_cast<void*>(kSet);
    }

    auto reset() -> void {
        void *expected_state = reinterpret_cast<void*>(kSet);
        m_state.compare_exchange_strong(expected_state, kUnsetNowaiter, std::memory_order_acq_rel);
    }

protected:
    void resume_awaiters() noexcept {
        void* old_state = m_state.exchange(reinterpret_cast<void*>(kSet), std::memory_order_acq_rel);
        if (old_state != reinterpret_cast<void*>(kSet)) {
            // 链表挂载等待的协程, 唤醒
            // base_awaiter* awaiter = static_cast<base_awaiter*>(old_state);
            base_awaiter* head = static_cast<base_awaiter*>(old_state);
            while (head != nullptr) {
                auto* next_awaiter = head->m_next;
                head->m_ctx.submit_task(head->m_handle);
                head = next_awaiter;
            }
        }
    }
protected:
    static constexpr void* kUnsetNowaiter = nullptr;
    static constexpr std::uintptr_t kSet = 1;
    std::atomic<void*> m_state;
};

auto try_lock() noexcept -> bool; // 普通调用
auto lock() noexcept -> awaiter; // 协程调用，awaiter 的 await_resume 返回 void
auto unlock() noexcept -> void; // 普通调用
auto lock_guard() noexcept -> awaiter; // 协程调用，awaiter 的 await_resume 返回 lock_guard

class mutex
{
public:
    struct mutex_awaiter {
        std::coroutine_handle<> handle;
        mutex& mtx;
        context& ctx;
        mutex_awaiter* next;

        explicit mutex_awaiter(mutex& mutex) noexcept : mtx(mutex), 
            ctx(local_context()), handle(nullptr), next(nullptr) {}

        virtual ~mutex_awaiter() {}    

        auto await_ready() noexcept -> bool {
            ctx.register_wait(); 
            void* expected = kUNLOCKED;
            return mtx.m_state.compare_exchange_strong(
                expected,
                 kLOCKED_NO_WATIER,
                 std::memory_order_acquire);
        }

        auto await_suspend(std::coroutine_handle<> awaiting_handle) noexcept -> bool {
            handle = awaiting_handle;
            return add_waiting_list();
        }

        auto await_resume() noexcept -> void {
            ctx.unregister_wait();
        }

        auto add_waiting_list() noexcept -> bool {
            void* old_state = mtx.m_state.load();
            while (true) {
                // 竞争情况：在我们准备排队时，锁被释放了。
                // 我们可以再次尝试获取锁，避免不必要的挂起。
                if (old_state == kUNLOCKED) {
                    if (mtx.m_state.compare_exchange_weak(
                            old_state, 
                            kLOCKED_NO_WATIER)) {
                        return false; // 成功获取锁，不要挂起。
                    }
                    // CAS 失败，意味着 old_state 被其他线程改变了，循环会用新值重试。
                    continue;
                }

                // 正常情况：锁被持有，我们需要把自己加入链表。
                // old_state 要么是 kLockedNoWaiters，要么是链表头。
                next = (old_state == kLOCKED_NO_WATIER) ? nullptr : static_cast<mutex_awaiter*>(old_state);

                // release 内存序确保了在我们挂起协程之前的代码，对唤醒我们的那个线程是可见的。
                if (mtx.m_state.compare_exchange_weak(
                        old_state, 
                        this)) {
                    return true; // 成功把自己加入队列，现在可以安全挂起了。
                }
                // CAS 失败，循环会用新的 old_state 值重试。
            }
        }

        virtual auto resume_coro() noexcept -> void {
            ctx.submit_task(handle);
        }
    };

    struct guard_awaiter : mutex_awaiter
    {
        using mutex_awaiter::mutex_awaiter;
        auto await_resume() -> detail::lock_guard<mutex> { 
            ctx.unregister_wait();
            return detail::lock_guard<mutex>(mtx); 
        }
    };
public:
    mutex() noexcept : m_state(kUNLOCKED), m_awaiter_head(nullptr) {}
    ~mutex() noexcept {
        assert(m_state.load() == kUNLOCKED);
    }
    mutex(const mutex&) = delete;
    mutex& operator=(const mutex&) = delete;

    auto try_lock() noexcept -> bool { 
        void* expected = kUNLOCKED;
        return m_state.compare_exchange_strong(
            expected, 
            kLOCKED_NO_WATIER);
    }


    auto lock() noexcept -> mutex_awaiter {
        return mutex_awaiter{*this}; 
    };

    auto unlock() noexcept -> void {
        assert(m_state != kUNLOCKED && "can't unlock unlocked mutex!");
        if (m_awaiter_head == nullptr) {
            // kLOCKED_NO_WAITER
            void* expected = kLOCKED_NO_WATIER;
            if (m_state.compare_exchange_strong(
                expected,
                kUNLOCKED,
                std::memory_order_acq_rel,
                std::memory_order_acq_rel)) {
                    // 当前没有等待者
                    return;
                }
            // HAVE WAITER
            void* waiting_list = m_state.exchange(kLOCKED_NO_WATIER, std::memory_order_acq_rel);
            assert(waiting_list != kUNLOCKED && "can't unlock unlocked mutex!");
            m_awaiter_head = static_cast<mutex_awaiter*>(waiting_list);
            // reverse the list, to implement FIFO
            mutex_awaiter* awaiter = m_awaiter_head;
            m_awaiter_head = nullptr;
            while (awaiter != nullptr) {
                auto next_awaiter = awaiter->next;
                awaiter->next = m_awaiter_head;
                m_awaiter_head = awaiter;
                awaiter = next_awaiter;
            }
        }
        // m_state is LOCKED_HAVE_WAITER
        mutex_awaiter* to_resume = m_awaiter_head;
        m_awaiter_head = to_resume->next;
        to_resume->resume_coro();
    };

    auto lock_guard() noexcept -> guard_awaiter { return guard_awaiter{*this};};
private:
    inline static void* kUNLOCKED = nullptr;
    inline static void* kLOCKED_NO_WATIER = reinterpret_cast<void*>(1);
    std::atomic<void*> m_state;
    mutex_awaiter* m_awaiter_head;
};