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 {
    public:
        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 {
            // We assume the event is
            awaiter* await = nullptr;
            void* old_state = m_state.load(std::memory_order::acquire);
            void* new_state;
            do {
                // If the state is already set we can just return
                if (old_state == this) return false;
                // Otherwise store the state as an awaiter
                await = static_cast<awaiter*>(old_state);
                // if the event had awaiters we update the state to "unset", if it was previously unset we transition to set
                new_state = (old_state == nullptr) ? this : nullptr;
            } while (!m_state.compare_exchange_weak(old_state, new_state, std::memory_order::release,
                                                    std::memory_order::acquire));
 
            // Execute the awaiters (if any)
            if (await == nullptr) return false;
            do {
                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;
            } while (await != nullptr);
            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 {
                // We assume the event is set
                void* old_state = m_parent;
                // And try to replace it with unset
                // NOTE: compare_exchange_weak can fail spuriously, but we don't care
                //       because in that case we just enter the main loop in await_suspend
                // Returns true if the value was updated (i.e. old_state equalled m_parent and cas succeeded)
                return m_parent->m_state.compare_exchange_weak( //
                    old_state, nullptr, std::memory_order::release, std::memory_order::acquire);
            }
            [[nodiscard]] bool await_suspend(coroutine_handle<> hdl) noexcept {
                m_handle = hdl;
                void* old_state = m_parent->m_state.load(std::memory_order::acquire);
                do {
                    // If the event is set, reset it and resume the coroutine
                    if (old_state == m_parent) {
                        if (m_parent->m_state.compare_exchange_weak( //
                                old_state, nullptr, std::memory_order::release, std::memory_order::acquire))
                            return false;
                        // the state changed between load and compare (either because another coroutine was faster or the event was reset).
                        // Retry from the start
                        continue;
                    }
                    // The event is unset, add the existing awaiter (if any) to the list and suspend the coroutine
                    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