notifier.hpp

// Copyright Take Vos 2022.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt)
 
#pragma once
 
#include "utility/module.hpp"
#include "generator.hpp"
#include "loop.hpp"
#include "callback_flags.hpp"
#include <vector>
#include <tuple>
#include <functional>
#include <coroutine>
#include <mutex>
 
namespace hi::inline v1 {

template<typename T = void()>
class notifier {
};
 
// The partial template specialization allows the use of a `std::function`-like template
// argument, that looks like a function prototype.
template<typename Result, typename... Args>
class notifier<Result(Args...)> {
public:
    static_assert(std::is_same_v<Result, void>, "Result of a notifier must be void.");
 
    using result_type = Result;
    using callback_proto = Result(Args...);
    using function_type = std::function<callback_proto>;
 
    using callback_token = std::shared_ptr<function_type>;
    using weak_callback_token = std::weak_ptr<function_type>;

    class awaiter_type {
    public:
        constexpr awaiter_type() noexcept = default;
        constexpr awaiter_type(awaiter_type const&) noexcept = default;
        constexpr awaiter_type(awaiter_type&&) noexcept = default;
        constexpr awaiter_type& operator=(awaiter_type const&) noexcept = default;
        constexpr awaiter_type& operator=(awaiter_type&&) noexcept = default;
 
        constexpr awaiter_type(notifier& notifier) noexcept : _notifier(&notifier) {}
 
        [[nodiscard]] constexpr bool await_ready() noexcept
        {
            return false;
        }
 
        void await_suspend(std::coroutine_handle<> handle) noexcept
        {
            hi_assert_not_null(_notifier);
 
            // We can use the this pointer in the callback, as `await_suspend()` is called by
            // the co-routine on the same object as `await_resume()`.
            _cbt = _notifier->subscribe(
                [this, handle](Args const&...args) {
                    // Copy the arguments received from the notifier into the awaitable object
                    // So that it can be read using `await_resume()`.
                    _args = {args...};
 
                    // Resume the co-routine.
                    handle.resume();
                },
                callback_flags::main | callback_flags::once);
        }
 
        constexpr void await_resume() const noexcept requires(sizeof...(Args) == 0) {}
 
        constexpr auto await_resume() const noexcept requires(sizeof...(Args) == 1)
        {
            return std::get<0>(_args);
        }
 
        constexpr auto await_resume() const noexcept requires(sizeof...(Args) > 1)
        {
            return _args;
        }
 
    private:
        notifier *_notifier = nullptr;
        callback_token _cbt;
        std::tuple<Args...> _args;
    };

    constexpr notifier() noexcept = default;
    constexpr notifier(notifier&&) noexcept = default;
    constexpr notifier(notifier const&) noexcept = default;
    constexpr notifier& operator=(notifier&&) noexcept = default;
    constexpr notifier& operator=(notifier const&) noexcept = default;

    awaiter_type operator co_await() const noexcept
    {
        return awaiter_type{const_cast<notifier&>(*this)};
    }

    [[nodiscard]] callback_token
    subscribe(forward_of<callback_proto> auto&& callback, callback_flags flags = callback_flags::synchronous) noexcept
    {
        auto token = std::make_shared<function_type>(hi_forward(callback));
 
        hilet lock = std::scoped_lock(_mutex);
        _callbacks.emplace_back(token, flags);
        return token;
    }

    void operator()(Args const&...args) const noexcept
    {
        hilet lock = std::scoped_lock(_mutex);
 
        for (auto& callback : _callbacks) {
            if (is_synchronous(callback.flags)) {
                if (auto func = callback.lock()) {
                    (*func)(args...);
                }
 
            } else if (is_local(callback.flags)) {
                loop::local().post_function([=] {
                    if (auto func = callback.lock()) {
                        (*func)(args...);
                    }
                });
 
            } else if (is_main(callback.flags)) {
                loop::main().post_function([=] {
                    if (auto func = callback.lock()) {
                        (*func)(args...);
                    }
                });
 
            } else if (is_timer(callback.flags)) {
                loop::timer().post_function([=] {
                    if (auto func = callback.lock()) {
                        (*func)(args...);
                    }
                });
 
            } else {
                hi_no_default();
            }
 
            // If the callback should only be triggered once, like inside an awaitable.
            // Then reset the weak_ptr in _callbacks so that it will be cleaned up.
            // In the lambda above the weak_ptr is copied first so that it callback will get executed
            // as long as the shared_ptr's use count does not go to zero.
            if (is_once(callback.flags)) {
                callback.reset();
            }
        }
        clean_up();
    }
 
private:
    struct callback_type {
        weak_callback_token token;
        callback_flags flags;
 
        [[nodiscard]] bool expired() const noexcept
        {
            return token.expired();
        }
 
        void reset() noexcept
        {
            token.reset();
        }
 
        [[nodiscard]] callback_token lock() const noexcept
        {
            return token.lock();
        }
    };
 
    mutable unfair_mutex _mutex;

    mutable std::vector<callback_type> _callbacks;
 
    void clean_up() const noexcept
    {
        hi_axiom(_mutex.is_locked());
 
        // Cleanup all callbacks that have expired, or when they may only be triggered once.
        std::erase_if(_callbacks, [](hilet& item) {
            return item.expired();
        });
    }
 
#ifndef NDEBUG
    mutable bool _notifying = false;
#endif
};
 
} // namespace hi::inline v1