1/* peer.rs: join_frame_and_checks, fleet_frames_round_trip, names_on_the_wire. */
2#include "mesh/abi/backend.h"
3#include "mesh/abi/group_state.h"
4#include "mesh/abi/peer_frames.h"
5#include "mesh/abi/peer_sizes.h"
6#include "mesh/abi/reconcile.h"
7#include "mesh/check.h"
8#include "mesh/peer_codec.h"
9#include "mesh/mt.h"
10
11#include <stdlib.h>
12
13static mpk_caps caps(void) {
14    mpk_caps c = {.cc = 121, .sms = 48, .sms_total = 48, .blocks_per_sm = 24, .gpu = "NVIDIA GB10", .arch = "aarch64", .host = "dgx-00", .max_threads = 1024, .max_shared = 101376};
15    return c;
16}
17
18static bool join_eq(const mpk_join *a, const mpk_join *b) { return a->system == b->system && a->version == b->version && a->systems == b->systems && mpk_caps_eq(&a->caps, &b->caps); }
19
20static bool rejects(const mpk_join *j, uint32_t me, uint32_t n, bool has, uint32_t exp, const char *needle) {
21    char *err = NULL;
22    if (mpk_join_check(j, me, n, has, exp, &err)) return false;
23    bool ok = !needle || mt_contains(err, needle);
24    free(err);
25    return ok;
26}
27
28int main(void) {
29    /* join_frame_and_checks */
30    mpk_join j = {.system = 2, .version = MESH_PEER_VERSION, .systems = 3, .caps = caps()};
31    j.caps.cc = 86, j.caps.blocks_per_sm = 16;
32    strcpy(j.caps.gpu, "NVIDIA GeForce RTX 3060");
33    strcpy(j.caps.arch, "x86_64");
34    mpk_join mj = j;
35    mj.caps = caps();
36    mj.caps.cc = 0, mj.caps.blocks_per_sm = 0, mj.caps.backend = mesh_backend_word(MESH_BACKEND_METAL, 9);
37    strcpy(mj.caps.gpu, "Apple M3 Ultra");
38    mpw_bytes f = mpk_join_encode(&mj);
39    mpk_join back = mpk_join_decode(f.b + 1);
40    CHECK(join_eq(&back, &mj));
41    mpw_bytes_free(&f);
42    f = mpk_join_encode(&j);
43    CHECK(f.b[0] == MESH_F_JOIN && f.n == 1 + MESH_JOIN_BYTES);
44    back = mpk_join_decode(f.b + 1);
45    CHECK(join_eq(&back, &j));
46    mpw_bytes_free(&f);
47    char *err = NULL;
48    CHECK(mpk_join_check(&j, 0, 3, true, 2, &err));
49    CHECK(rejects(&j, 0, 3, true, 1, NULL));
50    mpk_join j0 = j;
51    j0.system = 0;
52    CHECK(mpk_join_check(&j0, 1, 3, false, 0, &err));
53    CHECK(rejects(&j, 1, 3, false, 0, NULL));
54    CHECK(rejects(&j, 0, 4, true, 2, "cluster of 3"));
55    mpk_join j5 = j;
56    j5.version = 5;
57    CHECK(rejects(&j5, 0, 3, true, 2, "version 5"));
58
59    /* fleet_frames_round_trip */
60    mst_outcome moves[] = {{"pipehash", 41, true, 3100}, {"src", 42, false, 0}};
61    mpk_beat_status st = {.has_leader = true, .leader = 2, .applied = 9, .status = MESH_RS_FAILED, .up = 0xb, .group_epoch = 4, .group_state = MESH_GL_PAUSED,
62                          .reason = "load hello: jit: boom", .moves = moves, .nmoves = 2};
63    mfc_msg m = {.kind = MFC_BEAT, .term = 7, .leader = true, .commit = {6, 9}, .acc = {7, 10}};
64    f = mpk_beat_frame(&m, &st);
65    size_t n = strlen(st.reason);
66    CHECK(f.b[0] == MESH_F_BEAT && f.n == 1 + MESH_BEAT_HEAD + n + 4 + 2 * MESH_BEAT_MOVE_BYTES);
67    mfc_msg m2;
68    mpk_beat_status st2;
69    CHECK(mpk_beat_decode(f.b + 1, f.b + 1 + MESH_BEAT_HEAD, n, f.b + 1 + MESH_BEAT_HEAD + n + 4, 2 * MESH_BEAT_MOVE_BYTES, &m2, &st2, &err));
70    CHECK(m2.term == 7 && m2.leader && m2.commit.term == 6 && m2.commit.version == 9 && m2.acc.term == 7 && m2.acc.version == 10);
71    CHECK(mpk_beat_status_eq(&st2, &st));
72    mpk_beat_status_free(&st2);
73    mpw_bytes_free(&f);
74    mpk_hello spec = {.group = "train", .epoch = 3, .cycle = 5, .mask = 0xd, .order = {0, 2, 3, 0xff, 0xff, 0xff, 0xff, 0xff}, .hash = 0xfeed};
75    mpw_bytes h = mpk_hello_frame(2, &spec);
76    CHECK(h.n == 1 + MESH_HELLO_BYTES);
77    uint32_t sys, ver, r;
78    mpk_hello hb;
79    CHECK(mpk_hello_decode(h.b + 1, &sys, &ver, &hb, &err) && sys == 2 && ver == MESH_PEER_VERSION && mpk_hello_eq(&hb, &spec));
80    CHECK(mpk_hello_rank(&spec, 3, &r) && r == 2 && !mpk_hello_rank(&spec, 1, &r) && mpk_hello_n(&spec) == 3);
81    mpw_bytes_free(&h);
82    mfc_msg msgs[] = {{.kind = MFC_VOTE_REQ, .pre = true, .term = 3, .acc = {2, 8}}, {.kind = MFC_VOTE, .pre = false, .term = 3, .granted = true},
83                      {.kind = MFC_STATE_ACK, .term = 3, .ver = {3, 9}}, {.kind = MFC_COMMIT, .term = 3, .ver = {3, 9}}};
84    for (int i = 0; i < 4; i++) {
85        f = mpk_msg_frame(&msgs[i]);
86        size_t len = f.b[0] == MESH_F_VOTE_REQ ? MESH_VOTE_REQ_BYTES : f.b[0] == MESH_F_VOTE ? MESH_VOTE_BYTES : MESH_ACK_BYTES;
87        CHECK(f.n == 1 + len);
88        mfc_msg d = f.b[0] == MESH_F_VOTE_REQ ? mpk_vote_req_decode(f.b + 1) : f.b[0] == MESH_F_VOTE ? mpk_vote_decode(f.b + 1) : mpk_ack_decode(f.b[0], f.b + 1);
89        CHECK(d.kind == msgs[i].kind && d.term == msgs[i].term && d.pre == msgs[i].pre && d.granted == msgs[i].granted && d.acc.version == msgs[i].acc.version &&
90              d.ver.version == msgs[i].ver.version);
91        mpw_bytes_free(&f);
92    }
93    mfc_msg sm = {.kind = MFC_STATE, .term = 1};
94    char *empty = mfd_encode(&sm.doc);
95    f = mpk_msg_frame(&sm);
96    CHECK(f.n == 1 + MESH_STATE_HEAD + strlen(empty));
97    free(empty);
98    mpw_bytes_free(&f);
99    return mesh_check_done();
100}