test_unit.cpp

// Unit tests for the refactored control plane: pure-CPU algorithms and
// wire-codec round-trips. No GPU, no network, no agents required.
//
// Each `case_*` function is one focused check; `expect()` aborts the
// process with a clear message on failure so ctest sees a non-zero exit.

#include "control_plane/common/codec.h"
#include "control_plane/common/messages.h"
#include "control_plane/controller/cc/channel_creation.h"
#include "control_plane/controller/td/global_assembler.h"
#include "control_plane/controller/to/transfer_optimizer.h"
#include "control_plane/gig/forwarding.h"
#include "control_plane/gig/graph.h"
#include "control_plane/gig/max_flow.h"
#include "control_plane/gig/path.h"

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>

using namespace mgi;

namespace {
int g_pass = 0;
int g_fail = 0;

#define EXPECT(cond)                                                       \
  do                                                                       \
    {                                                                      \
      if (!(cond))                                                         \
        {                                                                  \
          std::fprintf (stderr, "FAIL %s:%d  %s\n", __FILE__, __LINE__,     \
              #cond);                                                       \
          ++g_fail;                                                        \
          return;                                                          \
        }                                                                  \
    }                                                                      \
  while (0)

#define CASE(name) static void name ()

// ─── Helpers ─────────────────────────────────────────────────────────────
gig::Gig
two_path_graph ()
{
  // Two parallel paths a -> sw1 -> b (cap 100) and a -> sw2 -> b (cap 50).
  gig::Gig g;
  g.add_vertex ({ 1, gig::VertexKind::GPU,      true,  "a" });
  g.add_vertex ({ 2, gig::VertexKind::GPU,      true,  "b" });
  g.add_vertex ({ 3, gig::VertexKind::NVSwitch, false, "sw1" });
  g.add_vertex ({ 4, gig::VertexKind::NVSwitch, false, "sw2" });
  g.add_edge ({ 1, 3, 100.0, gig::LinkType::NVLink });
  g.add_edge ({ 3, 1, 100.0, gig::LinkType::NVLink });
  g.add_edge ({ 3, 2, 100.0, gig::LinkType::NVLink });
  g.add_edge ({ 2, 3, 100.0, gig::LinkType::NVLink });
  g.add_edge ({ 1, 4, 50.0,  gig::LinkType::NVLink });
  g.add_edge ({ 4, 1, 50.0,  gig::LinkType::NVLink });
  g.add_edge ({ 4, 2, 50.0,  gig::LinkType::NVLink });
  g.add_edge ({ 2, 4, 50.0,  gig::LinkType::NVLink });
  return g;
}

agent::LocalSubgraph
fake_subgraph (uint32_t agent_id, double nic_bw)
{
  agent::LocalSubgraph s;
  s.agent_id = agent_id;
  auto vid = [agent_id] (uint32_t l) -> gig::VertexId {
    return ((uint64_t)agent_id << 32) | l;
  };
  s.gig.add_vertex ({ vid (0),     gig::VertexKind::CPU, true,  "cpu" });
  s.gig.add_vertex ({ vid (0x100), gig::VertexKind::GPU, true,  "gpu" });
  s.gig.add_vertex ({ vid (0x200), gig::VertexKind::NIC, false, "nic" });
  s.gig.add_edge ({ vid (0x100), vid (0),     64.0, gig::LinkType::PCIe });
  s.gig.add_edge ({ vid (0),     vid (0x100), 64.0, gig::LinkType::PCIe });
  s.gig.add_edge ({ vid (0x200), vid (0),     nic_bw, gig::LinkType::PCIe });
  s.gig.add_edge ({ vid (0),     vid (0x200), nic_bw, gig::LinkType::PCIe });
  return s;
}

// ─── Cases ───────────────────────────────────────────────────────────────

CASE (case_max_flow_two_paths)
{
  auto         g  = two_path_graph ();
  gig::PathSet ps = gig::edmonds_karp (g, 1, 2);
  EXPECT (ps.size () == 2);
  // Sum of bottlenecks should equal 100 + 50 = 150.
  double sum = 0;
  for (auto &p : ps) sum += p.capacity;
  EXPECT (std::abs (sum - 150.0) < 1e-6);
}

CASE (case_max_flow_self_and_missing)
{
  auto g = two_path_graph ();
  EXPECT (gig::edmonds_karp (g, 1, 1).empty ());
  EXPECT (gig::edmonds_karp (g, 1, 99).empty ());
}

CASE (case_concat_basic_and_mismatch)
{
  gig::Path a{ { 1, 3, 5 }, 100.0 };
  gig::Path b{ { 5, 7, 9 }, 50.0 };
  auto      c = gig::concat (a, b);
  EXPECT (c.vertices.size () == 5);
  EXPECT (c.vertices[0] == 1 && c.vertices[4] == 9);
  EXPECT (std::abs (c.capacity - 50.0) < 1e-6);

  gig::Path d{ { 8, 9 }, 30.0 };
  auto      bad = gig::concat (a, d);
  EXPECT (bad.empty ());
}

CASE (case_forwarding_gpu_only_merge)
{
  // a (GPU) -> sw1 -> b (GPU) and a -> sw2 -> b. Forwarding table for
  // a should have one bin (next=b) with 2 paths after §4.4 merge.
  auto         g  = two_path_graph ();
  gig::PathSet ps = gig::edmonds_karp (g, 1, 2);
  auto         t  = gig::build_from_paths (g, ps);
  EXPECT (t.count (1) == 1); // GPU a has a table
  EXPECT (t.count (3) == 0); // NVSwitch sw1 does not (GPU-only)
  EXPECT (t.count (4) == 0);
  const auto &bins = t.at (1).at (2);
  EXPECT (bins.size () == 1);
  EXPECT (bins.front ().next == 2);
  EXPECT (bins.front ().paths.size () == 2);
}

CASE (case_assemble_global_v_net_fanout)
{
  std::vector<agent::LocalSubgraph> subs;
  subs.push_back (fake_subgraph (0x0a000001, 32.0));
  subs.push_back (fake_subgraph (0x0a000002, 64.0));
  auto g = td::assemble_global (subs);

  // 3 vertices/agent * 2 + 1 v_net = 7
  EXPECT (g.vertices ().size () == 7);
  EXPECT (g.has_vertex (td::kNetCloudVertexId));

  size_t v_net_out = g.out_edges (td::kNetCloudVertexId).size ();
  EXPECT (v_net_out == 2); // one edge to each NIC
  // Capacity inferred from each agent's NIC<->CPU PCIe edge.
  std::vector<double> caps;
  for (const auto &e : g.out_edges (td::kNetCloudVertexId))
    caps.push_back (e.capacity);
  std::sort (caps.begin (), caps.end ());
  EXPECT (std::abs (caps[0] - 32.0) < 1e-6);
  EXPECT (std::abs (caps[1] - 64.0) < 1e-6);
}

CASE (case_to_cache_reuse_cross_server)
{
  std::vector<agent::LocalSubgraph> subs;
  subs.push_back (fake_subgraph (0x0a000001, 32.0));
  subs.push_back (fake_subgraph (0x0a000002, 64.0));
  auto                  g = td::assemble_global (subs);
  to::TransferOptimizer opt (g);

  auto gpu_a = (((uint64_t)0x0a000001) << 32) | 0x100;
  auto gpu_b = (((uint64_t)0x0a000002) << 32) | 0x100;

  (void)opt.paths (gpu_a, gpu_b);
  size_t after_first = opt.cache_size ();
  // Expected entries: (gpu_a, v_net), (v_net, gpu_b), (gpu_a, gpu_b).
  EXPECT (after_first == 3);

  (void)opt.paths (gpu_a, gpu_b);
  EXPECT (opt.cache_size () == 3); // pure cache hit, no growth

  // A second source crossing the same destination reuses (v_net, gpu_b).
  // We add a third agent with a single GPU and re-test.
  std::vector<agent::LocalSubgraph> subs2 = subs;
  subs2.push_back (fake_subgraph (0x0a000003, 16.0));
  auto                  g2 = td::assemble_global (subs2);
  to::TransferOptimizer opt2 (g2);
  auto gpu_c = (((uint64_t)0x0a000003) << 32) | 0x100;

  (void)opt2.paths (gpu_c, gpu_b);
  size_t e1 = opt2.cache_size ();   // 3 fresh entries
  (void)opt2.paths (gpu_a, gpu_b);
  size_t e2 = opt2.cache_size ();
  // Adds (gpu_a, v_net) and (gpu_a, gpu_b); reuses (v_net, gpu_b).
  EXPECT (e2 - e1 == 2);
}

CASE (case_codec_primitives_roundtrip)
{
  common::Writer w;
  w.put_u8 (0x12);
  w.put_u16 (0x1234);
  w.put_u32 (0x12345678u);
  w.put_u64 (0x123456789abcdef0ull);
  w.put_double (3.14159265358979323846);
  w.put_string ("hello, mgi");

  common::Reader r (w.buf);
  EXPECT (r.get_u8 () == 0x12);
  EXPECT (r.get_u16 () == 0x1234);
  EXPECT (r.get_u32 () == 0x12345678u);
  EXPECT (r.get_u64 () == 0x123456789abcdef0ull);
  EXPECT (std::abs (r.get_double () - 3.14159265358979323846) < 1e-12);
  EXPECT (r.get_string () == "hello, mgi");
  EXPECT (r.remaining () == 0);
}

CASE (case_codec_message_roundtrip)
{
  // Serialize a synthetic ReportSubgraph and DeliverFwdTable and check
  // they reconstruct byte-for-byte equal counts.
  common::ReportSubgraphMsg in;
  in.sub             = fake_subgraph (0x01020304, 25.0);
  in.sub.gdr_module  = true;
  in.sub.pcie_p2p_ok = false;
  in.sub.gdr_active  = false;
  in.push_port       = 0xCAFE;

  common::Writer w;
  common::write_report_subgraph (w, in);
  common::Reader            r (w.buf);
  common::ReportSubgraphMsg out;
  common::read_report_subgraph (r, out);
  EXPECT (out.sub.agent_id == in.sub.agent_id);
  EXPECT (out.sub.gdr_module == in.sub.gdr_module);
  EXPECT (out.sub.pcie_p2p_ok == in.sub.pcie_p2p_ok);
  EXPECT (out.sub.gdr_active == in.sub.gdr_active);
  EXPECT (out.push_port == in.push_port);
  EXPECT (out.sub.gig.vertices ().size () == in.sub.gig.vertices ().size ());
  EXPECT (out.sub.gig.edges ().size () == in.sub.gig.edges ().size ());
  EXPECT (r.remaining () == 0);
}

// ─── PCIe aggregator (sw_up / sw_dn) share-modelling cases ───────────────
//
// These are synthetic, exercising max-flow on the directed-split structure
// that local_discover would emit when a GPU and a NIC sit on the same PCIe
// switch. They do NOT call discover_local (no GPU/NIC needed in unit
// tier) -- they construct the graph in-memory.
//
// Layout (single agent, pretend GPU and NIC share switch X with 64 GB/s
// upstream link to CPU):
//
//   GPU ──→ sw_up ──→ CPU       (UP, cap 64 each)
//   CPU ──→ sw_dn ──→ GPU       (DOWN, cap 64 each)
//   NIC ──→ sw_up               (UP, cap 64)
//   sw_dn ──→ NIC               (DOWN, cap 64)
//
// GDR-on case adds a directly-bidirectional gate edge GPU↔NIC at cap 64
// (the peer-to-peer route through the switch fabric, bypassing CPU).
gig::Gig
pcie_aggregator_graph (bool gdr_on)
{
  gig::Gig g;
  // Vertices: cpu=1, gpu=2, nic=3, sw_up=4, sw_dn=5.
  g.add_vertex ({ 1, gig::VertexKind::CPU,            true,  "cpu" });
  g.add_vertex ({ 2, gig::VertexKind::GPU,            true,  "gpu" });
  g.add_vertex ({ 3, gig::VertexKind::NIC,            false, "nic" });
  g.add_vertex ({ 4, gig::VertexKind::PcieAggregator, false, "sw_up" });
  g.add_vertex ({ 5, gig::VertexKind::PcieAggregator, false, "sw_dn" });
  // Directed-only port edges (asymmetric -- this is the whole point).
  g.add_edge ({ 2, 4, 64.0, gig::LinkType::PCIe }); // gpu → sw_up
  g.add_edge ({ 3, 4, 64.0, gig::LinkType::PCIe }); // nic → sw_up
  g.add_edge ({ 4, 1, 64.0, gig::LinkType::PCIe }); // sw_up → cpu (shared)
  g.add_edge ({ 1, 5, 64.0, gig::LinkType::PCIe }); // cpu → sw_dn (shared)
  g.add_edge ({ 5, 2, 64.0, gig::LinkType::PCIe }); // sw_dn → gpu
  g.add_edge ({ 5, 3, 64.0, gig::LinkType::PCIe }); // sw_dn → nic
  if (gdr_on)
    {
      g.add_edge ({ 2, 3, 64.0, gig::LinkType::PCIe }); // gate gpu → nic
      g.add_edge ({ 3, 2, 64.0, gig::LinkType::PCIe }); // gate nic → gpu
    }
  return g;
}

CASE (case_pcie_share_gdr_off_routes_via_cpu)
{
  // Without the gate edge, GPU↔NIC must transit CPU. The split nodes
  // have no edge from sw_up to sw_dn (or back to peer devices), so
  // edmonds_karp can't shortcut across the aggregator.
  auto g  = pcie_aggregator_graph (/*gdr_on=*/false);
  auto ps = gig::edmonds_karp (g, 2, 3); // gpu → nic
  EXPECT (!ps.empty ());
  // Total flow = upstream cap = 64 (each path bottlenecked by sw_up→cpu
  // and cpu→sw_dn, both cap 64; only one such path exists).
  double sum = 0;
  for (const auto &p : ps) sum += p.capacity;
  EXPECT (std::abs (sum - 64.0) < 1e-6);
  // Every returned path must go through CPU (vertex id 1).
  for (const auto &p : ps)
    {
      bool through_cpu = false;
      for (gig::VertexId v : p.vertices)
        if (v == 1) through_cpu = true;
      EXPECT (through_cpu);
    }
}

CASE (case_pcie_share_gdr_on_prefers_gate)
{
  // With the gate, edmonds_karp finds both the 1-hop gate path and the
  // 4-hop CPU-bounce path. Total max-flow in our model is the sum of
  // the two -- 64 (gate) + 64 (cpu detour) = 128 -- which is the
  // accepted PCIe-lane over-estimation we already documented (the gate
  // edge models GDR's qualitative advantage, not a faithful upper bound
  // on a single physical PCIe x16 lane). What matters here:
  //   (a) at least one path is the 2-vertex gate route
  //   (b) the bottleneck distribution shows the gate is enabled
  auto g  = pcie_aggregator_graph (/*gdr_on=*/true);
  auto ps = gig::edmonds_karp (g, 2, 3);
  EXPECT (!ps.empty ());
  bool found_gate = false;
  for (const auto &p : ps)
    {
      // gate path: vertices == [gpu(2), nic(3)]
      if (p.vertices.size () == 2 && p.vertices[0] == 2 && p.vertices[1] == 3)
        {
          found_gate = true;
          EXPECT (std::abs (p.capacity - 64.0) < 1e-6);
        }
    }
  EXPECT (found_gate);
  double sum = 0;
  for (const auto &p : ps) sum += p.capacity;
  EXPECT (sum > 64.0 + 1e-6); // strictly more than GDR-off, gate adds bw
}

CASE (case_cc_all_to_all_local)
{
  // 3 GPUs in one server, fully-connected NVLink. CC should produce a
  // 3-vertex FwdTables, each table with 2 destinations.
  agent::LocalSubgraph s;
  s.agent_id = 0x01010101;
  auto vid = [&] (uint32_t l) -> gig::VertexId {
    return ((uint64_t)s.agent_id << 32) | l;
  };
  for (uint32_t i = 0; i < 3; ++i)
    s.gig.add_vertex (
        { vid (i), gig::VertexKind::GPU, true, "g" + std::to_string (i) });
  for (uint32_t i = 0; i < 3; ++i)
    for (uint32_t j = 0; j < 3; ++j)
      if (i != j)
        s.gig.add_edge ({ vid (i), vid (j), 100.0, gig::LinkType::NVLink });

  std::vector<agent::LocalSubgraph> subs{ std::move (s) };
  auto                  g = td::assemble_global (subs);
  to::TransferOptimizer opt (g);

  cc::ChannelSpec spec;
  for (const auto &[id, v] : g.vertices ())
    if (v.kind == gig::VertexKind::GPU)
      {
        spec.src_eps.push_back (id);
        spec.dst_eps.push_back (id);
      }
  auto plan = cc::build_channel (g, opt, spec);
  EXPECT (plan.tables.size () == 3);
  for (const auto &[owner, tbl] : plan.tables) EXPECT (tbl.size () == 2);

  // No CPU/NIC vertices in any table key (forwarding-table is GPU-only).
  for (const auto &[owner, tbl] : plan.tables)
    EXPECT (g.vertex (owner).kind == gig::VertexKind::GPU);
}

} // anonymous namespace

int
main ()
{
  struct {
    const char *name;
    void (*fn) ();
  } tests[] = {
    { "max_flow_two_paths", case_max_flow_two_paths },
    { "max_flow_self_and_missing", case_max_flow_self_and_missing },
    { "concat_basic_and_mismatch", case_concat_basic_and_mismatch },
    { "forwarding_gpu_only_merge", case_forwarding_gpu_only_merge },
    { "assemble_global_v_net_fanout", case_assemble_global_v_net_fanout },
    { "to_cache_reuse_cross_server", case_to_cache_reuse_cross_server },
    { "codec_primitives_roundtrip", case_codec_primitives_roundtrip },
    { "codec_message_roundtrip", case_codec_message_roundtrip },
    { "cc_all_to_all_local", case_cc_all_to_all_local },
    { "pcie_share_gdr_off_routes_via_cpu",
        case_pcie_share_gdr_off_routes_via_cpu },
    { "pcie_share_gdr_on_prefers_gate", case_pcie_share_gdr_on_prefers_gate },
  };
  for (auto &t : tests)
    {
      int prev = g_fail;
      t.fn ();
      if (prev == g_fail)
        {
          std::printf ("PASS  %s\n", t.name);
          ++g_pass;
        }
    }
  std::printf ("\nunit summary: %d passed, %d failed\n", g_pass, g_fail);
  return g_fail == 0 ? 0 : 1;
}