event.h

#pragma once
#include <asyncpp/detail/std_import.h>
#include <asyncpp/dispatcher.h>
#include <atomic>
#include <cassert>
 
namespace asyncpp {
    class single_consumer_event {
    public:
        explicit constexpr single_consumer_event(bool set_initially = false) noexcept
            : m_state(set_initially ? this : nullptr) {}
#ifndef NDEBUG
        ~single_consumer_event() noexcept { assert(!is_awaited()); }
#endif
        [[nodiscard]] bool is_set() const noexcept { return m_state.load(std::memory_order::acquire) == this; }
 
        [[nodiscard]] bool is_awaited() const noexcept {
            auto val = m_state.load(std::memory_order::acquire);
            return val != nullptr && val != this;
        }
 
        bool set(dispatcher* resume_dispatcher = nullptr) noexcept {
            auto state = m_state.exchange(this, std::memory_order::acq_rel);
            if (state != nullptr && state != this) {
                auto await = static_cast<awaiter*>(state);
                assert(await->m_parent == this);
                assert(await->m_handle);
                if (await->m_dispatcher != nullptr) {
                    await->m_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else if (resume_dispatcher != nullptr) {
                    resume_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else {
                    await->m_handle.resume();
                }
                return true;
            }
            return false;
        }
 
        void reset() noexcept {
            void* old_state = this;
            m_state.compare_exchange_strong(old_state, nullptr, std::memory_order::relaxed);
        }
 
        [[nodiscard]] auto operator co_await() noexcept { return awaiter{this, dispatcher::current()}; }
 
        [[nodiscard]] constexpr auto wait(dispatcher* resume_dispatcher = nullptr) noexcept {
            return awaiter{this, resume_dispatcher};
        }
 
    private:
        /* nullptr => unset
         * this => set
         * x => awaiter*
         */
        std::atomic<void*> m_state;
 
        struct [[nodiscard]] awaiter {
            explicit constexpr awaiter(single_consumer_event* parent, dispatcher* dispatcher) noexcept
                : m_parent(parent), m_dispatcher(dispatcher) {}
            [[nodiscard]] bool await_ready() const noexcept { return m_parent->is_set(); }
            [[nodiscard]] bool await_suspend(coroutine_handle<> hdl) noexcept {
                m_handle = hdl;
                void* old_state = nullptr;
                // If the current state is unset set it to this
                bool was_equal = m_parent->m_state.compare_exchange_strong(old_state, this, std::memory_order::release,
                                                                           std::memory_order::acquire);
                // If the state was not unset it has to be set,
                // otherwise we have a concurrent await, which is not supported
                assert(was_equal || old_state == m_parent);
                return was_equal;
            }
            constexpr void await_resume() const noexcept {}
 
            single_consumer_event* m_parent;
            dispatcher* m_dispatcher;
            coroutine_handle<> m_handle{};
        };
    };
 
    class single_consumer_auto_reset_event {
    public:
        explicit constexpr single_consumer_auto_reset_event(bool set_initially = false) noexcept
            : m_state(set_initially ? this : nullptr) {}
#ifndef NDEBUG
        ~single_consumer_auto_reset_event() noexcept { assert(!is_awaited()); }
#endif
        [[nodiscard]] bool is_set() const noexcept { return m_state.load(std::memory_order::acquire) == this; }
 
        [[nodiscard]] bool is_awaited() const noexcept {
            auto ptr = m_state.load(std::memory_order::acquire);
            return ptr != nullptr && ptr != this;
        }
 
        bool set(dispatcher* resume_dispatcher = nullptr) noexcept {
            auto state = m_state.exchange(this, std::memory_order::release);
            if (state != nullptr && state != this) {
                auto await = static_cast<awaiter*>(state);
 
                // Only modify the state if it has not been changed in between
                state = this;
                m_state.compare_exchange_strong(state, nullptr, std::memory_order::acq_rel);
 
                assert(await->m_parent == this);
                assert(await->m_handle);
                if (await->m_dispatcher != nullptr) {
                    await->m_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else if (resume_dispatcher != nullptr) {
                    resume_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else {
                    await->m_handle.resume();
                }
                return true;
            }
            return false;
        }
 
        void reset() noexcept {
            void* old_state = this;
            m_state.compare_exchange_strong(old_state, nullptr, std::memory_order::relaxed);
        }
 
        [[nodiscard]] auto operator co_await() noexcept { return awaiter{this, dispatcher::current()}; }
 
        [[nodiscard]] constexpr auto wait(dispatcher* resume_dispatcher = nullptr) noexcept {
            return awaiter{this, resume_dispatcher};
        }
 
    private:
        /* nullptr => unset
         * this => set
         * x => awaiter*
         */
        std::atomic<void*> m_state;
 
        struct [[nodiscard]] awaiter {
            explicit constexpr awaiter(single_consumer_auto_reset_event* parent, dispatcher* dispatcher) noexcept
                : m_parent(parent), m_dispatcher(dispatcher) {}
            [[nodiscard]] constexpr bool await_ready() const noexcept { return false; }
            [[nodiscard]] bool await_suspend(coroutine_handle<> hdl) noexcept {
                m_handle = hdl;
                void* old_state = nullptr;
                if (!m_parent->m_state.compare_exchange_strong(old_state, this, std::memory_order::release,
                                                               std::memory_order::relaxed)) {
                    // No duplicate awaiters allowed, so the only valid values are m_parent and nullptr
                    assert(m_parent == old_state);
                    m_parent->m_state.exchange(nullptr, std::memory_order::acquire);
                    return false;
                }
                return true;
            }
            constexpr void await_resume() const noexcept {}
 
            single_consumer_auto_reset_event* m_parent;
            dispatcher* m_dispatcher;
            coroutine_handle<> m_handle{};
        };
    };
 
    class multi_consumer_event {
    public:
        explicit constexpr multi_consumer_event(bool set_initially = false) noexcept
            : m_state(set_initially ? this : nullptr) {}
#ifndef NDEBUG
        ~multi_consumer_event() noexcept { assert(!is_awaited()); }
#endif
 
        [[nodiscard]] bool is_set() const noexcept { return m_state.load(std::memory_order::acquire) == this; }
 
        [[nodiscard]] bool is_awaited() const noexcept {
            auto ptr = m_state.load(std::memory_order::acquire);
            return ptr != nullptr && ptr != this;
        }
 
        bool set(dispatcher* resume_dispatcher = nullptr) noexcept {
            auto state = m_state.exchange(this, std::memory_order::acq_rel);
            if (state == this) return false;
            auto await = static_cast<awaiter*>(state);
            while (await != nullptr) {
                auto next = await->m_next;
                assert(await->m_parent == this);
                assert(await->m_handle);
                if (await->m_dispatcher != nullptr) {
                    await->m_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else if (resume_dispatcher != nullptr) {
                    resume_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else {
                    await->m_handle.resume();
                }
                await = next;
            }
            return true;
        }
 
        void reset() noexcept {
            void* old_state = this;
            m_state.compare_exchange_strong(old_state, nullptr, std::memory_order::relaxed);
        }
 
        [[nodiscard]] auto operator co_await() noexcept { return awaiter{this, dispatcher::current()}; }
 
        [[nodiscard]] constexpr auto wait(dispatcher* resume_dispatcher = nullptr) noexcept {
            return awaiter{this, resume_dispatcher};
        }
 
    private:
        /* nullptr => unset
         * this => set
         * x => head of awaiter* list
         */
        std::atomic<void*> m_state;
 
        struct [[nodiscard]] awaiter {
            explicit constexpr awaiter(multi_consumer_event* parent, dispatcher* dispatcher) noexcept
                : m_parent(parent), m_dispatcher(dispatcher) {}
            [[nodiscard]] bool await_ready() const noexcept { return m_parent->is_set(); }
            [[nodiscard]] bool await_suspend(coroutine_handle<> hdl) noexcept {
                m_handle = hdl;
                void* old_state = m_parent->m_state.load(std::memory_order::acquire);
                do {
                    // event became set
                    if (old_state == m_parent) return false;
                    m_next = static_cast<awaiter*>(old_state);
                } while (!m_parent->m_state.compare_exchange_weak( //
                    old_state, this, std::memory_order::release, std::memory_order::acquire));
                return true;
            }
            constexpr void await_resume() const noexcept {}
 
            multi_consumer_event* m_parent;
            dispatcher* m_dispatcher;
            awaiter* m_next{nullptr};
            coroutine_handle<> m_handle{};
        };
    };
 
    class multi_consumer_auto_reset_event {
        explicit constexpr multi_consumer_auto_reset_event(bool set_initially = false) noexcept
            : m_state(set_initially ? this : nullptr) {}
#ifndef NDEBUG
        ~multi_consumer_auto_reset_event() noexcept { assert(!is_awaited()); }
#endif
 
        [[nodiscard]] bool is_set() const noexcept { return m_state.load(std::memory_order::acquire) == this; }
 
        [[nodiscard]] bool is_awaited() const noexcept {
            auto ptr = m_state.load(std::memory_order::acquire);
            return ptr != nullptr && ptr != this;
        }
 
        bool set(dispatcher* resume_dispatcher = nullptr) noexcept {
            auto state = m_state.exchange(this, std::memory_order::acq_rel);
            if (state == this || state == nullptr) return false;
            auto await = static_cast<awaiter*>(state);
 
            // Only modify the state if it has not been changed in between
            state = this;
            m_state.compare_exchange_strong(state, nullptr, std::memory_order::acq_rel);
 
            while (await != nullptr) {
                auto next = await->m_next;
                assert(await->m_parent == this);
                assert(await->m_handle);
                if (await->m_dispatcher != nullptr) {
                    await->m_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else if (resume_dispatcher != nullptr) {
                    resume_dispatcher->push([hdl = await->m_handle]() mutable { hdl.resume(); });
                } else {
                    await->m_handle.resume();
                }
                await = next;
            }
            return true;
        }
 
        void reset() noexcept {
            void* old_state = this;
            m_state.compare_exchange_strong(old_state, nullptr, std::memory_order::relaxed);
        }
 
        [[nodiscard]] auto operator co_await() noexcept { return awaiter{this, dispatcher::current()}; }
 
        [[nodiscard]] constexpr auto wait(dispatcher* resume_dispatcher = nullptr) noexcept {
            return awaiter{this, resume_dispatcher};
        }
 
    private:
        /* nullptr => unset
         * this => set
         * x => head of awaiter* list
         */
        std::atomic<void*> m_state;
 
        struct [[nodiscard]] awaiter {
            explicit constexpr awaiter(multi_consumer_auto_reset_event* parent, dispatcher* dispatcher) noexcept
                : m_parent(parent), m_dispatcher(dispatcher) {}
            [[nodiscard]] bool await_ready() const noexcept { return m_parent->is_set(); }
            [[nodiscard]] bool await_suspend(coroutine_handle<> hdl) noexcept {
                m_handle = hdl;
                void* old_state = m_parent->m_state.load(std::memory_order::acquire);
                do {
                    // event became set
                    if (old_state == m_parent) return false;
                    m_next = static_cast<awaiter*>(old_state);
                } while (!m_parent->m_state.compare_exchange_weak( //
                    old_state, this, std::memory_order::release, std::memory_order::acquire));
                return true;
            }
            constexpr void await_resume() const noexcept {}
 
            multi_consumer_auto_reset_event* m_parent;
            dispatcher* m_dispatcher{};
            awaiter* m_next{nullptr};
            coroutine_handle<> m_handle{};
        };
    };
} // namespace asyncpp