libdpf/include/dpf/net/edge_mesh.hpp

/// @file dpf/net/edge_mesh.hpp
/// @brief N-edge RoundSink mesh for star / dealer / 4PC topologies.
#ifndef LIBDPF_INCLUDE_DPF_NET_EDGE_MESH_HPP__
#define LIBDPF_INCLUDE_DPF_NET_EDGE_MESH_HPP__

#include <cstddef>
#include <cstdint>
#include <memory>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "dpf/net/memory_sink.hpp"
#include "dpf/net/round_sink.hpp"

namespace dpf
{
namespace net
{

/// @brief Index of a duplex link in an `edge_mesh`.
using edge_id = std::uint16_t;

/// @brief Named edges for the common 2PC / RSS / dealer trio (also mesh ids 0..2).
/// @details For the dealer (party 2), `edge_dealer` is its link to party 0 and
///          `edge_dealer_p1` its link to party 1.
inline constexpr edge_id edge_peer = 0;
inline constexpr edge_id edge_rss_next = 1;
inline constexpr edge_id edge_dealer = 2;
inline constexpr edge_id edge_dealer_p1 = 3;

inline std::string edge_name(edge_id e)
{
    switch (e)
    {
    case edge_peer:
        return "peer";
    case edge_rss_next:
        return "rss_next";
    case edge_dealer:
        return "dealer";
    case edge_dealer_p1:
        return "dealer->p1";
    default:
        return "edge " + std::to_string(e);
    }
}

/// @brief Collection of duplex RoundSinks keyed by `edge_id`.
struct edge_mesh
{
    std::vector<RoundSink *> sinks;

    RoundSink & at(edge_id id) const
    {
        if (static_cast<std::size_t>(id) >= sinks.size() || sinks[id] == nullptr)
            throw std::logic_error("edge_mesh: edge not bound");
        return *sinks[id];
    }

    bool has(edge_id id) const noexcept
    {
        return static_cast<std::size_t>(id) < sinks.size()
            && sinks[id] != nullptr;
    }

    std::size_t size() const noexcept { return sinks.size(); }

    void flush_all()
    {
        for (auto * s : sinks)
        {
            if (s == nullptr)
                continue;
            s->flush();
            s->poll();
        }
    }

    /// @brief Sum of `progress()` over distinct bound sinks.
    std::uint64_t progress_total() const noexcept
    {
        std::uint64_t total = 0;
        for (std::size_t i = 0; i < sinks.size(); ++i)
        {
            const RoundSink * s = sinks[i];
            if (s == nullptr)
                continue;
            bool seen = false;
            for (std::size_t j = 0; j < i && !seen; ++j)
                seen = sinks[j] == s;
            if (!seen)
                total += s->progress();
        }
        return total;
    }

    /// @brief Block until at least one bound sink makes I/O progress.
    /// @details Returns true if any sink's `wait_io()` ran a real event (an
    ///          async sink slept in `epoll`); false if none did (memory sinks),
    ///          so the caller can fall back to its spin guard.
    bool wait_io_all()
    {
        bool progressed = false;
        for (auto * s : sinks)
        {
            if (s == nullptr)
                continue;
            if (s->wait_io())
                progressed = true;
        }
        return progressed;
    }
};

/// @brief In-process star: one client edge per server, matching server ends.
struct memory_star
{
    std::size_t servers = 0;
    std::vector<std::shared_ptr<memory_sink_hub>> hubs;
    std::vector<memory_sink> client;  ///< client side of edge i
    std::vector<memory_sink> server;  ///< server i side of edge i

    /// @brief Client mesh: edges `[0, servers)`.
    edge_mesh client_mesh()
    {
        edge_mesh m;
        m.sinks.resize(client.size());
        for (std::size_t i = 0; i < client.size(); ++i)
            m.sinks[i] = &client[i];
        return m;
    }

    /// @brief Server `i` mesh with a single live edge at `edge_id{i}` (sparse).
    ///        Prefer `server_edge(i)` when the schedule uses `edge_id{0}` locally.
    edge_mesh server_mesh_at(std::size_t i)
    {
        if (i >= server.size())
            throw std::out_of_range("memory_star server");
        edge_mesh m;
        m.sinks.assign(server.size(), nullptr);
        m.sinks[i] = &server[i];
        return m;
    }

    /// @brief Server `i` as a one-edge mesh (`edge_id` 0 → that duplex).
    edge_mesh server_edge(std::size_t i)
    {
        if (i >= server.size())
            throw std::out_of_range("memory_star server");
        return edge_mesh{{&server[i]}};
    }
};

/// @brief Build an in-process client↔N-server star.
/// @param slot_bytes Round widths shared by every edge (same schedule shape).
inline memory_star make_memory_star(std::size_t n_servers, std::size_t count,
    std::vector<std::size_t> slot_bytes)
{
    if (n_servers < 2)
        throw std::invalid_argument("make_memory_star needs >= 2 servers");
    memory_star star;
    star.servers = n_servers;
    star.hubs.reserve(n_servers);
    star.client.reserve(n_servers);
    star.server.reserve(n_servers);
    for (std::size_t i = 0; i < n_servers; ++i)
    {
        auto hub = std::make_shared<memory_sink_hub>(count, slot_bytes);
        star.hubs.push_back(hub);
        star.client.emplace_back(hub, true);
        star.server.emplace_back(hub, false);
    }
    return star;
}

/// @brief Fully connected memory clique of `n` roles (every unordered pair).
/// @details Edge id for ordered pair (a,b) with a<b is the combinatorial
///          index; both directions share one hub (a is side_a).
struct memory_clique
{
    std::size_t roles = 0;
    std::vector<std::shared_ptr<memory_sink_hub>> hubs;
    /// hubs[edge], ends[edge].first = lower role, .second = higher role
    std::vector<std::pair<memory_sink, memory_sink>> ends;

    static edge_id pair_edge(std::size_t a, std::size_t b, std::size_t n)
    {
        if (a == b || a >= n || b >= n)
            throw std::invalid_argument("memory_clique pair");
        if (a > b)
            std::swap(a, b);
        // Index among pairs (i,j) with i<j.
        edge_id e = 0;
        for (std::size_t i = 0; i < a; ++i)
            e = static_cast<edge_id>(e + (n - 1 - i));
        e = static_cast<edge_id>(e + (b - a - 1));
        return e;
    }

    RoundSink & end(std::size_t role, std::size_t peer)
    {
        const auto e = pair_edge(role, peer, roles);
        if (role < peer)
            return ends[e].first;
        return ends[e].second;
    }

    edge_mesh mesh_for(std::size_t role)
    {
        edge_mesh m;
        m.sinks.assign(ends.size(), nullptr);
        for (std::size_t p = 0; p < roles; ++p)
        {
            if (p == role)
                continue;
            m.sinks[pair_edge(role, p, roles)] = &end(role, p);
        }
        return m;
    }
};

inline memory_clique make_memory_clique(std::size_t n_roles, std::size_t count,
    std::vector<std::size_t> slot_bytes)
{
    if (n_roles < 2)
        throw std::invalid_argument("make_memory_clique needs >= 2 roles");
    memory_clique c;
    c.roles = n_roles;
    const std::size_t n_edges = n_roles * (n_roles - 1) / 2;
    c.hubs.reserve(n_edges);
    c.ends.reserve(n_edges);
    for (std::size_t e = 0; e < n_edges; ++e)
    {
        auto hub = std::make_shared<memory_sink_hub>(count, slot_bytes);
        c.hubs.push_back(hub);
        c.ends.emplace_back(memory_sink(hub, true), memory_sink(hub, false));
    }
    return c;
}

}  // namespace net
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_NET_EDGE_MESH_HPP__