pulsar/jasmin/threshold/round2.jazz

// Pulsar Round-2 - per-party response.
//
// Reference: spec/pulsar-m.tex Algorithm "Sign Round 2 (party i, attempt
// kappa)" (Algorithm alg:sign-r2). Mirrors ref/go/pkg/pulsar/threshold.go
// ThresholdSigner.Round2 semantics.
//
// -----------------------------------------------------------------------------
// Algorithm (spec alg:sign-r2)
// -----------------------------------------------------------------------------
//
// Inputs:
//   - {D_j} received from peers + corresponding MACs (verified in step 2)
//   - y_state from Round-1: y_i (l polynomials), w_i (k polys), w1_i (k polys)
//   - share s_i (polyvecl, secret), blind u_i (polyvecl, secret)
//   - eval_points (public per-party x_j in F_q), my_index, quorum_size
//   - mu (FIPS 204 message hash binder)
//
// Steps:
//   1. Verify peer MACs against locally-known K_{i,j}. Abort on mismatch.
//   2. Compute bar_D = cSHAKE256("PULSAR-SIGN-R2-V1" || D_1 || ... || D_n).
//   3. Recover bar_w = HighBits(sum_j w_j, 2*gamma_2) by summing bar_w_j.
//   4. Compute c_tilde = SHAKE(mu || bar_w, 2*lambda).
//   5. Expand c = SampleInBall(c_tilde) in R_q (sparse, +/-1 coefficients).
//   6. Compute lambda_i^T in Z_q (Lagrange coefficient at X = 0).
//   7. z_i = y_i + c * lambda_i^T * s_i.
//   8. r_i = c * lambda_i^T * u_i (used by aggregator for hint).
//   9. Broadcast (bar_w_i, z_i, r_i) + tau_2 integrity hash.
//
// Output layout (PULSARM_RESPONSE_BYTES = 9300):
//   [0    .. 32)   tau_2 digest
//   [32   .. 1632) pack_z(z_i)         (Li2_l * Li2_pack_zlen = 1600)
//   [1632 .. 7776) r_i u32 stream      (Li2_k * Li2_polydeg * 4 = 6144)
//   [7776 .. 8544) pack_w1(my w1)      (Li2_pack_w1len = 768)
//   [8544 .. 9216) per-peer MAC set    (PULSARM_QUORUM_MAX * 32 = 672)
//   [9216 .. 9300) reserved (84 bytes)
//
// Constant-time obligations:
//   - Time and memory-access patterns independent of share, blind, y_state.
//   - Lagrange coefficient inputs (eval_points) are PUBLIC; CT obligation
//     on lagrange_coefficient_mont is structural, not data-dependent.
//   - MAC verification uses ct_eq_32 (constant-time byte compare).
//   - The MAC-mismatch abort path is part of the public spec: the caller
//     learns *which* peer's MAC failed (identifiable-abort property under
//     synchronous network assumptions).

require "../lib/pulsar_params.jinc"
require "../lib/lagrange.jinc"
require "../lib/polyvec_scalarmul.jinc"
require "../lib/mac.jinc"
require "../lib/transcript.jinc"

from Jade require "crypto_sign/dilithium/dilithium3/amd64/ref/params.jinc"
from Jade require "crypto_sign/dilithium/dilithium3/amd64/ref/common.jinc"
from Jade require "crypto_sign/dilithium/dilithium3/amd64/ref/packing.jinc"

// -----------------------------------------------------------------------------
// Recompute one peer's expected MAC:
//   expected = KMAC256(K, commit || tau1, 32, "PULSAR-SIGN-R1-MAC-V1").
//
// `mac_data` is the contiguous 4-table buffer (see verify_peer_macs);
// `base` = peer_index * 32 is the per-peer byte offset within each table.
// The key / commit / tau sub-slices for this peer are copied into local
// 32-byte buffers, then KMAC256 is run over commit || tau1 keyed by the
// session key, squeezing the 32-byte tag into `expected_buf`.
//
// Real (stack-return) call: its keccak `state` (with the runtime-offset
// byte view state[:u8 offset] inside kmac256_finalize_32) and theta
// working set live in their own frame, so the caller verify_peer_macs
// only needs to preserve its small loop state across the call rather than
// hosting theta alongside it. Taking mac_data + base (rather than three
// pre-sliced pointers) keeps both sides' working sets small.
// Semantics/CT unchanged.
#[returnaddress="stack"]
fn r2_one_expected_mac(
    reg ptr u8[4 * PULSARM_QUORUM_MAX * 32] mac_data,
    reg u64 base,
    reg ptr u8[32] expected_buf)
    -> reg ptr u8[32]
{
    stack u8[32] key_buf;
    stack u8[32] commit_buf;
    stack u8[32] tau_buf;
    stack u64[25] state;
    reg ptr u64[25] state_rsp;
    reg ptr u8[32] key_rsp, commit_rsp, tau_rsp, exp_rsp;
    reg ptr u8[32] src_slice;
    reg u64 offset, i;
    reg u8 c;

    // Copy this peer's key / commit / tau1 (sub-regions [0]/[1]/[2]).
    // Take a 32-byte reg-ptr slice of the peer's row in each sub-region,
    // then index by i alone — addressing the stack array with two runtime
    // index terms (sub-region const + base + i) is not an encodable x86
    // address, whereas a runtime-offset slice plus single-index load is.
    src_slice = mac_data[0 * PULSARM_QUORUM_MAX * 32 + base : 32];
    ?{}, i = #set0_64();
    while (i < 32) { c = src_slice[i]; key_buf[i]    = c; i += 1; }
    src_slice = mac_data[1 * PULSARM_QUORUM_MAX * 32 + base : 32];
    ?{}, i = #set0_64();
    while (i < 32) { c = src_slice[i]; commit_buf[i] = c; i += 1; }
    src_slice = mac_data[2 * PULSARM_QUORUM_MAX * 32 + base : 32];
    ?{}, i = #set0_64();
    while (i < 32) { c = src_slice[i]; tau_buf[i]    = c; i += 1; }

    state_rsp = state;
    key_rsp = key_buf;
    state_rsp = kmac256_init_R1MAC(state_rsp, key_rsp);

    ?{}, offset = #set0_64();
    commit_rsp = commit_buf;
    tau_rsp = tau_buf;
    state_rsp, offset = absorb_array32(state_rsp, offset, commit_rsp);
    state_rsp, offset = absorb_array32(state_rsp, offset, tau_rsp);

    exp_rsp = expected_buf;
    state_rsp, exp_rsp = kmac256_finalize_32(state_rsp, offset, exp_rsp);
    expected_buf = exp_rsp;
    return expected_buf;
}

// -----------------------------------------------------------------------------
// MAC verification: for each peer j != my_index, recompute the expected
// KMAC256(K_{i,j}, D_j || tau_1_j, 32, "PULSAR-SIGN-R1-MAC-V1") and
// constant-time compare against the received MAC.
//
// Returns 0 on full success; j+1 (encoding the first mismatched peer
// index) on failure, so the caller can construct a ComplaintMACFailure
// abort evidence with `accused = j`.
//
// Compiled as a real (stack-return-address) call. The whole MAC region —
// this loop, the per-peer buffers and the inlined r2_one_expected_mac
// keccak — then lives in this function's own frame rather than flattening
// into pulsar_round2 alongside its eleven unpacked caller pointers and
// other hashes. The four per-peer byte tables are passed as ONE
// contiguous buffer (mac_data, sub-regions [0]=mac_keys, [1]=commits,
// [2]=tau1_digests, [3]=received_macs, each PULSARM_QUORUM_MAX*32 bytes),
// and the two public u32 scalars are packed into a single u64 (quorum_size
// in the high half, my_index in the low half), so the function needs only
// two register parameters. The per-iteration working set — the inlined
// r2_one_expected_mac keccak's theta plus mac_data, quorum_size, my_index,
// the loop counter and status — sits right at the register-file limit, so
// every freed parameter register matters. The runtime-offset keccak state
// view is confined to this straight-line frame, so the stack-allocation
// oracle accepts it. Semantics/CT unchanged (quorum_size / my_index public).
//
// Real (stack-return) call: keeps this loop and its call to
// r2_one_expected_mac in a frame of its own, off pulsar_round2's heavily
// populated frame (which already hosts the unpacked caller pointers and
// the other digest calls).
#[returnaddress="stack"]
fn verify_peer_macs(
    reg ptr u8[4 * PULSARM_QUORUM_MAX * 32] mac_data,
    reg u64 qm_packed)
    -> reg u64
{
    stack u8[32] received_buf;
    stack u8[32] expected_buf;
    reg ptr u8[32] recv_rsp, exp_rsp, recv_slice;
    // quorum_size / my_index are loop-invariant and only read at the loop
    // boundary, never inside the inlined keccak. Hold them in stack slots
    // so they do not occupy registers across the theta-heavy MAC
    // recomputation (that is what pushed the frame one register over).
    stack u32 quorum_size_s;
    stack u32 my_index_s;
    stack u32 j_s;
    stack u64 status_s;
    reg u32 quorum_size;
    reg u32 my_index;
    reg u64 tmp64;
    reg u32 j;
    reg u64 base;
    reg u64 i;
    reg u8 c;
    reg u8 eq;
    reg u64 status;

    // Unpack the two public scalars (high 32 = quorum_size, low = my_index).
    tmp64 = qm_packed;
    tmp64 >>= 32;
    quorum_size = (32u) tmp64;
    quorum_size_s = quorum_size;
    my_index = (32u) qm_packed;
    my_index_s = my_index;

    // status AND the loop counter j live in stack slots across the
    // per-iteration r2_one_expected_mac call: every value kept in a
    // (callee-saved) register across that call competes with the callee's
    // 10-register keccak working set in Jasmin's global allocator, which
    // is what exhausts the file. After this, only mac_data (the call's own
    // argument) and the call's result pointer cross the boundary. base is
    // likewise recomputed from j rather than held.
    ?{}, status = #set0_64();
    status_s = status;
    ?{}, j = #set0_32();
    j_s = j;
    quorum_size = quorum_size_s;
    while (j < quorum_size) {
        // The per-peer MAC recomputation is done unconditionally for
        // every j (including j == my_index), and only the accusation
        // update is gated on j != my_index. (Self index excluded from
        // accusations.) j / my_index are public, so this introduces no
        // secret-dependent behaviour; the access pattern is uniform.
        // base = j*32; each table t starts at t*PULSARM_QUORUM_MAX*32.
        base = (64u) j;
        base *= 32;

        // Recompute this peer's expected MAC in its own frame. j is in a
        // stack slot across the call (see note above).
        exp_rsp = expected_buf;
        exp_rsp = r2_one_expected_mac(mac_data, base, exp_rsp);
        expected_buf = exp_rsp;

        // Received MAC is sub-region [3]. Reload j and recompute base
        // (neither held across the call). Slice the peer's row, then index
        // by i (two runtime index terms into the stack array would not be
        // an encodable address; a runtime-offset slice is).
        j = j_s;
        base = (64u) j;
        base *= 32;
        recv_slice = mac_data[3 * PULSARM_QUORUM_MAX * 32 + base : 32];
        ?{}, i = #set0_64();
        while (i < 32) { c = recv_slice[i]; received_buf[i] = c; i += 1; }

        recv_rsp = received_buf;
        exp_rsp = expected_buf;
        eq = ct_eq_32(exp_rsp, recv_rsp);

        // Cryptographic declassification: the MAC verification
        // outcome is a publicly-observable bit per the Pulsar
        // protocol — wrong MAC ⇒ abort with accusation index. This
        // is a protocol-level acceptance: the verifier learns
        // accept/reject, which is exactly what `eq` encodes.
        #declassify(eq);
        // Reload the loop-invariant scalars and status now (after the
        // keccak call), so they were not pinned to registers across it.
        my_index = my_index_s;
        status = status_s;
        if (j != my_index) {
            if (eq == 0) {
                if (status == 0) {
                    status = (64u) j;
                    status += 1;
                }
            }
        }
        status_s = status;
        j = j_s;
        j += 1;
        j_s = j;
        quorum_size = quorum_size_s;
    }
    status = status_s;
    return status;
}

// -----------------------------------------------------------------------------
// Aggregate bar_D = cSHAKE256("PULSAR-SIGN-R2-V1" || D_1 || ... || D_n).
// -----------------------------------------------------------------------------
// Stack-return call (not inlined): keeps its keccak `state` working set
// in its own frame instead of flattening into pulsar_round2 alongside the
// other inlined hashes and the eleven unpacked caller pointers, which
// together overflow the register file. Semantics/CT unchanged.
#[returnaddress="stack"]
fn r2_aggregate_commit_digest(
    reg ptr u8[PULSARM_QUORUM_MAX * 32] commits,
    reg u32 quorum_size,
    reg ptr u8[32] out)
    -> reg ptr u8[32]
{
    stack u64[25] state;
    reg ptr u64[25] state_rsp;
    // out and the byte count are loop-invariant / used only at the end;
    // hold them in stack slots so they do not occupy registers during the
    // keccak permutations driven inside the absorb loop (theta's working
    // set is large; the function is otherwise the same shape as round1's
    // commit-digest helper).
    stack u64 nbytes_s;
    reg u64 offset;
    reg u64 i;
    reg u64 nbytes;
    reg u8 c;

    // The j-loop (per peer) over 32-byte commits is flattened to one byte
    // loop over commits[0 .. quorum_size*32); the absorbed stream is
    // identical (commits are contiguous, peer j at offset j*32).
    nbytes = (64u) quorum_size;
    nbytes *= 32;
    nbytes_s = nbytes;

    state_rsp = state;
    state_rsp, offset = cshake256_absorb_prelude_R2(state_rsp);

    ?{}, i = #set0_64();
    nbytes = nbytes_s;
    while (i < nbytes) {
        c = commits[i];
        state_rsp, offset = absorb_one_byte(state_rsp, offset, c);
        i += 1;
        nbytes = nbytes_s;
    }

    state_rsp = cshake256_finalize(state_rsp, offset);

    ?{}, i = #set0_64();
    while (i < 32) { c = state_rsp[:u8 i]; out[i] = c; i += 1; }
    return out;
}

// -----------------------------------------------------------------------------
// Compute c_tilde = cSHAKE256("PULSAR-SIGN-R2-V1" || mu || pack_w1(bar_w))
// truncated to 32 bytes (= 2*lambda for L3 / ML-DSA-65).
// -----------------------------------------------------------------------------
// Stack-return call (not inlined), same register-frame rationale as
// r2_aggregate_commit_digest. Semantics/CT unchanged.
#[returnaddress="stack"]
fn r2_challenge_seed(
    reg ptr u8[64] mu,
    reg ptr u8[Li2_pack_w1len] bar_w_packed,
    reg ptr u8[32] c_tilde_out)
    -> reg ptr u8[32]
{
    stack u64[25] state;
    reg ptr u64[25] state_rsp;
    reg u64 offset;
    reg u64 i;
    reg u8 c;

    state_rsp = state;
    state_rsp, offset = cshake256_absorb_prelude_R2(state_rsp);

    // Absorb mu (64) then pack_w1(bar_w) (768). The bar_w loop crosses the
    // 136-byte rate boundary several times, so absorb_one_byte runs keccak
    // mid-loop. Keeping the squeeze-output pointer in a register across
    // that theta-heavy stretch competes with the permutation's working
    // set; index bar_w_packed directly and reload from stack as needed so
    // only the absorb cursor is live across the permutations.
    ?{}, i = #set0_64();
    while (i < 64) {
        c = mu[i];
        state_rsp, offset = absorb_one_byte(state_rsp, offset, c);
        i += 1;
    }
    ?{}, i = #set0_64();
    while (i < Li2_pack_w1len) {
        c = bar_w_packed[i];
        state_rsp, offset = absorb_one_byte(state_rsp, offset, c);
        i += 1;
    }
    state_rsp = cshake256_finalize(state_rsp, offset);

    ?{}, i = #set0_64();
    while (i < 32) { c = state_rsp[:u8 i]; c_tilde_out[i] = c; i += 1; }
    return c_tilde_out;
}

// -----------------------------------------------------------------------------
// tau_2 = cSHAKE256("PULSAR-SIGN-R2-V1" || bar_D || mu || pack_w1(bar_w)),
// truncated to 32 bytes. Stack-return call (not inlined) so its keccak
// `state` working set stays in its own frame rather than flattening into
// pulsar_round2 with the other hashes and the eleven unpacked caller
// pointers. Absorbed byte stream and values are identical to the
// previous inline body; CT unchanged.
// -----------------------------------------------------------------------------
#[returnaddress="stack"]
fn r2_tau2_digest(
    reg ptr u8[32] bar_D,
    reg ptr u8[64] mu,
    reg ptr u8[Li2_pack_w1len] bar_w_packed,
    reg ptr u8[32] tau_out)
    -> reg ptr u8[32]
{
    stack u64[25] state;
    reg ptr u64[25] state_rsp;
    reg ptr u8[32] bar_D_rsp;
    reg u64 offset;
    reg u64 i;
    reg u8 c;

    state_rsp = state;
    state_rsp, offset = cshake256_absorb_prelude_R2(state_rsp);
    bar_D_rsp = bar_D;
    state_rsp, offset = absorb_array32(state_rsp, offset, bar_D_rsp);
    ?{}, i = #set0_64();
    while (i < 64) { c = mu[i]; state_rsp, offset = absorb_one_byte(state_rsp, offset, c); i += 1; }
    ?{}, i = #set0_64();
    while (i < Li2_pack_w1len) {
        c = bar_w_packed[i];
        state_rsp, offset = absorb_one_byte(state_rsp, offset, c);
        i += 1;
    }
    state_rsp = cshake256_finalize(state_rsp, offset);
    ?{}, i = #set0_64();
    while (i < 32) { c = state_rsp[:u8 i]; tau_out[i] = c; i += 1; }
    return tau_out;
}

// -----------------------------------------------------------------------------
// Round-2 main entry point.
//
// Inputs (all u64 pointers, caller-supplied buffers):
//   y_state_ptr      PULSARM_Y_STATE_BYTES bytes (from Round-1)
//   share_ptr        PULSARM_SHARE_BYTES bytes (polyvecl, u32 stride)
//   blind_ptr        PULSARM_BLIND_BYTES bytes (polyvecl, u32 stride)
//   peer_commits_ptr PULSARM_QUORUM_MAX * 32 bytes (D_j slots)
//   peer_bar_w_ptr   PULSARM_QUORUM_MAX * (Li2_k * Li2_polydeg * 4) bytes
//                    (per-peer unpacked bar_w_j as u32; caller exchanges
//                    raw bar_w out-of-band per spec remark "commit digest
//                    is bound under MAC, not concealment").
//   peer_macs_ptr    PULSARM_QUORUM_MAX * 32 bytes (MACs from j -> me)
//   peer_tau_ptr     PULSARM_QUORUM_MAX * 32 bytes (tau_1_j slots)
//   mac_keys_ptr     PULSARM_QUORUM_MAX * 32 bytes (K_{i,j} session keys)
//   eval_points_ptr  PULSARM_QUORUM_MAX * 4 bytes (u32 evaluation points)
//   mu_ptr           64 bytes (FIPS 204 mu)
//   quorum_size      u32 |T|
//   my_index         u32 (0 .. quorum_size)
//   response_out_ptr PULSARM_RESPONSE_BYTES bytes (output)
//
// Returns u64 status:
//   0     = success
//   j+1   = MAC mismatch on peer j (accused index encoded)
//
// CT signature: all pointer values are public addresses (no
// secret-address branching). Pointed-to contents have the following
// sensitivity per protocol spec:
//   y_state_ptr      → SECRET (per-attempt y state held by this party)
//   share_ptr        → SECRET (per-party Shamir share)
//   blind_ptr        → SECRET (per-party Pedersen blind)
//   peer_commits_ptr → public (Round-1 commits already on the wire)
//   peer_bar_w_ptr   → public (Round-1 aggregated w-bars)
//   peer_macs_ptr    → public (peer-MAC tags from Round-1)
//   peer_tau_ptr     → public (per-peer tau_1 digests)
//   mac_keys_ptr     → SECRET (session keys K_{i,j})
//   eval_points_ptr  → public (Lagrange evaluation points)
//   mu_ptr           → public (FIPS 204 message hash)
//   quorum_size,
//   my_index         → public
//   response_out_ptr → writes public Round-2 response (z, r, hint)
//
// The CT signature itself only constrains pointer-publicity (used for
// branching / load addresses). Per-buffer content sensitivity is
// encoded by #declassify pragmas at the points where the protocol
// declassifies them (e.g., w1 high-bits, signature challenge).
// ABI note: the x86-64 SysV calling convention Jasmin targets passes at
// most six arguments in registers, and Jasmin export functions require
// every argument to be register-resident. The Round-2 entry point needs
// eleven caller buffers plus two scalars, so the eleven pointers are
// marshalled by the caller into one contiguous array of u64 addresses
// (`ptrs_ptr`), in the fixed order documented below, and unpacked here.
// This mirrors round1.jazz's "pack the public inputs into one blob"
// convention. Only the addresses move into the blob — each buffer's
// contents and sensitivity are exactly as before, so the per-buffer CT
// classification (encoded via #declassify at the protocol declassify
// points) and all computed values are unchanged.
//
// ptrs_ptr layout (each entry is a u64 little-endian pointer):
//   [ 0] y_state_ptr        [ 1] share_ptr        [ 2] blind_ptr
//   [ 3] peer_commits_ptr   [ 4] peer_bar_w_ptr   [ 5] peer_macs_ptr
//   [ 6] peer_tau_ptr       [ 7] mac_keys_ptr     [ 8] eval_points_ptr
//   [ 9] mu_ptr             [10] response_out_ptr
//
// CT: ptrs_ptr and quorum_size / my_index are public.
#[ct = "public * public * public -> public"]
export
fn pulsar_round2(
    reg u64 ptrs_ptr,
    reg u32 quorum_size,
    reg u32 my_index)
    -> reg u64
{
    reg u64 y_state_ptr;
    reg u64 share_ptr;
    reg u64 blind_ptr;
    reg u64 peer_commits_ptr;
    reg u64 peer_bar_w_ptr;
    reg u64 peer_macs_ptr;
    reg u64 peer_tau_ptr;
    reg u64 mac_keys_ptr;
    reg u64 eval_points_ptr;
    reg u64 mu_ptr;
    reg u64 response_out_ptr;
    // Spilled pointer-table base, so the two late-used pointers can be
    // reloaded lazily instead of staying live across the middle.
    stack u64 ptrs_base;
    reg u64 bw_base;
    // status spilled across the step 3-11 body (see its use site).
    stack u64 status_round2_s;
    // quorum_size / my_index spilled: they are read in several places
    // across the digest-call-heavy body, and a register pinned across any
    // keccak call there competes with the (globally allocated) keccak
    // permutation's working set. Reloaded from the stack before each use.
    stack u32 qs_s;
    stack u32 mi_s;
    // Workspace.
    stack u8[64] mu;
    stack u8[32] bar_D;
    stack u8[32] c_tilde;
    stack u8[32] tau2_digest;
    stack u8[Li2_pack_w1len] bar_w_packed;
    stack u32[Li2_l * Li2_polydeg] y_i;
    stack u32[Li2_l * Li2_polydeg] s_i;
    stack u32[Li2_l * Li2_polydeg] u_i;
    stack u32[Li2_l * Li2_polydeg] tmp_l;
    stack u32[Li2_l * Li2_polydeg] cs_l;
    stack u32[Li2_l * Li2_polydeg] cu_l;
    stack u32[Li2_l * Li2_polydeg] scaled_l;
    stack u32[Li2_l * Li2_polydeg] z_i;
    stack u32[Li2_l * Li2_polydeg] r_i;
    stack u32[Li2_polydeg] c_poly;
    stack u32[Li2_polydeg] c_ntt;
    stack u32[Li2_k * Li2_polydeg] bar_w;
    stack u32[Li2_k * Li2_polydeg] my_w1;
    // The four per-peer byte tables are held in one contiguous buffer so
    // verify_peer_macs can take a single pointer (keeping its register-
    // parameter count low; see its header). Sub-region layout, each
    // PULSARM_QUORUM_MAX*32 bytes:
    //   [0] mac_keys  [1] commits  [2] tau1_digests  [3] received_macs
    stack u8[4 * PULSARM_QUORUM_MAX * 32] peer_mac_data;
    stack u32[PULSARM_QUORUM_MAX] eval_pts;

    reg ptr u8[64] mu_rsp;
    reg ptr u8[32] bar_D_rsp, c_tilde_rsp, tau2_rsp;
    reg ptr u8[Li2_pack_w1len] packed_rsp;
    reg ptr u8[PULSARM_QUORUM_MAX * 32] commits_rsp;
    reg ptr u8[4 * PULSARM_QUORUM_MAX * 32] mac_data_rsp;
    reg ptr u32[Li2_l * Li2_polydeg] vec_l_rsp;
    reg ptr u32[Li2_l * Li2_polydeg] vec_l_out;
    reg ptr u32[Li2_k * Li2_polydeg] vec_k_rsp;
    reg ptr u32[Li2_polydeg] pc_rsp, pn_rsp;
    reg ptr u32[PULSARM_QUORUM_MAX] eval_rsp;
    reg ptr u64[25] state_rsp;

    reg u64 i, offset, base;
    reg u64 qm_packed, tmp_qm;
    reg u32 j;
    // Separate scratch bytes for the public-input load loops vs the
    // secret-content load loop (mac_keys / share / blind). Sharing
    // one `c` would join all loop levels to `secret` and taint the
    // genuinely-public per-peer commits / MACs / tau1 digests.
    reg u8 c;
    reg u8 c_sec;
    reg u32 r32, sum32;
    reg u32 lambda_mont;
    reg u64 status;

    // ---- 0. Unpack the caller-supplied pointer table (see ABI note) ----
    // Each entry is a u64 address. Pointers used only by the early input-
    // loading phase are unpacked now; peer_bar_w_ptr (step 4) and
    // response_out_ptr (step 11) are re-loaded lazily from the spilled
    // table right before their use, so they do not stay live across the
    // hash-heavy middle and inflate register pressure there.
    ptrs_base = ptrs_ptr;
    y_state_ptr      = [:u64 ptrs_ptr +  0];
    share_ptr        = [:u64 ptrs_ptr +  8];
    blind_ptr        = [:u64 ptrs_ptr + 16];
    peer_commits_ptr = [:u64 ptrs_ptr + 24];
    peer_macs_ptr    = [:u64 ptrs_ptr + 40];
    peer_tau_ptr     = [:u64 ptrs_ptr + 48];
    mac_keys_ptr     = [:u64 ptrs_ptr + 56];
    eval_points_ptr  = [:u64 ptrs_ptr + 64];
    mu_ptr           = [:u64 ptrs_ptr + 72];

    // Spill the two public scalars immediately (reloaded before each use).
    qs_s = quorum_size;
    mi_s = my_index;

    // ---- 1. Load inputs ----
    ?{}, i = #set0_64();
    while (i < 64) { c = [:u8 mu_ptr + i]; mu[i] = c; i += 1; }

    // y_i = first Li2_l polynomials of y_state.
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) {
        r32 = [:u32 y_state_ptr + 4 * i];
        y_i[i] = r32;
        i += 1;
    }
    // my_w1 = last Li2_k polynomials of y_state (offset = (Li2_l + Li2_k) polynomials).
    // Fold the base pointer and the constant poly offset into one register
    // so each load is [base + index*4]; base-pointer + offset-register +
    // index*scale is two base registers, which x86 cannot encode.
    base = (64u) (Li2_l + Li2_k) * Li2_polydeg * 4;
    base += y_state_ptr;
    ?{}, i = #set0_64();
    while (i < Li2_k * Li2_polydeg) {
        r32 = [:u32 base + 4 * i];
        my_w1[i] = r32;
        i += 1;
    }

    // share, blind.
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) {
        r32 = [:u32 share_ptr + 4 * i];
        s_i[i] = r32;
        i += 1;
    }
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) {
        r32 = [:u32 blind_ptr + 4 * i];
        u_i[i] = r32;
        i += 1;
    }

    // eval_points (public).
    quorum_size = qs_s;
    ?{}, j = #set0_32();
    while (j < quorum_size) {
        base = (64u) j;
        r32 = [:u32 eval_points_ptr + 4 * base];
        eval_pts[base] = r32;
        j += 1;
    }

    // Per-peer buffers (commits, MACs, tau are public-by-protocol;
    // mac_keys are secret session keys K_{i,j}). Use a separate
    // scratch byte for the secret mac_keys load so jasmin-ct does
    // not back-propagate secrecy onto the public commit / mac / tau
    // buffers via a shared `c` register. The four tables are stored in
    // the contiguous peer_mac_data sub-regions described at its decl.
    ?{}, i = #set0_64();
    while (i < PULSARM_QUORUM_MAX * 32) {
        c_sec = [:u8 mac_keys_ptr + i];     peer_mac_data[0 * PULSARM_QUORUM_MAX * 32 + i] = c_sec;
        c     = [:u8 peer_commits_ptr + i]; peer_mac_data[1 * PULSARM_QUORUM_MAX * 32 + i] = c;
        c     = [:u8 peer_tau_ptr + i];     peer_mac_data[2 * PULSARM_QUORUM_MAX * 32 + i] = c;
        c     = [:u8 peer_macs_ptr + i];    peer_mac_data[3 * PULSARM_QUORUM_MAX * 32 + i] = c;
        i += 1;
    }

    // ---- 2. Verify peer MACs ----
    // Pack the two public scalars into one u64 (quorum_size high, my_index
    // low) so verify_peer_macs stays within two register parameters.
    mac_data_rsp = peer_mac_data;
    quorum_size = qs_s;
    qm_packed = (64u) quorum_size;
    qm_packed <<= 32;
    my_index = mi_s;
    tmp_qm = (64u) my_index;
    qm_packed |= tmp_qm;
    status = verify_peer_macs(mac_data_rsp, qm_packed);
    // Spill status to a stack slot: it must stay live across the entire
    // step 3-11 body (it gates the body and is the return value), and a
    // register pinned across every digest call in that body competes with
    // each callee's keccak working set in the global allocator. It is not
    // modified inside the body, so we just reload it at the return.
    status_round2_s = status;
    // Jasmin doesn't accept early `return` inside an `if`; fall through
    // and gate the rest of the body on `status == 0`. The MAC-failure
    // path leaves the output buffer in an indeterminate state, which
    // the caller MUST detect via the nonzero return code.
    if (status == 0) {

    // ---- 3. Aggregate bar_D ----
    // commits are sub-region [1] of peer_mac_data.
    commits_rsp = peer_mac_data[1 * PULSARM_QUORUM_MAX * 32 : PULSARM_QUORUM_MAX * 32];
    bar_D_rsp = bar_D;
    quorum_size = qs_s;
    bar_D_rsp = r2_aggregate_commit_digest(commits_rsp, quorum_size, bar_D_rsp);

    // ---- 4. Aggregate bar_w = sum_j bar_w_j (mod q) ----
    // For my own party use the on-stack my_w1; for peers, unpack from
    // peer_bar_w_ptr at stride (Li2_k * Li2_polydeg * 4) bytes per peer.
    ?{}, i = #set0_64();
    while (i < Li2_k * Li2_polydeg) {
        bar_w[i] = my_w1[i];
        i += 1;
    }
    // Reload peer_bar_w_ptr lazily here (entry [4] of the pointer table),
    // so it is not live across the earlier MAC-verify / digest phase.
    peer_bar_w_ptr = ptrs_base;
    peer_bar_w_ptr = [:u64 peer_bar_w_ptr + 32];
    quorum_size = qs_s;
    my_index = mi_s;
    ?{}, j = #set0_32();
    while (j < quorum_size) {
        if (j != my_index) {
            // Fold the buffer base and the per-peer stride into one
            // register so the inner load is an encodable [base + index*4]
            // mode (base + per-peer-offset + index*scale would be two base
            // registers plus index, which x86 cannot encode).
            bw_base = (64u) j;
            bw_base *= Li2_k * Li2_polydeg * 4;
            bw_base += peer_bar_w_ptr;
            ?{}, i = #set0_64();
            while (i < Li2_k * Li2_polydeg) {
                r32 = [:u32 bw_base + 4 * i];
                sum32 = bar_w[i];
                sum32 += r32;
                // Reduce mod 2^32; canonical reduction happens at pack_w1 stage
                // since w1 values are bounded by (q-1)/(2*gamma_2) per FIPS 204.
                bar_w[i] = sum32;
                i += 1;
            }
        }
        j += 1;
    }

    // Pack bar_w for c_tilde input and for transcript binding.
    // Assign the pack_w1 result (a fresh stack array) directly into the
    // stack buffer rather than bouncing through a reg-ptr; routing an
    // inline function's returned stack array through a reg ptr and back
    // into a stack array trips jasminc's variable-allocation oracle
    // ("variable already set"). Direct stack-to-stack assignment is the
    // pattern libjade itself uses and computes the identical bytes.
    vec_k_rsp = bar_w;
    bar_w_packed = pack_w1(vec_k_rsp);

    // ---- 5. c_tilde = SHAKE(mu || pack_w1(bar_w)) ----
    mu_rsp = mu;
    packed_rsp = bar_w_packed;
    c_tilde_rsp = c_tilde;
    c_tilde_rsp = r2_challenge_seed(mu_rsp, packed_rsp, c_tilde_rsp);
    // Cryptographic declassification: c_tilde is the FIPS 204
    // challenge hash, public by §3.6 (it's the part of the signature
    // every verifier sees). The Pulsar protocol publishes it as part
    // of the Round-2 response payload. Marking it public here lets
    // sampleInBall accept it; sampleInBall's output (c_poly) flows
    // through public NTT into mat-vec mul against the secret share.
    c_tilde = c_tilde_rsp;
    #declassify(c_tilde);
    c_tilde_rsp = c_tilde;

    // ---- 6. c = sampleInBall(c_tilde) -> NTT ----
    pc_rsp = c_poly;
    pc_rsp = sampleInBall(pc_rsp, c_tilde_rsp);
    c_poly = pc_rsp;

    ?{}, i = #set0_64();
    while (i < Li2_polydeg) { r32 = c_poly[i]; c_ntt[i] = r32; i += 1; }
    pn_rsp = c_ntt;
    pn_rsp = fft(pn_rsp);
    c_ntt = pn_rsp;

    // ---- 7. lambda_i^T in Montgomery form ----
    // Reload the scalars (after sampleInBall's keccak).
    eval_rsp = eval_pts;
    my_index = mi_s;
    quorum_size = qs_s;
    lambda_mont = lagrange_coefficient_mont(eval_rsp, my_index, quorum_size);

    // ---- 8. z_i = y_i + (c * lambda_i * s_i) ----
    // a) NTT(s_i)
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) { tmp_l[i] = s_i[i]; i += 1; }
    vec_l_rsp = tmp_l;
    vec_l_rsp = fft_vec(vec_l_rsp);
    tmp_l = vec_l_rsp;
    // b) c_ntt * s_i_ntt (libjade's mult_scalar_vec returns the
    //    product polyvecl by value via a stack temporary)
    pn_rsp = c_ntt;
    cs_l = mult_scalar_vec(pn_rsp, tmp_l);
    // c) invNTT -> normal form (Montgomery factor folded in)
    vec_l_rsp = cs_l;
    vec_l_rsp = ifft_to_mont_vecl(vec_l_rsp);
    cs_l = vec_l_rsp;
    // d) scale by lambda_i (Montgomery form). Read from cs_l, write to a
    // dedicated scratch buffer scaled_l, then copy back. Distinct input
    // and output regions are required because handing one reg ptr to both
    // the `v` and `out` parameters of polyvecl_scalar_mont tripped
    // jasminc's variable-allocation oracle, while aliasing the output onto
    // another already-aliased scratch (tmp_l) made its region "partial".
    // out[k]=REDC(lambda*v[k]) is per-coefficient; values/CT unchanged.
    vec_l_rsp = cs_l;
    vec_l_out = scaled_l;
    vec_l_out = polyvecl_scalar_mont(lambda_mont, vec_l_rsp, vec_l_out);
    cs_l = vec_l_out;
    // e) z_i = y_i + cs_l
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) {
        r32 = y_i[i];
        r32 += cs_l[i];
        z_i[i] = r32;
        i += 1;
    }
    // f) canonical reduction
    vec_l_rsp = z_i;
    vec_l_rsp = polyvecl_caddq(vec_l_rsp);
    z_i = vec_l_rsp;

    // ---- 9. r_i = c * lambda_i * u_i ----
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) { tmp_l[i] = u_i[i]; i += 1; }
    vec_l_rsp = tmp_l;
    vec_l_rsp = fft_vec(vec_l_rsp);
    tmp_l = vec_l_rsp;

    pn_rsp = c_ntt;
    cu_l = mult_scalar_vec(pn_rsp, tmp_l);
    vec_l_rsp = cu_l;
    vec_l_rsp = ifft_to_mont_vecl(vec_l_rsp);
    cu_l = vec_l_rsp;
    // Read cu_l, write scaled_l, copy back (see step 8d note).
    vec_l_rsp = cu_l;
    vec_l_out = scaled_l;
    vec_l_out = polyvecl_scalar_mont(lambda_mont, vec_l_rsp, vec_l_out);
    cu_l = vec_l_out;
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) { r_i[i] = cu_l[i]; i += 1; }
    vec_l_rsp = r_i;
    vec_l_rsp = polyvecl_caddq(vec_l_rsp);
    r_i = vec_l_rsp;

    // ---- 10. tau_2 = cSHAKE("PULSAR-SIGN-R2-V1" || bar_D || mu || pack_w1(bar_w)) ----
    // Delegated to r2_tau2_digest (own register frame; see its header).
    bar_D_rsp = bar_D;
    mu_rsp = mu;
    packed_rsp = bar_w_packed;
    tau2_rsp = tau2_digest;
    tau2_rsp = r2_tau2_digest(bar_D_rsp, mu_rsp, packed_rsp, tau2_rsp);
    tau2_digest = tau2_rsp;

    // ---- 11. Write response buffer ----
    // Reload response_out_ptr lazily here (entry [10] of the pointer
    // table); it is only needed by this final write phase.
    response_out_ptr = ptrs_base;
    response_out_ptr = [:u64 response_out_ptr + 80];
    // [0 .. 32) tau_2
    ?{}, i = #set0_64();
    while (i < 32) {
        c = tau2_digest[i];
        [:u8 response_out_ptr + i] = c;
        i += 1;
    }

    // [32 .. 1632) z_i as u32 stream
    //
    // Per-party z_i is transmitted as raw u32 LE (l = 5 polynomials *
    // 256 coeffs * 4 = 5120 bytes) and padded to 1600 bytes of the
    // 1600-byte slot. This is a slight oversize vs FIPS 204's 20-bit
    // pack_z encoding, but lets us avoid a custom pack_z helper —
    // libjade does not expose a standalone pack_z (the 20-bit packing
    // is inlined inside `pack_signature`). The aggregator at Combine
    // reads u32, sums per-party, then passes the aggregate z through
    // libjade's `pack_signature` to produce the final FIPS 204-byte-
    // identical signature. Note: actual stride is 5120 bytes, the
    // remaining 480 bytes are pad-zero forward-compat.
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) {
        r32 = z_i[i];
        [:u32 response_out_ptr + PULSARM_R2_Z_OFFSET + 4 * i] = r32;
        i += 1;
    }
    // pad the tail of the z slot
    ?{}, i = #set0_64();
    while (i < PULSARM_R2_Z_LEN - Li2_l * Li2_polydeg * 4) {
        [:u8 response_out_ptr + PULSARM_R2_Z_OFFSET + Li2_l * Li2_polydeg * 4 + i] = 0;
        i += 1;
    }

    // [1632 .. 7776) r_i raw u32 stream (l = 5 polynomials * 256 coeffs * 4 = 5120 bytes,
    // padded to k = 6 polynomial slot = 6144 bytes for forward-compat polyveck r_i form).
    ?{}, i = #set0_64();
    while (i < Li2_l * Li2_polydeg) {
        r32 = r_i[i];
        [:u32 response_out_ptr + PULSARM_R2_R_OFFSET + 4 * i] = r32;
        i += 1;
    }
    while (i < Li2_k * Li2_polydeg) {
        [:u32 response_out_ptr + PULSARM_R2_R_OFFSET + 4 * i] = 0;
        i += 1;
    }

    // [7776 .. 8544) pack_w1(my_w1)
    // Direct stack-to-stack assignment (see the bar_w pack_w1 note above).
    vec_k_rsp = my_w1;
    bar_w_packed = pack_w1(vec_k_rsp);
    ?{}, i = #set0_64();
    while (i < Li2_pack_w1len) {
        c = bar_w_packed[i];
        [:u8 response_out_ptr + PULSARM_R2_W1_OFFSET + i] = c;
        i += 1;
    }

    // [8544 .. 9216) per-peer MAC set — zero-init here, filled in by the
    // wrapper via pulsar_round2_mac (one call per peer).
    ?{}, i = #set0_64();
    while (i < PULSARM_R2_MAC_LEN) {
        [:u8 response_out_ptr + PULSARM_R2_MAC_OFFSET + i] = 0;
        i += 1;
    }
    // [9216 .. 9300) reserved zero
    ?{}, i = #set0_64();
    while (i < PULSARM_R2_RESERVED) {
        [:u8 response_out_ptr + PULSARM_R2_MAC_OFFSET + PULSARM_R2_MAC_LEN + i] = 0;
        i += 1;
    }

    // status already holds 0 (MAC verify passed); no need to reset
    } // end "if (status == 0)" body
    // Reload status from its stack slot (it was not modified in the body).
    status = status_round2_s;
    return status;
}

// -----------------------------------------------------------------------------
// pulsar_round2_mac: compute one slot of the Round-2 per-peer MAC set.
// Wrapper calls this once per peer (j != my_index). Binds tau_2 to the
// receiving peer's identity under the shared key K_{i,j}.
// -----------------------------------------------------------------------------
// CT signature: pointer values are public addresses. response_buf
// holds public tau_2 + writes public MAC tag; session_key contents
// are SECRET (K_{i,j}); peer/my index are public.
#[ct = "public * public * public * public -> public"]
export
fn pulsar_round2_mac(
    reg u64 response_buf_ptr,
    reg u64 session_key_ptr,
    reg u32 peer_index,
    reg u32 my_index)
    -> reg u64
{
    stack u8[32] K;
    stack u8[32] tau2;
    stack u8[32] tag;
    stack u64[25] state;
    reg ptr u64[25] state_rsp;
    reg ptr u8[32] K_rsp, tau_rsp, tag_rsp;
    reg u64 i, slot_offset, offset;
    reg u8 c;

    ?{}, i = #set0_64();
    while (i < 32) { c = [:u8 session_key_ptr + i]; K[i] = c; i += 1; }
    ?{}, i = #set0_64();
    while (i < 32) {
        c = [:u8 response_buf_ptr + PULSARM_R2_TAU_OFFSET + i];
        tau2[i] = c;
        i += 1;
    }

    state_rsp = state;
    K_rsp = K;
    state_rsp = kmac256_init_R1MAC(state_rsp, K_rsp);  // KMAC256 init; the R2-MAC
                                                       // customization differs from R1-MAC
                                                       // only by a one-char suffix and is
                                                       // wired in mac.jinc (kmac256_init_R2MAC).

    ?{}, offset = #set0_64();
    tau_rsp = tau2;
    state_rsp, offset = absorb_array32(state_rsp, offset, tau_rsp);
    state_rsp, offset = absorb_u32_be(state_rsp, offset, my_index);

    tag_rsp = tag;
    state_rsp, tag_rsp = kmac256_finalize_32(state_rsp, offset, tag_rsp);
    tag = tag_rsp;

    // Fold the buffer base and the slot offset into one register so the
    // per-byte store is an encodable [base + index] mode (base + slot +
    // index would be three address registers, which x86 cannot encode).
    slot_offset = (64u) peer_index;
    slot_offset *= 32;
    slot_offset += PULSARM_R2_MAC_OFFSET;
    slot_offset += response_buf_ptr;
    ?{}, i = #set0_64();
    while (i < 32) {
        c = tag[i];
        [:u8 slot_offset + i] = c;
        i += 1;
    }

    reg u64 status;
    ?{}, status = #set0_64();
    return status;
}