xbot/irc_coroutine.hpp

#pragma once

#include "connection.hpp"

#include <chrono>
#include <coroutine>
#include <initializer_list>
#include <utility>
#include <vector>

struct irc_promise;

/// A coroutine that can co_await on various IRC events
struct irc_coroutine : std::coroutine_handle<irc_promise>
{
    using promise_type = irc_promise;

    /// Start the coroutine and associate it with a specific connection.
    auto start(Connection &connection) -> void;

    /// Returns true when this coroutine is still waiting on events
    auto is_running() -> bool;

    /// Returns the exception that terminated this coroutine, if there is one.
    auto exception() -> std::exception_ptr;
};

struct irc_promise
{
    // Pointer to the connection while running. Cleared on termination.
    std::shared_ptr<Connection> connection_;

    // Pointer to exception that terminated this coroutine if there is one.
    std::exception_ptr exception_;

    auto get_return_object() -> irc_coroutine
    {
        return {irc_coroutine::from_promise(*this)};
    }

    // Suspend waiting for start() to initialize connection_
    auto initial_suspend() noexcept -> std::suspend_always { return {}; }

    // Suspend so that is_running() and exception() work
    auto final_suspend() noexcept -> std::suspend_always { return {}; }

    // Normal termination
    auto return_void() -> void
    {
        connection_.reset();
    }

    // Abnormal termination - remember the exception
    auto unhandled_exception() -> void
    {
        connection_.reset();
        exception_ = std::current_exception();
    }
};

template <typename... Ts>
class Wait;

/// Argument to a Wait that expects one or more IRC messages
class wait_ircmsg
{
    // Vector of commands this wait is expecting. Leave empty to accept all messages.
    std::vector<IrcCommand> want_cmds_;

    // Slot for the ircmsg event
    boost::signals2::scoped_connection ircmsg_slot_;

public:
    using result_type = std::pair<IrcCommand, const IrcMsg &>;

    wait_ircmsg(std::initializer_list<IrcCommand> want_cmds)
        : want_cmds_{want_cmds}
    {
    }

    template <size_t I, typename... Ts>
    auto start(Wait<Ts...> &command) -> void;
    auto stop() -> void { ircmsg_slot_.disconnect(); }
};

class wait_snote
{
    // Vector of tags this wait is expecting. Leave empty to accept all messages.
    std::vector<SnoteTag> want_tags_;

    // Slot for the snote event
    boost::signals2::scoped_connection snote_slot_;

public:
    using result_type = SnoteMatch;

    wait_snote(std::initializer_list<SnoteTag> want_tags)
        : want_tags_{want_tags}
    {
    }

    template <size_t I, typename... Ts>
    auto start(Wait<Ts...> &command) -> void;
    auto stop() -> void { snote_slot_.disconnect(); }
};

class wait_timeout
{
    std::optional<boost::asio::steady_timer> timer_;
    std::chrono::milliseconds timeout_;

public:
    struct result_type
    {
    };
    wait_timeout(std::chrono::milliseconds timeout)
        : timeout_{timeout}
    {
    }

    template <size_t I, typename... Ts>
    auto start(Wait<Ts...> &command) -> void;
    auto stop() -> void { timer_->cancel(); }
};

template <typename... Ts>
class Wait
{

    // State associated with each wait mode
    std::tuple<Ts...> modes_;

    // Result from any one of the wait modes
    std::optional<std::variant<typename Ts::result_type...>> result_;

    // Handle of the continuation to be resumed when one of the wait
    // modes is ready.
    std::coroutine_handle<irc_promise> handle_;

    // Slot for tearing down the irc_coroutine in case the connection
    // fails before any wait modes complete.
    boost::signals2::scoped_connection disconnect_slot_;

    template <size_t I>
    auto start_mode() -> void
    {
        std::get<I>(modes_).template start<I, Ts...>(*this);
    }

    template <std::size_t... Indices>
    auto start_modes(std::index_sequence<Indices...>) -> void
    {
        (start_mode<Indices>(), ...);
    }

    template <std::size_t... Indices>
    auto stop_modes(std::index_sequence<Indices...>) -> void
    {
        (std::get<Indices>(modes_).stop(), ...);
    }

public:
    Wait(Ts &&...modes)
        : modes_{std::forward<Ts>(modes)...}
    {
    }

    // Get the connection that this coroutine was started with.
    auto get_connection() const -> Connection &
    {
        return *handle_.promise().connection_;
    }

    // Store a successful result and resume the coroutine
    template <size_t I, typename... Args>
    auto complete(Args &&...args) -> void
    {
        result_.emplace(std::in_place_index<I>, std::forward<Args>(args)...);
        handle_.resume();
    }

    // The coroutine always needs to wait for a message. It will never
    // be ready immediately.
    auto await_ready() noexcept -> bool { return false; }

    /// Install event handles in the connection that will resume this coroutine.
    auto await_suspend(std::coroutine_handle<irc_promise> handle) -> void;

    auto await_resume() -> std::variant<typename Ts::result_type...>;
};

template <size_t I, typename... Ts>
auto wait_ircmsg::start(Wait<Ts...> &command) -> void
{
    ircmsg_slot_ = command.get_connection().sig_ircmsg.connect([this, &command](auto cmd, auto &msg) {
        const auto wanted = want_cmds_.empty() || std::find(want_cmds_.begin(), want_cmds_.end(), cmd) != want_cmds_.end();
        if (wanted)
        {
            command.template complete<I>(cmd, msg);
        }
    });
}

template <size_t I, typename... Ts>
auto wait_snote::start(Wait<Ts...> &command) -> void
{
    snote_slot_ = command.get_connection().sig_snote.connect([this, &command](auto &match) {
        const auto tag = match.get_tag();
        const auto wanted = want_tags_.empty() || std::find(want_tags_.begin(), want_tags_.end(), tag) != want_tags_.end();
        if (wanted)
        {
            command.template complete<I>(match);
        }
    });
}

template <size_t I, typename... Ts>
auto wait_timeout::start(Wait<Ts...> &command) -> void
{
    timer_.emplace(command.get_connection().get_executor());
    timer_->expires_after(timeout_);
    timer_->async_wait([this, &command](const auto &error) {
        if (not error)
        {
            timer_.reset();
            command.template complete<I>();
        }
    });
}

template <typename... Ts>
auto Wait<Ts...>::await_suspend(std::coroutine_handle<irc_promise> handle) -> void
{
    handle_ = handle;

    const auto tuple_size = std::tuple_size_v<decltype(modes_)>;
    start_modes(std::make_index_sequence<tuple_size>{});

    disconnect_slot_ = get_connection().sig_disconnect.connect([this]() {
        handle_.resume();
    });
}

template <typename... Ts>
auto Wait<Ts...>::await_resume() -> std::variant<typename Ts::result_type...>
{
    const auto tuple_size = std::tuple_size_v<decltype(modes_)>;
    stop_modes(std::make_index_sequence<tuple_size>{});

    disconnect_slot_.disconnect();

    if (result_)
    {
        return std::move(*result_);
    }
    else
    {
        throw std::runtime_error{"connection terminated"};
    }
}

/// Start the coroutine and associate it with a specific connection.
inline auto irc_coroutine::start(Connection &connection) -> void
{
    promise().connection_ = connection.shared_from_this();
    resume();
}

/// Returns true when this coroutine is still waiting on events
inline auto irc_coroutine::is_running() -> bool
{
    return promise().connection_ != nullptr;
}

inline auto irc_coroutine::exception() -> std::exception_ptr
{
    return promise().exception_;
}