// mlkem_top.v - ML-KEM KeyGen top-level integration. Runtime-selectable
// parameter set via k_i: k=2 (ML-KEM-512, eta1=3), k=3 (768), k=4 (1024).
// Storage is sized for KMAX (worst case = ML-KEM-1024); k_i picks the
// active sub-range at start_i.
//
// Streaming valid/ready interface. Given seeds d and z, computes the
// ML-KEM key pair per FIPS 203 Algorithm 16 (KeyGen_internal):
//   (rho,sigma) = G(d || K)
//   A_hat[i][j] = SampleNTT(rho || j || i)        i,j in 0..K-1
//   s[i] = CBD3(PRF(sigma, i)),  e[i] = CBD3(PRF(sigma, K+i))
//   s_hat[i] = NTT(s[i]);  e_hat[i] = NTT(e[i])
//   t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j]
//   ek = byteEncode12(t_hat[0..K-1]) || rho
//   dk = byteEncode12(s_hat[0..K-1]) || ek || H(ek) || z
//
// Built incrementally and verified stage-by-stage against ml-kem-r golden
// vectors (test_framework/modules/mlkem_keygen/golden) and NIST KAT.
//
// Uses independent (verified) leaf modules, each with its own keccak_core:
// sha3_top, sample_ntt_sync, sample_cbd_sync, ntt_core, poly_mul_sync,
// mod_add_sync. No shared-keccak arbiter.

`include "sync_rtl/common/defines.vh"

module mlkem_top #(
    parameter KMAX = 4               // storage sizing (worst case = ML-KEM-1024)
) (
    input              clk,
    input              rst_n,
    input  [2:0]       k_i,          // RUNTIME ML-KEM param: 2=512, 3=768, 4=1024
    input  [1:0]       op_i,          // 0=KeyGen, 1=Encaps, 2=Decaps (captured at start_i)
    input  [255:0]     d_i,          // KeyGen seed d (byte 0 in d_i[7:0])
    input  [255:0]     z_i,          // implicit-rejection seed z
    input  [255:0]     msg_i,        // Encaps message m (byte 0 in msg_i[7:0])
    input              start_i,      // pulse to begin KeyGen
    output             busy_o,       // high while running
    output             done_o,       // pulse when ek/dk ready

    // Encaps ek input: stream ek bytes (=pk) into ek_bram before start_i with
    // op_i=1. ek_in_we writes byte ek_in_byte at address ek_in_addr.
    input              ek_in_we,     // write one ek byte (only honored in ST_IDLE for Encaps preload)
    input  [10:0]      ek_in_addr,   // ek byte address 0..ek_bytes-1
    input  [7:0]       ek_in_byte,   // ek byte value
    // Decaps dk input: stream dk bytes (=sk, 768K+96) into the design before
    // start_i with op_i=2. The load logic routes each byte by region:
    //   [0,384K)            -> dk_pke (dkp_bram)   = s_hat encoding
    //   [384K, 768K+32)     -> ek_pke (ek_bram)    = t_hat encoding || rho
    //   [768K+32, 768K+64)  -> H(ek)  (hek_r)
    //   [768K+64, 768K+96)  -> z      (z_r)
    input              dk_in_we,     // write one dk byte (ST_IDLE Decaps preload)
    input  [11:0]      dk_in_addr,   // dk byte address 0..(768K+96)-1
    input  [7:0]       dk_in_byte,   // dk byte value
    // Decaps ciphertext input: stream c bytes (=ct) into c_in_bram before
    // start_i. Kept separate from ct_bram so the computed c' can be compared
    // against the original c (and J(z||c) can read c) without aliasing.
    input              c_in_we,      // write one c byte (ST_IDLE Decaps preload)
    input  [10:0]      c_in_addr,    // c byte address 0..ct_bytes-1
    input  [7:0]       c_in_byte,    // c byte value
    // Encaps shared secret output (= K), valid at done_o
    output [255:0]     ss_o,
    // Encaps ciphertext readback tap: ct byte by index 0..ct_bytes-1
    input  [10:0]      dbg_ct_idx_i,
    output [7:0]       dbg_ct_o,

    // Debug readback tap: read one stored coefficient by (poly slot, index).
    // Lets stage TBs verify intermediates without wide buses.
    input  [5:0]       dbg_slot_i,   // poly slot (see localparams below)
    input  [7:0]       dbg_idx_i,    // coefficient index 0..255
    output [11:0]      dbg_coeff_o,
    // Debug byte readback: ek (sel=0, 0..799) / dk_pke (sel=1, 0..767)
    input              dbg_byte_sel_i,
    input  [10:0]      dbg_byte_idx_i,
    output [7:0]       dbg_byte_o,
    // Debug full-dk readback: dk = dk_pke(768) || ek(800) || H(ek)(32) || z(32)
    // = 1632 bytes. Index 0..1631.
    input  [11:0]      dbg_dk_idx_i,
    output [7:0]       dbg_dk_o,
    // Debug taps for hash outputs
    output [255:0]     dbg_rho_o,
    output [255:0]     dbg_sigma_o,
    // Encaps debug taps: r (G high half) and H(ek)
    output [255:0]     dbg_r_o,
    output [255:0]     dbg_hek_o,
    // Decaps debug taps: m' (Decrypt output), K-bar (J output), parsed z/h.
    output [255:0]     dbg_mprime_o,
    output [255:0]     dbg_kbar_o,
    output [255:0]     dbg_decz_o,   // parsed z (dk[768K+64:+32])
    output [255:0]     dbg_dech_o    // parsed H(ek) (dk[768K+32:+32])
);

    localparam Q = `Q;   // 3329

    // Runtime ML-KEM parameter, captured at start_i.
    reg  [2:0] k_r;
    // FIPS 203: eta1 = 3 for ML-KEM-512 (k=2), else 2 (k=3/4).
    wire [1:0] eta1_rt = (k_r == 3'd2) ? 2'd3 : 2'd2;

    // Runtime-derived sizes (k_r in {2,3,4}). Small multiplies are cheap.
    wire [5:0] kk_rt     = k_r * k_r;            // 4/9/16
    wire [5:0] slot_s_rt = kk_rt;                // s_hat base slot
    wire [5:0] slot_e_rt = kk_rt + k_r;          // e_hat base slot
    wire [5:0] slot_t_rt = kk_rt + {1'b0, k_r} + {1'b0, k_r};  // t_hat base = kk+2k
    wire [11:0] ek_bytes_rt = 12'd384 * {9'b0, k_r} + 12'd32;  // 800/1184/1568
    wire [11:0] dk_bytes_rt = 12'd384 * {9'b0, k_r};            // 768/1152/1536
    // full dk = dk_pke(dk_bytes) || ek(ek_bytes) || H(ek)(32) || z(32). Region
    // boundaries (used by Decaps dk load routing and dk readback).
    wire [11:0] dk_ek_end   = dk_bytes_rt + ek_bytes_rt;        // ek_pke region end
    wire [11:0] dk_hek_end  = dk_ek_end + 12'd32;               // H(ek) region end
    // LOAD-TIME boundaries: dk is streamed in BEFORE start_i, so k_r is not yet
    // captured. Route the dk load by the LIVE k_i input instead of k_r.
    wire [11:0] dkp_bytes_ld = 12'd384 * {9'b0, k_i};                 // 768/1152/1536
    wire [11:0] ek_bytes_ld  = 12'd384 * {9'b0, k_i} + 12'd32;        // 800/1184/1568
    wire [11:0] dk_ek_end_ld = dkp_bytes_ld + ek_bytes_ld;           // ek_pke region end
    wire [11:0] dk_hek_end_ld= dk_ek_end_ld + 12'd32;                // H(ek) region end
    // H(ek) block count = ceil((ek_bytes+1)/136): 6/9/12 for k=2/3/4 (table)
    wire [3:0] h_nblk_rt = (k_r == 3'd2) ? 4'd6 : (k_r == 3'd3) ? 4'd9 : 4'd12;
    wire [11:0] h_last_rt = {6'b0, h_nblk_rt} * 12'd136 - 12'd1;  // final padded byte index

    // ---- Encaps runtime params ----
    reg  [1:0]   op_r;             // 0=KeyGen 1=Encaps 2=Decaps (captured at start)
    reg  [255:0] m_r;              // Encaps message m (captured at start); Decaps reuses for m'
    reg  [255:0] ss_r;             // Encaps shared secret K (= G output low half)
    reg  [255:0] r_r;              // Encaps PRF seed r (= G output high half)
    // ---- Decaps runtime params (parsed from dk during load) ----
    reg  [255:0] z_r;              // implicit-rejection seed z (dk tail), captured at load
    reg  [255:0] kbar_r;           // K-bar = J(z||c) (D5)
    // FIPS 203: eta2 = 2 for all parameter sets.
    wire [1:0]  eta2_rt = 2'd2;
    // Compression params: (du,dv) = (10,4) for k=2/3, (11,5) for k=4.
    wire [4:0]  du_rt = (k_r == 3'd4) ? 5'd11 : 5'd10;
    wire [4:0]  dv_rt = (k_r == 3'd4) ? 5'd5  : 5'd4;
    // Ciphertext byte sizes: c1 = 32*du*k, c2 = 32*dv, ct = c1+c2.
    wire [11:0] c1_bytes_rt = 12'd32 * {7'b0, du_rt} * {9'b0, k_r};  // 640/960/1408
    wire [11:0] c2_bytes_rt = 12'd32 * {7'b0, dv_rt};                 // 128/128/160
    wire [11:0] ct_bytes_rt = c1_bytes_rt + c2_bytes_rt;             // 768/1088/1568
    wire [11:0] ct_bytes_rt = c1_bytes_rt + c2_bytes_rt;             // 768/1088/1568
    // Shared secret out: Encaps drives ss_r = K directly. Decaps implicit reject:
    // ss = (c' == c) ? K' : K-bar. K' is in ss_r, K-bar in kbar_r; dec_reject is
    // the latched c' != c flag from D7 (0 for KeyGen/Encaps, so ss_r passes through).
    assign ss_o = dec_reject ? kbar_r : ss_r;

    // ---- E1: rho-load + byteDecode12 bookkeeping ----
    // rho load: read 32 bytes from ek_bram at offset 384*k into rho_r.
    reg  [5:0]  rl_idx;            // 0..33 read-ahead pointer over rho bytes
    reg  [5:0]  rl_widx;           // 0..31 byte being written (1 cyc behind rl_idx)
    reg         rl_vld;            // a byte is arriving this cycle
    wire [11:0] rl_ek_off = dk_bytes_rt + {6'b0, rl_idx};   // ek addr = 384k + idx
    // byteDecode12: read 3 ek bytes -> 2 coeffs into bank_t. Walk K polys * 128
    // triples. 5-cycle micro-phase per triple: ph0..2 fetch b0,b1,b2 (read-ahead,
    // 1-cycle bram latency), ph3 write c0, ph4 write c1.
    reg  [2:0]  td_poly;           // 0..k-1 t_hat poly index
    reg  [7:0]  td_trip;           // 0..127 triple within poly
    reg  [2:0]  td_ph;             // 0..4 micro-phase
    reg  [7:0]  td_b0, td_b1, td_b2;  // captured ek bytes
    reg         td_done;
    // ek byte address for the triple: poly*384 + trip*3 + (ph for ph<3)
    wire [11:0] td_base = {td_poly, 7'd0} * 12'd3;          // poly*384 (=poly*256*1.5)... see below
    // poly*384 + trip*3 + byteidx; byteidx = ph (0,1,2) during fetch phases
    wire [11:0] td_byteidx = {2'd0, td_trip, 1'b0} + {3'd0, td_trip};  // trip*3
    wire [11:0] td_ekaddr  = ({4'd0,td_poly}*12'd384) + td_byteidx + {9'd0, td_ph};
    // decoded coeffs: c0 = b0 | ((b1&0xF)<<8); c1 = (b1>>4) | (b2<<4)
    wire [11:0] td_c0 = {td_b1[3:0], td_b0};
    wire [11:0] td_c1 = {td_b2, td_b1[7:4]};
    // TDEC byte source: Encaps reads ek (ek_rd_data); Decaps D2 (ST_DEC_SDEC)
    // reuses the same machine to byteDecode12 dk_pke -> s_hat, reading dkp_bram.
    wire        td_dec_s   = (st == ST_DEC_SDEC);
    wire [7:0]  td_byte    = td_dec_s ? dkp_rd_data : ek_rd_data;
    // bank_a write for TDEC (registered); t_hat[td_poly] -> bank_a slot
    // td_poly*K. td_wa is a bank_a byte address (PA_AW=12).
    reg         td_we;
    reg  [11:0] td_wa;             // bank_a address (PA_AW=12)
    reg  [11:0] td_wd;

    // ---- ct_bram: ciphertext byte buffer (<=1568 B). Written by E5/E7
    // (compress + byteEncode_du/dv), read back via dbg_ct tap. ----
    wire [10:0] ct_rd_addr;
    wire [7:0]  ct_rd_data;
    reg         ct_we;
    reg  [10:0] ct_wa;
    reg  [7:0]  ct_wd;
    sd_bram #(.W(8), .D(2048), .A(11)) u_ct_bram (
        .clk(clk),
        .rd_addr(ct_rd_addr), .rd_data(ct_rd_data),
        .wr_en(ct_we), .wr_addr(ct_wa), .wr_data(ct_wd)
    );
    // ct_bram read addr: D7 compare walks c' via cmp_idx; else debug readback.
    assign ct_rd_addr = (st == ST_DEC_CMP) ? cmp_idx : dbg_ct_idx_i;
    assign dbg_ct_o   = ct_rd_data;

    // ---- c_in_bram: Decaps input ciphertext c (<=1568 B). Preloaded via
    // c_in_* in ST_IDLE. Kept separate from ct_bram (which holds the computed
    // c') so D7 can compare c' vs c and D5 J(z||c) can read c. Read port is
    // muxed (D5 J-feed / D7 compare); for D0 only the load write path exists.
    wire [10:0] cin_rd_addr;
    wire [7:0]  cin_rd_data;
    reg  [10:0] cin_rd_addr_r;       // D5/D7 read address (tied 0 until then)
    sd_bram #(.W(8), .D(2048), .A(11)) u_c_in_bram (
        .clk(clk),
        .rd_addr(cin_rd_addr), .rd_data(cin_rd_data),
        .wr_en(c_in_we), .wr_addr(c_in_addr), .wr_data(c_in_byte)
    );
    // c_in read addr: D5 J presents dj_c_idx combinationally (no extra register
    // stage) so the registered BRAM read yields the byte 1 cycle later, matching
    // the assemble/writeback pipeline. Otherwise the registered cin_rd_addr_r
    // (D1 walker / idle) drives it.
    assign cin_rd_addr = (st == ST_DEC_J)   ? dj_c_idx[10:0] :
                         (st == ST_DEC_CMP) ? cmp_idx        : cin_rd_addr_r;

    // ================================================================
    // Polynomial storage, sized for KMAX (worst case). Runtime k uses a
    // sub-range. Slot layout (each slot = 256 coeffs):
    //   A_hat[i][j] : slots 0 .. k*k-1        at index i*k + j
    //   s_hat[i]    : slots slot_s_rt .. +k-1
    //   e_hat[i]    : slots slot_e_rt .. +k-1
    //   t_hat[i]    : slots slot_t_rt .. +k-1
    // NUM_SLOTS = KMAX*KMAX + 3*KMAX = 28 for KMAX=4.
    // ================================================================
    localparam NUM_SLOTS = KMAX*KMAX + 3*KMAX;
    localparam SAW       = 5;             // slot-address width (>=clog2(28))

    // ================================================================
    // Coefficient storage as 3 banks (Phase 2). Was one async-read reg
    // array polymem[0:28*256-1]; split by polynomial so each becomes a
    // registered-read sd_bram (infers BRAM / ASIC SRAM) in the final step.
    //   bank_a : A_hat[i][j]  slots 0..K^2-1   -> D=KMAX*KMAX*256=4096
    //   bank_se: s_hat||e_hat slots slot_s..   -> D=2*KMAX*256=2048 (rel 0..2K-1)
    //   bank_t : t_hat[i]     slots slot_t..   -> D=KMAX*256=1024  (rel 0..K-1)
    // Addresses are base-relative: bank index = abs_slot - base_slot.
    // Stage 2b (in progress): per-consumer registered read-ahead. Banks are
    // still async reg arrays here; converted consumers read via their own
    // 1-cycle pipeline reg (== sd_bram timing). Once every consumer of a
    // bank is registered, the array is replaced by an sd_bram with its read
    // port muxed across the (phase-disjoint) consumers.
    // ================================================================
    localparam PA_AW = 12;                // bank_a addr width (4096)
    localparam PSE_AW = 11;               // bank_se addr width (2048)
    localparam PT_AW = 10;                // bank_t addr width (1024)

    // bank_a: A_hat. Reader = ST_M load (pm_a) + dbg. Writer = ST_A.
    // Read port muxed by phase; sd_bram's internal rd_addr_r register is the
    // 1-cycle read pipeline (replaces the explicit pm_a_rd reg). Write port is
    // driven combinationally (sd_bram registers the write at posedge, matching
    // the old reg-array nonblocking write timing).
    wire [PA_AW-1:0] ba_rd_addr;
    wire [11:0]      ba_rd_data;
    wire             ba_we;
    wire [PA_AW-1:0] ba_wa;
    wire [11:0]      ba_wd;
    sd_bram #(.W(12), .D(1<<PA_AW), .A(PA_AW)) u_bank_a (
        .clk(clk), .rd_addr(ba_rd_addr), .rd_data(ba_rd_data),
        .wr_en(ba_we), .wr_addr(ba_wa), .wr_data(ba_wd)
    );
    // ST_A / ST_ENC_A write: commit snt_coeff to A_hat[a_slot] when accepted.
    // Encaps E6 also writes t_hat here via byteDecode12 (TDEC): t_hat[j] lands
    // in bank_a at slot j*K so V's MAC reuses E4's u_aslot=u_j*K (u_row=0) addr.
    // ST_ENC_E2MV writes the relocated e2 into bank_a[E2_ASLOT] (em_we).
    // D1 DECOMP writes v' into bank_a[DEC_VASLOT] (dec_v_we) so D3 can read it
    // in parallel with the bank_t psum during the w-SUB.
    assign ba_we = ((st == ST_A || st == ST_ENC_A) && a_busy && snt_vo && snt_ack)
                   || td_we || em_we || dec_v_we;
    assign ba_wa = td_we ? td_wa[PA_AW-1:0]
                 : em_we ? em_wa[PA_AW-1:0]
                 : dec_v_we ? dec_v_wr[PA_AW-1:0]
                         : (a_slot*256 + a_widx) & ((1<<PA_AW)-1);
    assign ba_wd = td_we ? td_wd :
                   em_we ? bt_rd_data :     // e2 read from bank_t rel slot 0
                   dec_v_we ? cd_out :      // v' decompressed coeff
                           snt_coeff;

    // bank_se: s_hat || e_hat. Readers = ST_N load, ST_M load (pm_b, s_hat[j]),
    // ST_M acc (e_hat, j==0), ST_E dk-half, dbg. Writers = ST_C (CBD), ST_N
    // writeback (NTT in place). All readers already registered; read port muxed
    // by phase (within ST_M, load vs accumulate never overlap). Write port
    // combinational (ST_C and ST_N are disjoint states).
    wire [PSE_AW-1:0] bse_rd_addr;
    wire [11:0]       bse_rd_data;
    wire              bse_we;
    wire [PSE_AW-1:0] bse_wa;
    wire [11:0]       bse_wd;
    sd_bram #(.W(12), .D(1<<PSE_AW), .A(PSE_AW)) u_bank_se (
        .clk(clk), .rd_addr(bse_rd_addr), .rd_data(bse_rd_data),
        .wr_en(bse_we), .wr_addr(bse_wa), .wr_data(bse_wd)
    );

    // bank_t: t_hat. Readers = ST_M acc (j>0 running t_hat) + ST_E ek-half + dbg.
    // Writer = ST_M acc. Read port muxed by phase; sd_bram's rd_addr_r is the
    // 1-cycle read register. RMW is safe: the acc read addr leads the write
    // addr by 1 (m_oidx+1 vs m_oidx), so read/write never alias in a cycle and
    // sd_bram write-first == reg-array read-old. Write port combinational.
    wire [PT_AW-1:0] bt_rd_addr;
    wire [11:0]      bt_rd_data;
    wire             bt_we;
    wire [PT_AW-1:0] bt_wa;
    wire [11:0]      bt_wd;
    sd_bram #(.W(12), .D(1<<PT_AW), .A(PT_AW)) u_bank_t (
        .clk(clk), .rd_addr(bt_rd_addr), .rd_data(bt_rd_data),
        .wr_en(bt_we), .wr_addr(bt_wa), .wr_data(bt_wd)
    );
    // bank_t write: KeyGen ST_M accumulate, or Encaps e2 (ST_ENC_C c_poly==2K
    // -> bank_t rel slot 0, free during C/N/U). (TDEC writes bank_a now.)
    wire e2_we = (st == ST_ENC_C) && c_busy && cbd_vo && cbd_ack &&
                 (c_poly == {1'b0, k_r, 1'b0});   // c_poly == 2K
    // Encaps E4/E6 bank_t writes: MAC psum (u_sub0, on pm_vo) and INTT in-place
    // (u_sub1, on ntt_vo). Both target bank_t rel slot UPSUM. ST_ENC_V (E6) adds
    // the V-ADD writeback (psum + e2 + mu) at UPSUM via u_v_we below.
    wire u_psum_we = (st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && (u_sub == 2'd0) && pm_vo;
    wire u_intt_we = (st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && (u_sub == 2'd1) && ntt_vo;
    // D1 DECOMP writeback: cd_vo while decoding c2 writes v' to bank_a[DEC_VASLOT].
    wire dec_v_we  = (st == ST_DEC_DECOMP) && dec_in_c2 && cd_vo;
    wire [13:0] dec_v_wr = DEC_VASLOT*256 + dec_widx;
    assign bt_we = ((st == ST_M) && pm_vo) || e2_we ||
                   u_psum_we || u_intt_we || u_v_we || u_w_we;
    assign bt_wa = e2_we ? (c_widx & ((1<<PT_AW)-1)) :
                   u_psum_we ? ((UPSUM*256 + u_oidx) & ((1<<PT_AW)-1)) :
                   u_intt_we ? (u_intt_wr[PT_AW-1:0]) :
                   u_v_we    ? (u_vadd_wr[PT_AW-1:0]) :
                   u_w_we    ? (u_wsub_wr[PT_AW-1:0]) :
                           ((m_i*256 + m_oidx) & ((1<<PT_AW)-1));
    assign bt_wd = e2_we ? cbd_modq :
                   u_psum_we ? u_accq :
                   u_intt_we ? ntt_coeff :
                   u_v_we    ? u_vq :
                   u_w_we    ? u_wq :
                           m_accq;

    // Debug readback (registered for timing)
    reg [11:0] dbg_coeff_r;
    // Route debug read by absolute slot range (not on the KAT correctness
    // path — TB reads ek/dk BRAMs only). A: 0..K^2-1; s/e: slot_s..slot_t-1;
    // t: slot_t..  Index within bank is (abs_slot-base)*256 + idx.
    wire [13:0] dbg_a_addr  = dbg_slot_i*256 + dbg_idx_i;
    wire [13:0] dbg_se_addr = (dbg_slot_i - slot_s_rt)*256 + dbg_idx_i;
    wire [13:0] dbg_t_addr  = (dbg_slot_i - slot_t_rt)*256 + dbg_idx_i;
    // bank_a read port: ST_M / Encaps ST_ENC_U MAC drive pm_a addr; ST_ENC_V
    // MAC reads t_hat[j] (u_pm_a_full, u_row=0) and ADD reads e2 (E2_ASLOT); else dbg.
    assign ba_rd_addr = (st == ST_M)     ? pm_a_full[PA_AW-1:0]    :
                        (st == ST_ENC_U) ? u_pm_a_full[PA_AW-1:0]  :
                        (st == ST_ENC_V) ?
                            ((u_sub == 2'd2) ? u_v_e2rd[PA_AW-1:0]    // ADD: e2 read
                                             : u_pm_a_full[PA_AW-1:0]) : // MAC: t_hat[u_j]
                        (st == ST_DEC_W) ?
                            ((u_sub == 2'd2) ? u_w_vrd[PA_AW-1:0]     // SUB: v' read
                                             : u_pm_a_full[PA_AW-1:0]) : // MAC: s_hat[u_j]
                                           dbg_a_addr[PA_AW-1:0];
    // bank_t read port: ST_M acc (t_hat), ST_E ek-half (byteEncode), Encaps
    // ST_ENC_U/V + Decaps ST_DEC_W (MAC psum acc, INTT load, or ADD/SUB psum
    // read by sub-phase), ST_ENC_E2MV (e2 read from rel slot 0); else dbg.
    assign bt_rd_addr = (st == ST_M) ? m_tacc_full[PT_AW-1:0] :
                        (st == ST_E) ? e_rd_full[PT_AW-1:0]   :
                        (st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) ?
                            ((u_sub == 2'd0) ? u_psum_full[PT_AW-1:0] :  // MAC acc
                             (u_sub == 2'd1) ? u_intt_rd[PT_AW-1:0]   :  // INTT load
                                               u_add_prd[PT_AW-1:0]) :   // ADD/SUB psum read
                        (st == ST_ENC_E2MV) ? em_rd[PT_AW-1:0] :         // e2 relocate read
                        (st == ST_ENC_C2) ? cp_bt_full[PT_AW-1:0] :      // C2: v[cp_idx]
                        (st == ST_DEC_MENC) ? men_rd[PT_AW-1:0] :        // D4: w[men_idx]
                                       dbg_t_addr[PT_AW-1:0];
    // bank_se read port: ST_N load (s/e NTT), ST_M load (pm_b s_hat[j]) vs
    // accumulate (e_hat, mutually exclusive via m_loading), ST_E dk-half, dbg.
    assign bse_rd_addr = (st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? ntt_rd_full[PSE_AW-1:0] :
                         (st == ST_M) ? (m_loading ? pm_b_full[PSE_AW-1:0]
                                                   : m_eacc_full[PSE_AW-1:0]) :
                         (st == ST_E) ? e_rd_full[PSE_AW-1:0] :
                         (st == ST_ENC_U) ?
                             ((u_sub == 2'd0) ? u_pm_b_full[PSE_AW-1:0] :  // MAC: y_hat[u_j]
                                                u_add_e1rd[PSE_AW-1:0]) :  // ADD: e1[u_row] read
                         (st == ST_ENC_V) ? u_pm_b_full[PSE_AW-1:0] :     // V MAC: y_hat[u_j]
                         (st == ST_DEC_W && u_sub == 2'd0) ? u_pm_b_full[PSE_AW-1:0] : // D3 MAC: u_hat[u_j]
                         (st == ST_ENC_C1) ? cp_se_full[PSE_AW-1:0] :     // C1: u[cp_poly]
                                        dbg_se_addr[PSE_AW-1:0];
    // bank_se write: KeyGen ST_C CBD (s/e), Encaps ST_ENC_C CBD (y/e1, c_poly<2K;
    // e2 at c_poly==2K goes to bank_t instead), ST_N/ST_ENC_N NTT writeback,
    // Encaps E4 ADD u-writeback (u_sub==2).
    // Encaps E4 ADD: write u[u_row] over e1 in bank_se rel slot (K+u_row).
    wire u_add_we = (st == ST_ENC_U) && (u_sub == 2'd2) && u_avalid;
    // D1 DECOMP writeback: cd_vo while decoding c1 writes u'[dec_poly] to
    // bank_se rel slot dec_poly (0..K-1), at coeff index dec_widx.
    wire dec_u_we  = (st == ST_DEC_DECOMP) && !dec_in_c2 && cd_vo;
    wire [13:0] dec_u_wr = {2'b0, dec_poly}*256 + dec_widx;
    assign bse_we = ((st == ST_C) && c_busy && cbd_vo && cbd_ack) ||
                    ((st == ST_ENC_C) && c_busy && cbd_vo && cbd_ack &&
                     (c_poly < {1'b0, k_r, 1'b0})) ||
                    (((st == ST_N) || (st == ST_ENC_N) || (st == ST_DEC_NTT)) && ntt_vo) ||
                    u_add_we || dec_u_we;
    assign bse_wa = u_add_we ? u_add_uwr[PSE_AW-1:0] :
                    dec_u_we  ? dec_u_wr[PSE_AW-1:0] :
                    (st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
                                 : ((c_poly*256 + c_widx) & ((1<<PSE_AW)-1));
    assign bse_wd = u_add_we ? u_uq :
                    dec_u_we  ? cd_out :
                    (st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? ntt_coeff : cbd_modq;
    always @(posedge clk) begin
        if (dbg_slot_i >= slot_t_rt)      dbg_coeff_r <= bt_rd_data;     // bank_t (sd_bram)
        else if (dbg_slot_i >= slot_s_rt) dbg_coeff_r <= bse_rd_data;    // bank_se (sd_bram)
        else                              dbg_coeff_r <= ba_rd_data;     // bank_a (sd_bram)
    end
    assign dbg_coeff_o = dbg_coeff_r;

    // ek and dk_pke byte memories sized for KMAX.
    localparam EK_MAX = 384*KMAX + 32;   // 1568
    localparam DK_MAX = 384*KMAX;         // 1536

    // ek / dk_pke byte storage in BRAM (sd_bram: 1R/1W, 1-cycle read latency).
    // Each single read port is shared time-disjointly: ST_H assemble reads ek
    // (in ST_H), debug readback reads ek/dk (in ST_DONE). Write port driven by
    // ST_E (1 byte/cycle) and the rho-tail copy.
    wire [10:0] ek_rd_addr, dkp_rd_addr;
    wire [7:0]  ek_rd_data, dkp_rd_data;
    reg         ek_we, dkp_we;
    reg  [10:0] ek_wa, dkp_wa;
    reg  [7:0]  ek_wd, dkp_wd;

    // ek BRAM write port: KeyGen ST_E drives ek_we/ek_wa/ek_wd; Encaps preloads
    // ek from the external ek_in_* port (TB streams ek=pk before start_i). Decaps
    // routes the dk's ek_pke region (dk[384K:768K+32]) here at load time. The
    // three never overlap (preloads happen in ST_IDLE before a run).
    wire        dk_ld_ekpke = dk_in_we && (dk_in_addr >= dkp_bytes_ld)
                                       && (dk_in_addr <  dk_ek_end_ld);  // ek_pke region
    wire        dk_ld_dkpke = dk_in_we && (dk_in_addr <  dkp_bytes_ld);  // dk_pke region
    wire [10:0] dk_ek_off   = (dk_in_addr - dkp_bytes_ld);               // offset within ek
    wire        ek_we_mux = ek_in_we   ? 1'b1       :
                            dk_ld_ekpke ? 1'b1       : ek_we;
    wire [10:0] ek_wa_mux = ek_in_we   ? ek_in_addr :
                            dk_ld_ekpke ? dk_ek_off  : ek_wa;
    wire [7:0]  ek_wd_mux = ek_in_we   ? ek_in_byte :
                            dk_ld_ekpke ? dk_in_byte : ek_wd;

    sd_bram #(.W(8), .D(2048), .A(11)) u_ek_bram (
        .clk(clk),
        .rd_addr(ek_rd_addr), .rd_data(ek_rd_data),
        .wr_en(ek_we_mux), .wr_addr(ek_wa_mux), .wr_data(ek_wd_mux)
    );
    // dk_pke BRAM write port: KeyGen ST_E (s_hat byteEncode) drives dkp_we/wa/wd;
    // Decaps routes the dk's dk_pke region (dk[0:384K]) here at load time. The
    // two never overlap (preload in ST_IDLE before a Decaps run).
    wire        dkp_we_mux = dk_ld_dkpke ? 1'b1                 : dkp_we;
    wire [10:0] dkp_wa_mux = dk_ld_dkpke ? dk_in_addr[10:0]     : dkp_wa;
    wire [7:0]  dkp_wd_mux = dk_ld_dkpke ? dk_in_byte           : dkp_wd;
    sd_bram #(.W(8), .D(2048), .A(11)) u_dkp_bram (
        .clk(clk),
        .rd_addr(dkp_rd_addr), .rd_data(dkp_rd_data),
        .wr_en(dkp_we_mux), .wr_addr(dkp_wa_mux), .wr_data(dkp_wd_mux)
    );

    // full dk = dk_pke(dk_bytes) || ek(ek_bytes) || H(ek)(32) || z(32).
    // (dk_ek_end / dk_hek_end declared near the top with the other size wires.)

    // Debug-region selects for dk readback (combinational region decode).
    wire dbgdk_in_dkp = (dbg_dk_idx_i < dk_bytes_rt);
    wire dbgdk_in_ek  = (dbg_dk_idx_i >= dk_bytes_rt) && (dbg_dk_idx_i < dk_ek_end);

    // Read-address muxes (declared here; depend on ST_H h_g defined below via
    // the function-free expression h_blk*136 + h_byte). See assigns after the
    // H-stage register declarations.
    // dbg_byte_o / dbg_dk_o are combinational taps over the (registered) BRAM
    // read data; net 1-cycle latency, well within the TB's 2-cycle read wait.
    assign dbg_byte_o = dbg_byte_sel_i ? dkp_rd_data : ek_rd_data;
    assign dbg_dk_o   = dbgdk_in_dkp ? dkp_rd_data :
                        dbgdk_in_ek  ? ek_rd_data  :
                        (dbg_dk_idx_i < dk_hek_end) ? hek_r[(dbg_dk_idx_i - dk_ek_end)*8 +: 8]
                                                    : z_i[(dbg_dk_idx_i - dk_hek_end)*8 +: 8];

    // ================================================================
    // Top-level FSM (built incrementally). Stage 2a: G only.
    // ================================================================
    localparam ST_IDLE   = 5'd0;
    localparam ST_G      = 5'd1;   // run G(d||K), capture rho/sigma
    localparam ST_A      = 5'd2;   // generate A_hat[i][j] via SampleNTT
    localparam ST_C      = 5'd3;   // generate s[i],e[i] via CBD
    localparam ST_N      = 5'd4;   // forward NTT of s[i],e[i] in place
    localparam ST_M      = 5'd5;   // matrix accumulate t_hat = e_hat + sum A o s_hat
    localparam ST_E      = 5'd6;   // byteEncode12 -> ek/dk BRAM
    localparam ST_H      = 5'd7;   // H(ek) via multi-block SHA3-256
    // ---- Encaps states ----
    localparam ST_ENC_LOAD = 5'd8;  // ek already preloaded; placeholder/settle
    localparam ST_ENC_H    = 5'd9;  // H(ek) via multi-block SHA3-256 (reuses ST_H logic)
    localparam ST_ENC_G    = 5'd10; // (K,r) = G(m || H(ek)), 64-byte single block
    localparam ST_ENC_A    = 5'd11; // regenerate A_hat (transpose used in U)
    localparam ST_ENC_TDEC = 5'd12; // byteDecode12: ek -> t_hat (bank_a slot j*K)
    localparam ST_ENC_C    = 5'd13; // sample y,e1,e2 via CBD (eta1/eta2)
    localparam ST_ENC_N    = 5'd14; // forward NTT of y in place
    localparam ST_ENC_U    = 5'd15; // u = INTT(sum A^T o y_hat) + e1
    localparam ST_ENC_C1   = 5'd16; // Compress_du + byteEncode_du -> ct c1
    localparam ST_ENC_V    = 5'd17; // v = INTT(sum t_hat o y_hat) + e2 + mu
    localparam ST_ENC_C2   = 5'd18; // Compress_dv + byteEncode_dv -> ct c2
    localparam ST_ENC_E2MV = 5'd19; // relocate e2 bank_t[0] -> bank_a[E2_ASLOT]
    // ---- Decaps states ----
    localparam ST_DEC_LOAD   = 5'd20; // dk/c already streamed in; parse/settle (D0)
    localparam ST_DEC_DECOMP = 5'd21; // D1: byteDecode_d + Decompress c1->u', c2->v'
    localparam ST_DEC_SDEC   = 5'd22; // D2: byteDecode12 dk_pke -> s_hat (bank_a slot j*K)
    localparam ST_DEC_NTT    = 5'd23; // D2: u_hat[i] = NTT(u'[i]) in place (bank_se rel i)
    localparam ST_DEC_W      = 5'd24; // D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j])
    localparam ST_DEC_MENC   = 5'd25; // D4: m' = byteEncode_1(Compress_1(w))
    localparam ST_DEC_G      = 5'd26; // D5: (K',r') = G(m' || h), SHA3-512 single block
    localparam ST_DEC_J      = 5'd27; // D5: K-bar = J(z || c), multi-block SHAKE256
    localparam ST_DEC_CMP    = 5'd28; // D7: c' == c compare + implicit-reject ss select
    localparam ST_DONE   = 5'd31;

    reg [4:0]   st, st_next;
    reg [255:0] rho_r, sigma_r;

    // A-generation bookkeeping: explicit i/j counters (avoid runtime divide)
    reg  [2:0]   a_i;          // row 0..k-1
    reg  [2:0]   a_j;          // col 0..k-1
    reg  [4:0]   a_pair;       // 0..k*k pairs done (for done test)
    reg  [7:0]   a_widx;       // write index 0..255 within current poly
    reg          a_busy;       // 1 once current pair's request accepted (gates collect)
    wire [SAW-1:0] a_slot = a_i*k_r + a_j;   // A_hat[i][j] slot = i*k + j

    // C-generation bookkeeping: 2*k polys (s[0..k-1] then e[0..k-1])
    reg  [4:0]   c_poly;       // 0..2k
    reg  [7:0]   c_widx;
    reg          c_busy;
    wire [7:0]   c_nonce = {3'b0, c_poly};       // s:0..k-1  e:k..2k-1  == nonce
    // slot: c_poly < k -> s_hat[c_poly], else e_hat[c_poly-k]
    wire [SAW-1:0] c_slot = (c_poly < {2'b0, k_r}) ? (slot_s_rt + c_poly)
                                                   : (slot_e_rt + (c_poly - {2'b0, k_r}));

    assign busy_o = (st != ST_IDLE);
    assign done_o = (st == ST_DONE);
    assign dbg_rho_o   = rho_r;
    assign dbg_sigma_o = sigma_r;
    assign dbg_r_o     = r_r;
    assign dbg_hek_o   = hek_r;
    // Decaps taps: m' from D4 (written into m_r, reused by V/mu in D6 re-encrypt);
    // z/h parsed from dk at load.
    assign dbg_mprime_o = m_r;
    assign dbg_kbar_o   = kbar_r;
    assign dbg_decz_o   = z_r;
    assign dbg_dech_o   = hek_r;   // Decaps parses dk's H(ek) into hek_r at load

    // ---- sha3_top in G mode: data_i = {K_byte, d} (d byte0 in [7:0]) ----
    reg          sha3_valid;
    wire         sha3_ready;
    wire [511:0] sha3_hash;
    wire         sha3_vo;
    reg          sha3_ack;       // consumer ready for hash
    wire [511:0] kg_g_data  = {248'b0, 5'b0, k_r, d_i};   // KeyGen G: [263:256]=k, [255:0]=d
    wire [511:0] enc_g_data = {hek_r, m_r};               // Encaps G: m || H(ek), 64 bytes
    wire [511:0] dec_g_data = {hek_r, m_r};               // Decaps D5 G: m' || h, 64 bytes
    wire [511:0] g_data = (st == ST_ENC_G) ? enc_g_data :
                          (st == ST_DEC_G) ? dec_g_data : kg_g_data;

    // ================================================================
    // Shared keccak_core + phase mux (3 consumers -> 1 core)
    //
    // G/H (u_sha3), SampleNTT (u_snt) and CBD (u_cbd) each need a Keccak
    // permutation, but they run in DISJOINT top-FSM phases (ST_G/ST_H, ST_A,
    // ST_C respectively), so at most one is ever active. One keccak_core
    // serves all three: kc_state_i/kc_valid_i are muxed by phase; kc_valid_o
    // is broadcast but GATED per consumer so an inactive consumer never
    // latches a permutation result meant for another (the samplers latch
    // squeeze state unconditionally on their kc_valid_o input).
    // ================================================================
    wire [1599:0] kc_state_o;        // shared core output (to all consumers)
    wire          kc_valid_o;        // shared core output-valid
    wire          kc_ready_o;        // shared core input-ready

    // per-consumer drives toward the core
    wire [1599:0] sha3_kc_state_i, snt_kc_state_i, cbd_kc_state_i;
    wire          sha3_kc_valid_i, snt_kc_valid_i, cbd_kc_valid_i;
    /* verilator lint_off UNUSEDSIGNAL */
    wire          sha3_kc_ready_i, snt_kc_ready_i, cbd_kc_ready_i; // all 1'b1
    /* verilator lint_on UNUSEDSIGNAL */

    // phase selects (mutually exclusive). Encaps adds ST_ENC_H/G (sha3),
    // ST_ENC_A (snt), ST_ENC_C (cbd).
    wire sel_sha3 = (st == ST_G) || (st == ST_H) ||
                    (st == ST_ENC_H) || (st == ST_ENC_G) ||
                    (st == ST_DEC_G) || (st == ST_DEC_J);
    wire sel_snt  = (st == ST_A) || (st == ST_ENC_A);
    wire sel_cbd  = (st == ST_C) || (st == ST_ENC_C);

    // gated output-valid: only the active consumer sees kc_valid_o
    wire kc_valid_o_sha3 = kc_valid_o & sel_sha3;
    wire kc_valid_o_snt  = kc_valid_o & sel_snt;
    wire kc_valid_o_cbd  = kc_valid_o & sel_cbd;

    // input mux: route the active consumer's request to the core
    wire [1599:0] kc_state_i_mux = sel_snt ? snt_kc_state_i :
                                   sel_cbd ? cbd_kc_state_i :
                                             sha3_kc_state_i;
    wire kc_valid_i_mux = sel_snt ? snt_kc_valid_i :
                          sel_cbd ? cbd_kc_valid_i :
                          sel_sha3 ? sha3_kc_valid_i : 1'b0;

    keccak_core #(.ROUNDS(24)) u_keccak (
        .clk(clk), .rst_n(rst_n),
        .state_i(kc_state_i_mux),
        .valid_i(kc_valid_i_mux),
        .ready_o(kc_ready_o),
        .state_o(kc_state_o),
        .valid_o(kc_valid_o),
        .ready_i(1'b1)               // consumers always accept (kc_ready_i=1)
    );

    // G/H/J share one sha3_top. Single-block (mb_en=0): KeyGen G (ST_G, mode 00),
    // Encaps G (ST_ENC_G, mode 11), Decaps D5 G (ST_DEC_G, mode 11). Multi-block
    // absorb (mb_en=1): H(ek) (ST_H/ST_ENC_H, SHA3-256 pad) and Decaps D5 J
    // (ST_DEC_J, SHAKE256 pad over z||c). mb_* inputs are muxed by phase.
    wire sha3_mb_en = (st == ST_H) || (st == ST_ENC_H) || (st == ST_DEC_J);
    wire [1:0] sha3_mode = (st == ST_ENC_G || st == ST_DEC_G) ? 2'b11 : 2'b00;
    // multi-block feed mux: J (ST_DEC_J) drives dj_*, else H drives h_*.
    wire [1087:0] mb_block_mux = (st == ST_DEC_J) ? dj_block_r : h_block_r;
    wire          mb_valid_mux = (st == ST_DEC_J) ? dj_mbvalid : h_mbvalid;
    wire          mb_last_mux  = (st == ST_DEC_J) ? dj_mblast  : h_mblast;
    wire          mb_ack_mux   = (st == ST_DEC_J) ? dj_ack     : h_ack;
    sha3_top_shared u_sha3 (
        .clk(clk), .rst_n(rst_n),
        .mode(sha3_mode),             // G = SHA3-512 (only used when mb_en=0)
        .data_i(g_data),
        .valid_i(sha3_valid),
        .ready_o(sha3_ready),
        .hash_o(sha3_hash),
        .valid_o(sha3_vo),
        .ready_i(sha3_mb_en ? mb_ack_mux : sha3_ack),
        .mb_en(sha3_mb_en),
        .mb_block_i(mb_block_mux),
        .mb_valid_i(mb_valid_mux),
        .mb_last_i(mb_last_mux),
        .mb_ready_o(h_mbready),
        // shared keccak_core interface (gated by phase in the mux below)
        .kc_state_o(kc_state_o),
        .kc_valid_o(kc_valid_o_sha3),
        .kc_ready_o(kc_ready_o),
        .kc_state_i(sha3_kc_state_i),
        .kc_valid_i(sha3_kc_valid_i),
        .kc_ready_i(sha3_kc_ready_i)
    );

    // ---- multi-block H(ek) state (SHA3-256, 6/9/12 blocks); fed to shared u_sha3 ----
    reg  [1087:0] h_block_r;       // current pre-padded rate block
    reg           h_mbvalid;
    reg           h_mblast;
    wire          h_mbready;
    wire [511:0]  h_hash;
    wire          h_vo;
    reg           h_ack;
    reg  [255:0]  hek_r;           // captured H(ek)
    reg  [3:0]    h_blk;           // 0..H_NBLK-1 block index (up to 11 for K=4)
    reg  [7:0]    h_byte;          // 0..136 byte address being presented (assemble)
    reg  [1:0]    h_phase;         // 0=assemble 1=feed 2=wait-perm 3=done
    // ek BRAM read is registered (1-cycle latency); assemble presents the
    // address for h_byte this cycle and writes the byte that arrived (for the
    // address presented last cycle) into h_block_r. Writeback pipeline regs:
    reg           h_wb_vld;        // a byte is arriving this cycle
    reg  [7:0]    h_wb_idx;        // its position within the 136-byte block
    reg  [11:0]   h_wb_g;          // its global ek byte index
    reg  [7:0]    h_wb_pad;        // pad constant to use if g is out of ek range
    reg           h_wb_inek;       // 1 if g is within ek range (use BRAM data)

    // ---- Decaps D5 J(z||c) multi-block SHAKE256 state ----
    // Message = z (32 B, from z_r) || c (ct_bytes_rt B, from c_in_bram). Rate=136.
    // SHAKE256 pad: byte at msg_len = 0x1F, last byte of last block |= 0x80.
    // Reuses the shared u_sha3 mb_* port (muxed by ST_DEC_J above).
    reg  [1087:0] dj_block_r;      // current pre-padded rate block
    reg           dj_mbvalid;
    reg           dj_mblast;
    reg           dj_ack;
    reg  [3:0]    dj_blk;          // block index 0..DJ_NBLK-1 (z||c up to ~1600B -> <=12 blks)
    reg  [7:0]    dj_byte;         // 0..136 byte being assembled
    reg  [1:0]    dj_phase;        // 0=assemble 1=feed 2=wait-perm 3=done
    reg           dj_done;
    // writeback pipeline (c_in_bram registered read, 1-cyc latency)
    reg           dj_wb_vld;
    reg  [7:0]    dj_wb_idx;       // position within 136-byte block
    reg           dj_wb_inc;       // 1 if global byte g is within c range (use BRAM)
    reg           dj_wb_inz;       // 1 if g < 32 (use z_r byte)
    reg  [7:0]    dj_wb_pad;       // pad constant if g out of msg range
    reg  [4:0]    dj_wb_zidx;      // z byte index (g) when inz
    // J message length = 32 (z) + ct_bytes_rt; total padded blocks.
    wire [11:0] dj_msglen = 12'd32 + {1'b0, ct_bytes_rt};
    wire [11:0] dj_last_byte = (dj_msglen / 12'd136) * 12'd136 + 12'd135; // last byte of final block
    wire [3:0]  dj_nblk   = (dj_msglen / 12'd136) + 4'd1;
    // global byte index for the address presented this cycle (assemble)
    wire [11:0] dj_g = {4'd0, dj_blk} * 12'd136 + {4'd0, dj_byte};
    // c_in_bram byte index (when g>=32): g-32
    wire [11:0] dj_c_idx = dj_g - 12'd32;
    // SHAKE256 pad constant for global byte g
    function [7:0] dj_padconst(input [11:0] g);
        begin
            if (g == dj_last_byte && g == dj_msglen) dj_padconst = 8'h9f; // 0x1F|0x80
            else if (g == dj_msglen)   dj_padconst = 8'h1f;
            else if (g == dj_last_byte) dj_padconst = 8'h80;
            else                        dj_padconst = 8'h00;
        end
    endfunction

    // h_hash / h_vo are now served by the shared u_sha3 above (mb_en=1 during
    // ST_H). They alias the single core's outputs; the H consumer logic below
    // already gates on st==ST_H, and u_sha3's valid_o/hash_o are mb-selected.
    assign h_hash = sha3_hash;
    assign h_vo   = sha3_vo;

    // SHA3-256 over ek (ek_bytes_rt bytes): rate=136. Padded length = h_nblk_rt*136.
    //   pad: byte ek_bytes_rt = 0x06 (domain + first pad bit), last byte |= 0x80.
    // byte b (0..135) of block blk: global g = blk*136 + b.
    // ek bytes come from the ek BRAM (registered read); only the pad/zero bytes
    // are constants, returned by h_padconst. Reads runtime ek_bytes_rt/h_last_rt.
    function [7:0] h_padconst(input [3:0] blk, input [7:0] b);
        integer g;
        begin
            g = blk*136 + b;
            if (g == h_last_rt && g == ek_bytes_rt) h_padconst = 8'h86; // 0x06|0x80
            else if (g == ek_bytes_rt)  h_padconst = 8'h06;
            else if (g == h_last_rt)    h_padconst = 8'h80;
            else                        h_padconst = 8'h00;
        end
    endfunction

    // Global ek byte index for the address currently presented in ST_H assemble.
    wire [11:0] h_g_addr = {8'd0, h_blk} * 12'd136 + {4'd0, h_byte};

    // ek BRAM read-address mux: ST_H assemble drives it; otherwise debug readback.
    // dbg_dk ek-region takes priority (its loop runs after the ek loop), so the
    // single read port serves dbg_byte_o (ek loop) then dbg_dk_o (dk loop).
    wire [11:0] dbgdk_ek_off = dbg_dk_idx_i - dk_bytes_rt;   // offset within ek region
    assign ek_rd_addr = (st == ST_H || st == ST_ENC_H) ? h_g_addr[10:0] :
                        (st == ST_ENC_LOAD) ? rl_ek_off[10:0] :   // rho load
                        (st == ST_ENC_TDEC) ? td_ekaddr[10:0] :   // byteDecode12
                        dbgdk_in_ek   ? dbgdk_ek_off[10:0]
                                      : dbg_byte_idx_i;
    // dkp BRAM read-address mux: D2 byteDecode12 (ST_DEC_SDEC) walks dk_pke via
    // td_ekaddr (poly*384 + trip*3 + ph); else dbg_byte (sel=1) / dbg_dk.
    assign dkp_rd_addr = (st == ST_DEC_SDEC) ? td_ekaddr[10:0] :
                         dbg_byte_sel_i ? dbg_byte_idx_i : dbg_dk_idx_i[10:0];

    // ---- sample_ntt_sync: Â[i][j] = SampleNTT(rho || j || i) ----
    reg          snt_valid;
    wire         snt_ready;
    wire [11:0]  snt_coeff;
    wire         snt_vo;
    wire         snt_last;
    reg          snt_ack;       // we accept coeffs

    sample_ntt_sync_shared #(.K(KMAX)) u_snt (
        .clk(clk), .rst_n(rst_n),
        .rho_i(rho_r),
        .k_i(k_r),
        .i_idx(a_i[1:0]),
        .j_idx(a_j[1:0]),
        .valid_i(snt_valid),
        .ready_o(snt_ready),
        .coeff_o(snt_coeff),
        .valid_o(snt_vo),
        .ready_i(snt_ack),
        .last_o(snt_last),
        // shared keccak_core interface
        .kc_state_o(kc_state_o),
        .kc_valid_o(kc_valid_o_snt),
        .kc_ready_o(kc_ready_o),
        .kc_state_i(snt_kc_state_i),
        .kc_valid_i(snt_kc_valid_i),
        .kc_ready_i(snt_kc_ready_i)
    );

    // ---- sample_cbd_sync: s[i]=CBD3(PRF(sigma,i)), e[i]=CBD3(PRF(sigma,K+i)) ----
    reg          cbd_valid;
    wire         cbd_ready;
    wire [11:0]  cbd_coeff;       // 12-bit signed (two's complement)
    wire         cbd_vo;
    wire         cbd_last;
    reg          cbd_ack;

    // CBD inputs muxed KeyGen vs Encaps:
    //   KeyGen ST_C: seed=sigma, eta=eta1 (all polys), nonce=c_poly.
    //   Encaps ST_ENC_C: seed=r, nonce=c_poly; eta = eta1 for y (c_poly<K),
    //   eta2 for e1/e2 (c_poly>=K). nonce sequence 0..2K matches FIPS.
    wire        cbd_enc   = (st == ST_ENC_C);
    wire [255:0] cbd_seed = cbd_enc ? r_r : sigma_r;
    wire [1:0]  cbd_eta   = cbd_enc ? ((c_poly < {2'b0, k_r}) ? eta1_rt : eta2_rt)
                                    : eta1_rt;
    sample_cbd_sync_shared u_cbd (
        .clk(clk), .rst_n(rst_n),
        .seed_i(cbd_seed),
        .nonce_i(c_nonce),
        .eta_i(cbd_eta),
        .valid_i(cbd_valid),
        .ready_o(cbd_ready),
        .coeff_o(cbd_coeff),
        .valid_o(cbd_vo),
        .ready_i(cbd_ack),
        .last_o(cbd_last),
        // shared keccak_core interface
        .kc_state_o(kc_state_o),
        .kc_valid_o(kc_valid_o_cbd),
        .kc_ready_o(kc_ready_o),
        .kc_state_i(cbd_kc_state_i),
        .kc_valid_i(cbd_kc_valid_i),
        .kc_ready_i(cbd_kc_ready_i)
    );

    // signed (two's complement) -> [0,Q): add Q when negative
    wire [11:0] cbd_modq = cbd_coeff[11] ? (cbd_coeff + 12'(Q)) : cbd_coeff;

    // ---- ntt_core: forward NTT (mode=0, no scaling) of s[i],e[i] in place ----
    // N-stage bookkeeping: process slots S0,S1,E0,E1 (= SLOT_S0 + n_slot).
    reg  [4:0]  n_slot;        // 0..2K (process s_hat[0..K-1] then e_hat[0..K-1])
    reg  [8:0]  n_ridx;        // load read-AHEAD pointer 0..256 (leads consume by 1)
    reg  [7:0]  n_widx;        // output write index 0..255
    reg         n_valid;       // feeding coeffs to ntt_core (delayed 1 cyc vs n_ridx)
    reg         n_loading;     // 1 while presenting load addresses to bank_se
    reg         n_pending;     // waiting for ntt_core IDLE to start next slot
    wire [SAW-1:0] n_slot_addr = slot_s_rt + n_slot;   // s_hat then e_hat contiguous
    // NTT slot count: KeyGen processes 2K (s,e); Encaps processes K (y only).
    wire [4:0]     n_slot_max  = (st == ST_ENC_N || st == ST_DEC_NTT) ? {2'b0, k_r} : {1'b0, k_r, 1'b0};

    wire        ntt_ready;
    wire [11:0] ntt_coeff;
    wire        ntt_vo;
    wire        ntt_done;
    // bank_se read addr for the NTT load (relative slot = n_slot); sd_bram
    // registers it into bse_rd_data, which feeds ntt_core 1 cycle later.
    wire [13:0] ntt_rd_full = n_slot*256 + n_ridx[7:0];
    // ntt_core inputs muxed: KeyGen/Encaps fwd NTT feeds bse_rd_data (mode 0);
    // Encaps E4/E6 INTT sub-phase feeds psum (bt_rd_data) with mode 1.
    wire        u_intt_act = (st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && (u_sub == 2'd1);
    wire [11:0] ntt_in   = u_intt_act ? bt_rd_data : bse_rd_data;
    wire        ntt_vin  = u_intt_act ? u_nvalid : n_valid;
    wire        ntt_mode = u_intt_act ? 1'b1 : 1'b0;

    ntt_core u_ntt (
        .clk(clk), .rst_n(rst_n),
        .coeff_in(ntt_in),
        .valid_i(ntt_vin),
        .ready_o(ntt_ready),
        .mode(ntt_mode),              // 0 = fwd (NTT), 1 = inv (INTT, x3303)
        .coeff_out(ntt_coeff),
        .valid_o(ntt_vo),
        .ready_i(1'b1),               // always accept output
        .done_o(ntt_done)
    );

    // ---- poly_mul_sync: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] ----
    // M-stage bookkeeping. For each (i,j): LOAD 256 (A,shat) pairs, then accumulate
    // 256 products into T_i (init from E_i when j==0, else from running T_i).
    reg  [2:0]  m_i;           // row 0..K (needs to reach K to exit)
    reg  [2:0]  m_j;           // col 0..K-1
    reg  [8:0]  m_ld;          // load index 0..256
    reg  [7:0]  m_oidx;        // output/accum index 0..255
    reg         m_loading;     // 1 while streaming pairs into poly_mul
    reg         m_pending;     // wait for poly_mul IDLE before next (i,j)

    // ---- Stage 2f: byteEncode12 serializer ----
    // Pack each poly (2 coeffs -> 3 bytes, LSB-first 12-bit). ek = t_hat[0..K-1]
    // bytes || rho; dk_pke = s_hat[0..K-1] bytes. Walk coeff pairs per poly.
    reg  [4:0]  e_poly;        // 0..2K-1: [0,K) = t_hat -> ek; [K,2K) = s_hat -> dk_pke
    reg  [7:0]  e_pair;        // 0..127 coeff-pair within poly
    reg  [9:0]  e_rho;         // 0..31 rho byte copy index (ek tail)
    reg         e_done;        // serialization complete
    // source poly slot: t_hat[e_poly] for ek half, s_hat[e_poly-K] for dk half
    wire        e_is_dk = (e_poly >= {1'b0, k_r});
    wire [4:0]  e_pidx  = e_is_dk ? (e_poly - {1'b0, k_r}) : e_poly;   // index within target
    // Registered single-port read (sd_bram timing). One bank read per cycle:
    // ek half reads bank_t (t_hat), dk half reads bank_se (s_hat). The two
    // coeffs of a pair are serialized across a 4-cycle micro-phase e_ph:
    //   0 = fetch c0
    //   1 = c0 ready: write b0, save c0[11:8], fetch c1
    //   2 = c1 ready: write b1, save c1[11:4]
    //   3 = write b2, advance pair
    reg  [1:0]  e_ph;          // per-pair micro-phase 0..3
    reg  [3:0]  e_c0_hi;       // saved c0[11:8] for b1
    reg  [7:0]  e_c1_hi;       // saved c1[11:4] for b2
    wire        e_rd_half = (e_ph == 2'd1);          // ph0 -> c0, ph1 -> c1
    wire [13:0] e_rd_full = e_pidx*256 + {e_pair, e_rd_half};   // bse/bt rd addr in ST_E
    // dk half reads bank_se (bse_rd_data), ek half reads bank_t (bt_rd_data);
    // both registered inside their sd_bram with the same 1-cycle latency, so
    // the coeff for the addr presented last cycle is selected here.
    wire [11:0] e_rd_coeff = e_is_dk ? bse_rd_data : bt_rd_data;
    // byteEncode write byte offset within the target memory: pair*3 + byte index.
    wire [11:0] e_base = e_pidx * 12'd384;           // poly index *384 (=128 pairs*3)
    wire [11:0] e_boff = e_base + {e_pair, 1'b0} + {2'b0, e_pair};  // pair*3
    wire [1:0]  e_wb   = e_ph - 2'd1;                // ph1->byte0, ph2->byte1, ph3->byte2
    wire [7:0]  e_wbyte = (e_ph == 2'd1) ? e_rd_coeff[7:0]             // b0 = c0[7:0]
                        : (e_ph == 2'd2) ? {e_rd_coeff[3:0], e_c0_hi}  // b1 = {c1[3:0],c0[11:8]}
                        :                  e_c1_hi;                      // b2 = c1[11:4]

    wire [SAW-1:0]  m_aslot = m_i*k_r + m_j;     // A_hat[i][j] slot = i*k + j
    wire [SAW-1:0]  m_sslot = slot_s_rt + m_j;     // s_hat[j]
    wire [SAW-1:0]  m_eslot = slot_e_rt + m_i;     // e_hat[i]
    wire [SAW-1:0]  m_tslot = slot_t_rt + m_i;     // t_hat[i]

    // ================================================================
    // E4: u[i] = INTT(sum_j A^T[i][j] o y_hat[j]) + e1[i]  (Encaps ST_ENC_U)
    // Per row i, three sub-phases (u_sub):
    //   0 MAC : sum_j A_hat[j][i] o y_hat[j] (transpose) -> psum bank_t[UPSUM]
    //   1 INTT: INTT(psum) mode=1 in place in bank_t[UPSUM]
    //   2 ADD : u[i][w] = psum[w] + e1[i][w] mod Q -> bank_se rel (K+i) (over e1)
    // y_hat (bank_se 0..K-1) survives for V. Reuses shared u_pmul + u_ntt.
    // ================================================================
    localparam UPSUM = 10'd1;          // bank_t rel slot for the NTT-domain psum
    reg  [2:0]  u_row;                 // output row i 0..K-1
    reg  [1:0]  u_sub;                 // 0=MAC 1=INTT 2=ADD
    // MAC bookkeeping (mirrors ST_M)
    reg  [2:0]  u_j;                   // term col j 0..K-1
    reg  [8:0]  u_ld;                  // load read-ahead 0..256
    reg  [7:0]  u_oidx;                // accumulate index 0..255
    reg         u_loading;             // streaming pairs into poly_mul
    reg         u_pending;             // wait poly_mul IDLE before next j
    reg         u_j0q;                 // (u_j==0) delayed 1 cyc (first-term=0 select)
    // INTT bookkeeping (mirrors ST_N)
    reg  [8:0]  u_ridx;                // INTT load read-ahead 0..256
    reg  [7:0]  u_widx;                // INTT output index 0..255
    reg         u_nvalid;              // feed psum to ntt_core (mode=1)
    reg         u_nloading;            // presenting INTT load addresses
    reg         u_npending;            // (unused placeholder for symmetry)
    // ADD bookkeeping
    reg  [8:0]  u_aidx;                // ADD read-ahead 0..256
    reg  [7:0]  u_awidx;               // ADD write index 0..255
    reg         u_avalid;              // ADD write enable pipeline
    // MAC addresses: A^T -> A_hat[u_j][u_row] slot = u_j*k+u_row; y_hat[u_j] rel u_j
    wire [SAW-1:0] u_aslot   = u_j*k_r + u_row;
    wire [13:0] u_pm_a_full  = u_aslot*256 + u_ld[7:0];   // bank_a
    wire [13:0] u_pm_b_full  = u_j*256 + u_ld[7:0];        // bank_se y_hat[u_j]
    wire [7:0]  u_acc_radr   = pm_vo ? (u_oidx + 8'd1) : u_oidx;
    wire [13:0] u_psum_full  = UPSUM*256 + u_acc_radr;     // bank_t psum acc addr (MAC)
    wire [11:0] u_acc_src    = u_j0q ? 12'd0 : bt_rd_data; // first term init 0
    wire [12:0] u_sum        = {1'b0, u_acc_src} + {1'b0, pm_coeff};
    wire [11:0] u_accq       = (u_sum >= 13'(Q)) ? (u_sum - 13'(Q)) : u_sum[11:0];
    // INTT load addr (psum read) + in-place writeback to bank_t[UPSUM]
    wire [13:0] u_intt_rd    = UPSUM*256 + u_ridx[7:0];
    wire [13:0] u_intt_wr    = UPSUM*256 + u_widx;
    // ADD: read psum (bank_t) + e1 (bank_se K+u_row), write u over e1 (read-ahead 1)
    wire [13:0] u_add_prd    = UPSUM*256 + u_aidx[7:0];                       // bank_t psum
    wire [13:0] u_add_e1rd   = ({2'b0,k_r}+{2'b0,u_row})*256 + u_aidx[7:0];   // bank_se e1 (lead)
    wire [13:0] u_add_uwr    = ({2'b0,k_r}+{2'b0,u_row})*256 + u_awidx;       // bank_se u (write)
    wire [12:0] u_usum       = {1'b0, bt_rd_data} + {1'b0, bse_rd_data};       // psum + e1
    wire [11:0] u_uq         = (u_usum >= 13'(Q)) ? (u_usum - 13'(Q)) : u_usum[11:0];

    // ================================================================
    // E6: v = INTT(sum_j t_hat[j] o y_hat[j]) + e2 + mu  (Encaps ST_ENC_V).
    // Reuses the u_* MAC/INTT/ADD machine with u_row tied to 0 (single output
    // poly). Differences from E4:
    //   - MAC reads t_hat[j] from bank_a slot j*K (== u_aslot=u_j*K, u_row=0),
    //     so the bank_a MAC address is unchanged.
    //   - ADD computes psum + e2 + mu (not psum + e1); e2 lives in bank_a slot
    //     E2_ASLOT (relocated from bank_t[0] in ST_ENC_E2MV), v writes back to
    //     bank_t[UPSUM] in place.
    //   - mu[w] = Decompress_1(m bit w) = m_r[w] ? 1665 : 0 (1665 = round(Q/2)).
    // u_row_max bounds the row loop: K for U, 1 for V.
    localparam E2_ASLOT = 12'd1;       // bank_a slot holding relocated e2 (never a t_hat slot, K>=2)
    wire [2:0]  u_row_max = (st == ST_ENC_V || st == ST_DEC_W) ? 3'd1 : k_r;
    // V-ADD: read psum (bank_t UPSUM) + e2 (bank_a E2_ASLOT) at u_aidx (lead),
    // add mu (by write index u_awidx), write v to bank_t[UPSUM] at u_awidx.
    wire [13:0] u_v_e2rd     = E2_ASLOT*256 + u_aidx[7:0];                    // bank_a e2 (lead)
    wire [13:0] u_vadd_wr    = UPSUM*256 + u_awidx;                            // bank_t v (write)
    wire [11:0] u_mu         = m_r[u_awidx] ? 12'd1665 : 12'd0;               // mu by msg bit
    wire [13:0] u_vsum       = {2'b0, bt_rd_data} + {2'b0, ba_rd_data} + {2'b0, u_mu}; // psum+e2+mu
    wire [13:0] u_vsub1      = (u_vsum  >= 14'(Q)) ? (u_vsum  - 14'(Q)) : u_vsum;
    wire [11:0] u_vq         = (u_vsub1 >= 14'(Q)) ? (u_vsub1 - 14'(Q)) : u_vsub1[11:0];
    wire        u_v_we       = (st == ST_ENC_V) && (u_sub == 2'd2) && u_avalid;

    // ================================================================
    // D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j])  (Decaps ST_DEC_W).
    // Reuses the u_* MAC/INTT machine (u_row tied to 0, single poly):
    //   MAC : s_hat[j] (bank_a slot j*K, == u_aslot with u_row=0) o u_hat[j]
    //         (bank_se rel j, == u_pm_b_full) -> psum bank_t[UPSUM].  (same
    //         addressing as Encaps V MAC, which reads t_hat[j] o y_hat[j].)
    //   INTT: INTT(psum) in place in bank_t[UPSUM].
    //   SUB : w[i] = v'[i] - psum[i] mod Q. v' is in bank_a[DEC_VASLOT] (D1),
    //         psum in bank_t[UPSUM]; read both in parallel (like V-ADD reads
    //         psum + e2). w written back to bank_t[UPSUM] in place.
    wire [13:0] u_w_vrd     = DEC_VASLOT*256 + u_aidx[7:0];                   // bank_a v' (lead)
    wire [13:0] u_wsub_wr   = UPSUM*256 + u_awidx;                            // bank_t w (write)
    wire [12:0] u_wdiff     = {1'b0, ba_rd_data} - {1'b0, bt_rd_data};        // v' - psum
    wire [11:0] u_wq        = u_wdiff[12] ? (u_wdiff[11:0] + 12'(Q)) : u_wdiff[11:0];
    wire        u_w_we      = (st == ST_DEC_W) && (u_sub == 2'd2) && u_avalid;

    // ================================================================
    // D4: m' = byteEncode_1(Compress_1(w))  (Decaps ST_DEC_MENC).
    // Compress_1(x) = round(2x/Q) mod 2 = 1 iff x in (Q/4, 3Q/4], i.e.
    // 832 < x <= 2496 for Q=3329. Pack 256 bits LSB-first into 32 bytes.
    // w is read from bank_t[UPSUM] (in place from D3), 1 coeff per 2 cycles.
    reg  [7:0]  men_idx;        // coeff 0..255
    reg  [1:0]  men_ph;         // micro-phase
    reg         men_done;
    // D4 packs the recovered message bits directly into m_r (reused by D6's
    // V/mu re-encrypt and exposed on dbg_mprime_o).

    // ================================================================
    // D7: implicit-reject compare (Decaps ST_DEC_CMP).
    // Walk all ct bytes, reading c' (ct_bram) and c (c_in_bram) in parallel;
    // OR every byte difference into cmp_neq. After the last byte, ss = cmp_neq
    // ? K-bar : K'. The walk visits every byte regardless of early differences
    // (constant work / no early-out), matching the constant-time spec intent.
    //   ph0: present both addrs; ph1: compare registered reads, accumulate.
    reg  [10:0] cmp_idx;        // byte index 0..ct_bytes
    reg         cmp_ph;         // 0=present addr, 1=compare
    reg         cmp_neq;        // sticky: any byte differed
    reg         cmp_done;
    reg         dec_reject;     // latched cmp_neq at end (drives ss select)
    wire        cmp_byte_neq = (ct_rd_data != cin_rd_data);
    wire [13:0] men_rd   = UPSUM*256 + men_idx;          // bank_t w read addr
    wire        men_w_bit = (bt_rd_data > 12'd832) && (bt_rd_data <= 12'd2496); // Compress_1

    // ST_ENC_E2MV: copy e2 (bank_t rel slot 0, 256 coeffs) into bank_a[E2_ASLOT].
    // bank_t read leads (em_ridx), bank_a write trails 1 cycle (em_we/em_widx).
    reg  [8:0]  em_ridx;               // 0..256 read-ahead over e2
    reg  [7:0]  em_widx;               // bank_a write index
    reg         em_we;                 // a coeff is being written this cycle
    reg         em_done;
    wire [13:0] em_rd        = 14'd0 + em_ridx[7:0];          // bank_t rel slot 0 (e2)
    wire [11:0] em_wa        = E2_ASLOT*256 + em_widx;        // bank_a E2_ASLOT

    // ================================================================
    // E5/E7: Compress_d + byteEncode_d -> ciphertext (Encaps ST_ENC_C1/C2).
    // Per coeff: read poly coeff -> comp_decomp (mode 0 compress, d=du/dv) ->
    // bit-packer (LSB-first) -> emit bytes to ct_bram. c1 = K polys of u
    // (d=du), then c2 = 1 poly of v (d=dv). Per poly = 256 coeffs -> 32*d
    // bytes (whole), so the bit buffer empties at each poly boundary.
    //   micro-phase cp_ph: 0 present coeff addr; 1 feed comp_decomp (cd_valid);
    //   2 wait pipe; 3 capture compressed + accumulate bits; 4..n drain bytes.
    // ================================================================
    wire        cd_active = (st == ST_ENC_C1) || (st == ST_ENC_C2);
    reg  [11:0] cd_coeff;          // coeff presented to comp_decomp
    reg         cd_valid;          // 1-cyc pulse to comp_decomp
    wire        cd_ready;
    wire [11:0] cd_out;            // compressed value (low d bits valid)
    wire        cd_vo;
    wire [4:0]  cp_d = (st == ST_ENC_C2) ? dv_rt : du_rt;   // compress width
    // comp_decomp is shared: Encaps C1/C2 compress (mode 0); Decaps DECOMP
    // decompresses (mode 1). In DECOMP the decoder feeds dec_dc_coeff/valid/d.
    wire        cd_dec   = (st == ST_DEC_DECOMP);
    wire        cd_mode  = cd_dec ? 1'b1 : 1'b0;             // 1=decompress
    wire [4:0]  cd_d_mux = cd_dec ? dec_dc_d : cp_d;
    wire [11:0] cd_in_mux= cd_dec ? dec_dc_coeff : cd_coeff;
    wire        cd_vi_mux= cd_dec ? dec_dc_valid : cd_valid;
    comp_decomp_sync u_comp (
        .clk(clk), .rst_n(rst_n),
        .coeff_in(cd_in_mux),
        .d(cd_d_mux),
        .mode(cd_mode),
        .valid_i(cd_vi_mux),
        .ready_o(cd_ready),
        .coeff_out(cd_out),
        .valid_o(cd_vo),
        .ready_i(1'b1)
    );
    // bit-packer / ct walk bookkeeping
    reg  [2:0]  cp_poly;           // c1: 0..K-1 (u rows); c2: single v
    reg  [7:0]  cp_idx;            // coeff 0..255 within poly
    reg  [2:0]  cp_ph;             // micro-phase
    reg  [24:0] cp_buf;            // bit accumulator (LSB-first)
    reg  [5:0]  cp_nbits;          // valid bits in cp_buf
    reg  [11:0] cp_wa;             // ct_bram byte write address (runs c1 then c2)
    reg         cp_done;           // serialization complete (this region)
    // coeff source: c1 (ST_ENC_C1) reads u[cp_poly] in bank_se rel (K+cp_poly);
    // c2 (ST_ENC_C2) reads v in bank_t rel slot UPSUM (single poly, cp_poly=0).
    wire [13:0] cp_se_full = ({2'b0,k_r}+{2'b0,cp_poly})*256 + cp_idx;  // bank_se u[cp_poly]
    wire [13:0] cp_bt_full = UPSUM*256 + cp_idx;                        // bank_t v (UPSUM)
    // packer generalization (E5 c1 vs E7 c2):
    //   coeff source: c1 = bank_se (u), c2 = bank_t (v)
    //   poly count:   c1 = K (u rows), c2 = 1 (single v)
    //   bit width cp_d already = dv_rt for C2 else du_rt (declared above).
    wire [11:0] cp_coeff_src = (st == ST_ENC_C2) ? bt_rd_data : bse_rd_data;
    wire [2:0]  cp_poly_max  = (st == ST_ENC_C2) ? 3'd1 : k_r;

    // ================================================================
    // D1: K-PKE byteDecode_d + Decompress (decrypt c1/c2 -> u'/v').
    // Reverse of the C1/C2 packer: walk c_in_bram bytes, accumulate LSB-first
    // into a bit buffer, extract d-bit symbols, feed comp_decomp (mode=1
    // decompress), and write each result (mod q) to a bank.
    //   c1 = K polys, d=du -> u'[i] in bank_se rel i (0..K-1)
    //   c2 = 1 poly,  d=dv -> v'   in bank_t rel slot DEC_VSLOT
    //   c1 = K polys, d=du -> u'[i] in bank_se rel i (0..K-1)
    //   c2 = 1 poly,  d=dv -> v'   in bank_a slot DEC_VASLOT (so D3's w-SUB can
    //        read psum (bank_t) and v' (bank_a) in parallel, like Encaps V-ADD).
    // micro-phase dec_ph:
    //   0 present c_in byte addr;   1 capture byte, append 8 bits to buf;
    //   while (nbits>=d): extract d-bit symbol -> comp_decomp (decompress);
    //   2 wait pipe; 3 cd_vo -> write coeff to bank, advance coeff index.
    // ================================================================
    localparam DEC_VASLOT = 12'd1;      // bank_a slot for v' (free: s_hat at j*K, slot 1 unused K>=2)
    reg  [2:0]  dec_poly;               // c1: 0..K-1; c2: single (=K marker)
    reg  [7:0]  dec_cidx;               // coeff 0..255 within poly (feed side)
    reg  [7:0]  dec_widx;               // coeff 0..255 within poly (writeback side, lags pipe)
    reg  [1:0]  dec_ph;                 // micro-phase
    reg  [11:0] dec_ba;                 // c_in_bram byte read address (c1 then c2)
    reg  [24:0] dec_buf;                // bit accumulator (LSB-first)
    reg  [5:0]  dec_nbits;              // valid bits in dec_buf
    reg         dec_dc_done;            // D1 complete (both c1 and c2 decoded)
    reg         dec_in_c2;              // 0 = decoding c1 (u'), 1 = decoding c2 (v')
    // comp_decomp feed (decompress): the extracted d-bit symbol.
    reg  [11:0] dec_dc_coeff;
    reg         dec_dc_valid;
    wire [4:0]  dec_dc_d = dec_in_c2 ? dv_rt : du_rt;   // c2 uses dv, c1 uses du
    // (c_in_bram read addr driven by cin_rd_addr_r, set to dec_ba in the walker)

    reg          pm_valid;
    wire         pm_ready;
    wire [11:0]  pm_coeff;
    wire         pm_vo;
    // pm_a: A_hat[i][j] in bank_a (abs slot m_aslot). pm_b: s_hat[j] in
    // bank_se (relative slot = m_sslot - slot_s_rt = m_j). m_ld is a read-ahead
    // pointer; both bank reads come from their sd_bram (registered, 1-cycle
    // latency) and feed poly_mul one cycle later (pm_valid delayed 1 cyc).
    wire [13:0] pm_a_full = m_aslot*256 + m_ld[7:0];
    wire [13:0] pm_b_full = m_j*256 + m_ld[7:0];
    wire [11:0]  pm_a_in = ba_rd_data;   // bank_a sd_bram registered read
    wire [11:0]  pm_b_in = bse_rd_data;  // bank_se sd_bram registered read (load phase)

    poly_mul_sync u_pmul (
        .clk(clk), .rst_n(rst_n),
        .coeff_a_in(pm_a_in),
        .coeff_b_in(pm_b_in),
        .valid_i(pm_valid),
        .ready_o(pm_ready),
        .coeff_out(pm_coeff),
        .valid_o(pm_vo),
        .ready_i(1'b1)
    );

    // accumulator source: e_hat[i] for first term (j==0), else running t_hat[i].
    // e_hat[i] lives in bank_se at relative slot (slot_e_rt-slot_s_rt + m_i) = K+m_i.
    // t_hat[i] lives in bank_t at relative slot m_i. Both reads come from their
    // sd_bram (bse_rd_data / bt_rd_data) with the same 1-cycle latency.
    // Registered read-ahead: present the index the NEXT pm_vo will consume
    //   m_acc_radr = pm_vo ? m_oidx+1 : m_oidx
    // The j-select is applied on the registered outputs (m_jq = (m_j==0)
    // delayed 1 cyc to align with the read latency). RMW read-old holds: read
    // addr (m_oidx+1) leads write addr (m_oidx). Cadence CALC/C0/C1.
    wire [7:0]  m_acc_radr = pm_vo ? (m_oidx + 8'd1) : m_oidx;
    wire [13:0] m_eacc_full = ({2'b0, k_r} + {2'b0, m_i})*256 + m_acc_radr;  // K+m_i (bse_rd_addr in acc)
    wire [13:0] m_tacc_full = m_i*256 + m_acc_radr;                          // bt_rd_addr
    reg         m_jq;          // (m_j==0) delayed 1 cyc to match read latency
    // selected accumulator source aligned with pm_coeff
    wire [11:0]  m_acc_src = m_jq ? bse_rd_data : bt_rd_data;
    // (a + b) mod Q  (both < Q, sum < 2Q): one conditional subtract
    wire [12:0]  m_sum  = {1'b0, m_acc_src} + {1'b0, pm_coeff};
    wire [11:0]  m_accq = (m_sum >= 13'(Q)) ? (m_sum - 13'(Q)) : m_sum[11:0];

    always @(*) begin
        st_next = st;
        case (st)
            ST_IDLE: if (start_i) st_next = (op_i == 2'd2) ? ST_DEC_LOAD :
                                            (op_i == 2'd1) ? ST_ENC_H    : ST_G;
            // D0: settle after dk/c parse, then D1 Decompress. D1 lands in DONE
            // once u'/v' are computed so the stage can be dbg-checked.
            ST_DEC_LOAD: st_next = ST_DEC_DECOMP;
            ST_DEC_DECOMP: if (dec_dc_done) st_next = ST_DEC_SDEC;
            // D2: s_hat decode (reuse TDEC machine), then u_hat = NTT(u'). D2
            // lands in DONE once u_hat is ready so the stage can be dbg-checked.
            ST_DEC_SDEC: if (td_done) st_next = ST_DEC_NTT;
            ST_DEC_NTT:  if (n_slot >= {2'b0, k_r}) st_next = ST_DEC_W;
            // D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j]). Single output poly
            // (u_row 0..0); MAC->INTT->SUB, then D4 encodes m'.
            ST_DEC_W:    if (u_row >= 3'd1) st_next = ST_DEC_MENC;
            ST_DEC_MENC: if (men_done) st_next = ST_DEC_G;
            // D5: (K',r') = G(m'||h) single-block, then K-bar = J(z||c) multi-block.
            ST_DEC_G:    if (sha3_vo) st_next = ST_DEC_J;
            // D5 J done -> D6 re-encrypt: c' = K-PKE.Encrypt(ek_pke, m', r').
            // Reuse the entire Encaps pipeline (rho load -> A -> C -> ... -> C2).
            // r' is in r_r (CBD seed), m' in m_r (V/mu), ek_pke in ek_bram.
            ST_DEC_J:    if (dj_done) st_next = ST_ENC_LOAD;
            ST_G:    if (sha3_vo) st_next = ST_A;
            ST_A:    if (a_pair >= kk_rt) st_next = ST_C;
            ST_C:    if (c_poly >= {1'b0, k_r, 1'b0}) st_next = ST_N;
            ST_N:    if (n_slot >= {1'b0, k_r, 1'b0}) st_next = ST_M;
            ST_M:    if (m_i >= k_r) st_next = ST_E;
            ST_E:    if (e_done) st_next = ST_H;
            ST_H:    if (h_phase == 2'd3) st_next = ST_DONE;
            // ---- Encaps ----
            ST_ENC_H: if (h_phase == 2'd3) st_next = ST_ENC_G;
            ST_ENC_G: if (sha3_vo) st_next = ST_ENC_LOAD;   // (K,r) captured
            ST_ENC_LOAD: if (rl_idx >= 6'd33) st_next = ST_ENC_A;  // rho loaded
            ST_ENC_A:    if (a_pair >= kk_rt) st_next = ST_ENC_C;
            ST_ENC_C:    if (c_poly >= {k_r, 1'b1}) st_next = ST_ENC_N;  // 2K+1 polys done
            ST_ENC_N:    if (n_slot >= {2'b0, k_r}) st_next = ST_ENC_U;     // K slots (y_hat)
            ST_ENC_U:    if (u_row >= k_r) st_next = ST_ENC_C1;            // u[0..K-1] done
            ST_ENC_C1:   if (cp_done) st_next = ST_ENC_TDEC;               // E5 done -> E6
            ST_ENC_TDEC: if (td_done) st_next = ST_ENC_E2MV;   // t_hat decoded -> relocate e2
            ST_ENC_E2MV: if (em_done) st_next = ST_ENC_V;      // e2 relocated -> compute v
            ST_ENC_V:    if (u_row >= 3'd1) st_next = ST_ENC_C2;            // E6: v done -> E7
            ST_ENC_C2:   if (cp_done) st_next = (op_r == 2'd2) ? ST_DEC_CMP : ST_DONE;
            // D7: walk ct bytes comparing c' (ct_bram) vs c (c_in_bram); when the
            // last byte is compared, select ss = (c'==c) ? K' : K-bar and finish.
            ST_DEC_CMP:  if (cmp_done) st_next = ST_DONE;
            ST_DONE: st_next = ST_IDLE;
            default: st_next = ST_IDLE;
        endcase
    end

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            st         <= ST_IDLE;
            k_r        <= 3'd0;
            op_r       <= 2'd0;
            m_r        <= 256'd0;
            ss_r       <= 256'd0;
            r_r        <= 256'd0;
            rl_idx     <= 6'd0;
            rl_widx    <= 6'd0;
            rl_vld     <= 1'b0;
            td_poly    <= 3'd0;
            td_trip    <= 8'd0;
            td_ph      <= 3'd0;
            td_b0      <= 8'd0;
            td_b1      <= 8'd0;
            td_b2      <= 8'd0;
            td_done    <= 1'b0;
            td_we      <= 1'b0;
            td_wa      <= 12'd0;
            td_wd      <= 12'd0;
            rho_r      <= 256'd0;
            sigma_r    <= 256'd0;
            sha3_valid <= 1'b0;
            sha3_ack   <= 1'b0;
            snt_valid  <= 1'b0;
            snt_ack    <= 1'b0;
            a_pair     <= 5'd0;
            a_i        <= 3'd0;
            a_j        <= 3'd0;
            a_widx     <= 8'd0;
            a_busy     <= 1'b0;
            cbd_valid  <= 1'b0;
            cbd_ack    <= 1'b0;
            c_poly     <= 3'd0;
            c_widx     <= 8'd0;
            c_busy     <= 1'b0;
            n_slot     <= 3'd0;
            n_ridx     <= 9'd0;
            n_widx     <= 8'd0;
            n_valid    <= 1'b0;
            n_loading  <= 1'b0;
            n_pending  <= 1'b0;
            m_i        <= 2'd0;
            m_j        <= 2'd0;
            m_ld       <= 9'd0;
            m_oidx     <= 8'd0;
            m_loading  <= 1'b0;
            m_pending  <= 1'b0;
            pm_valid   <= 1'b0;
            m_jq       <= 1'b0;
            u_row      <= 3'd0;
            u_sub      <= 2'd0;
            u_j        <= 3'd0;
            u_ld       <= 9'd0;
            u_oidx     <= 8'd0;
            u_loading  <= 1'b0;
            u_pending  <= 1'b0;
            u_j0q      <= 1'b0;
            u_ridx     <= 9'd0;
            u_widx     <= 8'd0;
            u_nvalid   <= 1'b0;
            u_nloading <= 1'b0;
            u_npending <= 1'b0;
            u_aidx     <= 9'd0;
            u_awidx    <= 8'd0;
            u_avalid   <= 1'b0;
            em_ridx    <= 9'd0;
            em_widx    <= 8'd0;
            em_we      <= 1'b0;
            em_done    <= 1'b0;
            cd_coeff   <= 12'd0;
            cd_valid   <= 1'b0;
            cp_poly    <= 3'd0;
            cp_idx     <= 8'd0;
            cp_ph      <= 3'd0;
            cp_buf     <= 25'd0;
            cp_nbits   <= 6'd0;
            cp_wa      <= 12'd0;
            cp_done    <= 1'b0;
            ct_we      <= 1'b0;
            ct_wa      <= 11'd0;
            ct_wd      <= 8'd0;
            e_poly     <= 3'd0;
            e_pair     <= 8'd0;
            e_ph       <= 2'd0;
            e_c0_hi    <= 4'd0;
            e_c1_hi    <= 8'd0;
            e_rho      <= 10'd0;
            e_done     <= 1'b0;
            ek_we      <= 1'b0;
            dkp_we     <= 1'b0;
            ek_wa      <= 11'd0;
            dkp_wa     <= 11'd0;
            ek_wd      <= 8'd0;
            dkp_wd     <= 8'd0;
            h_block_r  <= 1088'd0;
            h_mbvalid  <= 1'b0;
            h_mblast   <= 1'b0;
            h_ack      <= 1'b0;
            hek_r      <= 256'd0;
            z_r        <= 256'd0;
            kbar_r     <= 256'd0;
            cin_rd_addr_r <= 11'd0;
            dec_poly   <= 3'd0;
            dec_cidx   <= 8'd0;
            dec_widx   <= 8'd0;
            dec_ph     <= 2'd0;
            dec_ba     <= 12'd0;
            dec_buf    <= 25'd0;
            dec_nbits  <= 6'd0;
            dec_in_c2  <= 1'b0;
            dec_dc_coeff <= 12'd0;
            dec_dc_valid <= 1'b0;
            dec_dc_done  <= 1'b0;
            men_idx    <= 8'd0;
            men_ph     <= 2'd0;
            men_done   <= 1'b0;
            cmp_idx    <= 11'd0;
            cmp_ph     <= 1'b0;
            cmp_neq    <= 1'b0;
            cmp_done   <= 1'b0;
            dec_reject <= 1'b0;
            dj_blk     <= 4'd0;
            dj_byte    <= 8'd0;
            dj_phase   <= 2'd0;
            dj_mbvalid <= 1'b0;
            dj_mblast  <= 1'b0;
            dj_ack     <= 1'b0;
            dj_done    <= 1'b0;
            dj_wb_vld  <= 1'b0;
            dj_block_r <= 1088'd0;
            h_blk      <= 3'd0;
            h_byte     <= 8'd0;
            h_phase    <= 2'd0;
            h_wb_vld   <= 1'b0;
            h_wb_idx   <= 8'd0;
            h_wb_g     <= 12'd0;
            h_wb_pad   <= 8'd0;
            h_wb_inek  <= 1'b0;
        end else begin
            st <= st_next;

            // BRAM write-enables default low; pulsed where a byte is written.
            ek_we  <= 1'b0;
            dkp_we <= 1'b0;
            td_we  <= 1'b0;    // TDEC bank_a write default low
            ct_we  <= 1'b0;    // ct_bram byte write default low (E5/E7)
            em_we  <= 1'b0;    // e2-relocate bank_a write default low (E6)

            // Kick off when entering from IDLE: KeyGen starts G; Encaps captures
            // op/m and arms the H(ek) machinery (ST_ENC_H reuses the ST_H FSM);
            // Decaps captures op (dk/c already streamed; z/h captured at load).
            if (st == ST_IDLE && start_i) begin
                k_r  <= k_i;                // capture runtime ML-KEM param
                op_r <= op_i;
                if (op_i == 2'd1) begin
                    m_r      <= msg_i;      // capture Encaps message
                    // arm H(ek) (same fields the ST_E->ST_H arming sets)
                    h_blk    <= 3'd0;
                    h_byte   <= 8'd0;
                    h_phase  <= 2'd0;       // assemble
                    h_mbvalid<= 1'b0;
                    h_mblast <= 1'b0;
                    h_ack    <= 1'b1;       // ready to consume final digest
                    h_wb_vld <= 1'b0;
                end else if (op_i == 2'd2) begin
                    // Decaps: dk_pke/ek_pke already in BRAM; z/H(ek) captured into
                    // z_r/hek_r during load (below). Nothing else to arm in D0.
                end else begin
                    sha3_valid <= 1'b1;
                    sha3_ack   <= 1'b1;
                end
            end

            // Decaps dk load: capture the H(ek) and z byte regions into registers
            // as they stream in (the dk_pke/ek_pke regions go to BRAM via the
            // write muxes). Byte i within a region lands in [8i +: 8]. Uses the
            // LIVE k_i boundaries (k_r not yet captured during preload).
            if (dk_in_we) begin
                if (dk_in_addr >= dk_ek_end_ld && dk_in_addr < dk_hek_end_ld)
                    hek_r[(dk_in_addr - dk_ek_end_ld)*8 +: 8] <= dk_in_byte;   // H(ek)
                else if (dk_in_addr >= dk_hek_end_ld)
                    z_r[(dk_in_addr - dk_hek_end_ld)*8 +: 8]  <= dk_in_byte;   // z
            end
            // Drop valid once accepted
            if (sha3_valid && sha3_ready) sha3_valid <= 1'b0;

            // Capture rho/sigma when G completes; arm A stage
            if (st == ST_G && sha3_vo) begin
                rho_r   <= sha3_hash[255:0];     // rho = G output bytes 0..31
                sigma_r <= sha3_hash[511:256];   // sigma = bytes 32..63
                sha3_ack <= 1'b0;
                snt_valid <= 1'b1;               // start first SampleNTT
                snt_ack   <= 1'b1;
                a_pair    <= 5'd0;
                a_i       <= 3'd0;
                a_j       <= 3'd0;
                a_widx    <= 8'd0;
                a_busy    <= 1'b0;
            end

            // ---- ST_A / ST_ENC_A: drive SampleNTT, store 256 coeffs per pair ----
            if (st == ST_A || st == ST_ENC_A) begin
                // mark busy once this pair's request accepted
                if (snt_valid && snt_ready) begin
                    snt_valid <= 1'b0;
                    a_busy    <= 1'b1;
                end

                // store each output coefficient only while busy (ignore stale
                // last coeff from prior poly). The bank_a write itself is the
                // combinational ba_we/ba_wa/ba_wd assigns above; here we only
                // advance the write index and (i,j) bookkeeping.
                if (a_busy && snt_vo && snt_ack) begin
                    if (snt_last) begin
                        // finished this poly; advance (i,j) in row-major order
                        a_pair <= a_pair + 5'd1;
                        a_widx <= 8'd0;
                        a_busy <= 1'b0;
                        if (a_j + 3'd1 < k_r) begin
                            a_j <= a_j + 3'd1;
                        end else begin
                            a_j <= 3'd0;
                            a_i <= a_i + 3'd1;
                        end
                        // start next SampleNTT if more pairs remain
                        if (a_pair + 5'd1 < kk_rt) snt_valid <= 1'b1;
                    end else begin
                        a_widx <= a_widx + 8'd1;
                    end
                end
            end

            // Arm C stage when A finishes (KeyGen ST_A or Encaps ST_ENC_A)
            if ((st == ST_A && st_next == ST_C) ||
                (st == ST_ENC_A && st_next == ST_ENC_C)) begin
                cbd_valid <= 1'b1;
                cbd_ack   <= 1'b1;
                c_poly    <= 3'd0;
                c_widx    <= 8'd0;
                c_busy    <= 1'b0;
            end

            // ---- ST_C / ST_ENC_C: drive CBD, store 256 mod-q coeffs per poly ----
            // KeyGen: 2K polys (s,e). Encaps: 2K+1 polys (y,e1,e2). Loop bound
            // differs: KeyGen restarts while c_poly+1 < 2K; Encaps while < 2K+1.
            if (st == ST_C || st == ST_ENC_C) begin
                if (cbd_valid && cbd_ready) begin
                    cbd_valid <= 1'b0;
                    c_busy    <= 1'b1;
                end

                if (c_busy && cbd_vo && cbd_ack) begin
                    if (cbd_last) begin
                        c_poly <= c_poly + 3'd1;
                        c_widx <= 8'd0;
                        c_busy <= 1'b0;
                        // restart for next poly while more remain (state-dependent bound)
                        if (st == ST_ENC_C) begin
                            if (c_poly + 3'd1 < {k_r, 1'b1}) cbd_valid <= 1'b1;  // < 2K+1
                        end else begin
                            if (c_poly + 3'd1 < {1'b0, k_r, 1'b0}) cbd_valid <= 1'b1; // < 2K
                        end
                    end else begin
                        c_widx <= c_widx + 8'd1;
                    end
                end
            end

            // Arm N stage when C finishes: prime load of slot S0. n_ridx is a
            // read-ahead pointer; bank_se read is registered inside sd_bram (bse_rd_data)
            // and fed to ntt_core one cycle later, so valid starts low (priming).
            if ((st == ST_C && st_next == ST_N) ||
                (st == ST_ENC_C && st_next == ST_ENC_N) ||
                (st == ST_DEC_SDEC && st_next == ST_DEC_NTT)) begin
                n_slot    <= 3'd0;
                n_ridx    <= 9'd0;
                n_widx    <= 8'd0;
                n_valid   <= 1'b0;
                n_loading <= 1'b1;       // begin presenting load addresses
                n_pending <= 1'b0;
            end

            // ---- ST_N / ST_ENC_N / ST_DEC_NTT: forward NTT in place. KeyGen: 2K
            // slots (s,e). Encaps: K slots (y). Decaps D2: K slots (u' -> u_hat). ----
            if (st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) begin
                // slot-count bound: 2K for KeyGen, K for Encaps
                // (n_slot_max below); same LOAD/OUTPUT cadence either way.
                if (n_loading) begin
                    if (n_ridx == 9'd256) begin
                        // 256th coeff consumed this cycle; stop presenting addr
                        n_loading <= 1'b0;
                        n_valid   <= 1'b0;
                    end else begin
                        n_ridx    <= n_ridx + 9'd1;
                        n_valid   <= 1'b1;     // data presented last cycle is valid
                    end
                end

                // OUTPUT phase: collect 256 results, write back to same slot.
                if (ntt_vo) begin
                    n_widx <= n_widx + 8'd1;   // wraps 255->0 after last
                end

                // Slot complete when ntt_core returns to DONE
                if (ntt_done) begin
                    if (n_slot + 3'd1 < n_slot_max) begin
                        n_slot    <= n_slot + 3'd1;
                        n_widx    <= 8'd0;
                        n_pending <= 1'b1;     // wait one cycle for core IDLE
                    end else begin
                        n_slot <= n_slot + 3'd1;  // == n_slot_max -> exit
                    end
                end

                // Kick next slot's load once core is back IDLE (re-prime)
                if (n_pending && ntt_ready && !ntt_done) begin
                    n_ridx    <= 9'd0;
                    n_valid   <= 1'b0;
                    n_loading <= 1'b1;
                    n_pending <= 1'b0;
                end
            end

            // Arm M stage when N finishes: prime first (i=0,j=0) poly_mul load.
            // m_ld read-ahead pointer; bank_a/bank_se sd_bram reads land 1 cyc
            // later, pm_valid asserted one cycle after an address is presented.
            if (st == ST_N && st_next == ST_M) begin
                m_i       <= 2'd0;
                m_j       <= 2'd0;
                m_ld      <= 9'd0;
                m_oidx    <= 8'd0;
                m_loading <= 1'b1;
                m_pending <= 1'b0;
                pm_valid  <= 1'b0;
            end

            // ---- ST_M: t_hat[i] = e_hat[i] + sum_j A[i][j] o s_hat[j] ----
            if (st == ST_M) begin
                // accumulator reads come from sd_bram (bse_rd_data e_hat /
                // bt_rd_data t_hat). m_jq aligns the j-select with the 1-cycle
                // read latency. (bse_rd_addr muxes load vs acc by m_loading.)
                m_jq <= (m_j == 2'd0);

                // LOAD: present read-ahead addr to bank_a/bank_se via their
                // sd_bram read ports (ba_rd_addr=pm_a_full, bse_rd_addr=pm_b_full
                // while m_loading); ba_rd_data/bse_rd_data land next cycle and
                // are consumed by poly_mul (pm_valid). poly_mul holds ready high
                // through LOAD, so a fixed 1-cycle skew suffices.
                if (m_loading) begin
                    if (m_ld == 9'd256) begin
                        pm_valid  <= 1'b0;     // 256th pair consumed this cycle
                        m_loading <= 1'b0;
                        m_ld      <= 9'd0;
                        m_oidx    <= 8'd0;
                    end else begin
                        m_ld    <= m_ld + 9'd1;
                        pm_valid <= 1'b1;      // pair presented last cycle is valid
                    end
                end

                // ACCUMULATE: each product coeff += e_hat (j==0) or running t_hat.
                // The bank_t write is the combinational bt_we/bt_wa/bt_wd assigns
                // below; here we only advance the accumulate index / (i,j).
                if (pm_vo) begin
                    if (m_oidx == 8'd255) begin
                        // finished this (i,j) term; advance
                        if (m_j + 2'd1 < k_r) begin
                            m_j       <= m_j + 2'd1;
                            m_pending <= 1'b1;     // next term, same row
                        end else begin
                            m_j <= 2'd0;
                            m_i <= m_i + 2'd1;     // next row (or == K -> DONE)
                            if (m_i + 2'd1 < k_r) m_pending <= 1'b1;
                        end
                    end else begin
                        m_oidx <= m_oidx + 8'd1;
                    end
                end

                // Start next (i,j) poly_mul load once core is IDLE again (re-prime)
                if (m_pending && pm_ready && !pm_vo) begin
                    m_loading <= 1'b1;
                    m_ld      <= 9'd0;
                    m_oidx    <= 8'd0;
                    pm_valid  <= 1'b0;
                    m_pending <= 1'b0;
                end
            end

            // Arm Encaps U stage when N finishes: prime first row MAC (i=0,j=0).
            if (st == ST_ENC_N && st_next == ST_ENC_U) begin
                u_row     <= 3'd0;
                u_sub     <= 2'd0;       // MAC
                u_j       <= 3'd0;
                u_ld      <= 9'd0;
                u_oidx    <= 8'd0;
                u_loading <= 1'b1;
                u_pending <= 1'b0;
                pm_valid  <= 1'b0;
            end

            // ---- ST_ENC_U/V sub-phase 0: MAC sum_j (A^T or t_hat)[j] o y_hat[j] ----
            // psum (init 0 at j==0) accumulates into bank_t[UPSUM]. Mirrors ST_M
            // load/accumulate cadence (read-ahead by 1, j-select via u_j0q).
            // U: A^T[u_row][j] (bank_a u_j*K+u_row). V: t_hat[j] (bank_a u_j*K, u_row=0).
            if ((st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && u_sub == 2'd0) begin
                u_j0q <= (u_j == 3'd0);

                if (u_loading) begin
                    if (u_ld == 9'd256) begin
                        pm_valid  <= 1'b0;
                        u_loading <= 1'b0;
                        u_ld      <= 9'd0;
                        u_oidx    <= 8'd0;
                    end else begin
                        u_ld     <= u_ld + 9'd1;
                        pm_valid <= 1'b1;
                    end
                end

                if (pm_vo) begin
                    // bank_t psum write is combinational (u_psum_we); advance idx
                    if (u_oidx == 8'd255) begin
                        if (u_j + 3'd1 < k_r) begin
                            u_j       <= u_j + 3'd1;
                            u_pending <= 1'b1;       // next term, same row
                        end else begin
                            // row's MAC done -> INTT sub-phase
                            u_sub      <= 2'd1;
                            u_ridx     <= 9'd0;
                            u_widx     <= 8'd0;
                            u_nvalid   <= 1'b0;
                            u_nloading <= 1'b1;
                        end
                    end else begin
                        u_oidx <= u_oidx + 8'd1;
                    end
                end

                // re-prime next term's poly_mul load
                if (u_pending && pm_ready && !pm_vo) begin
                    u_loading <= 1'b1;
                    u_ld      <= 9'd0;
                    u_oidx    <= 8'd0;
                    pm_valid  <= 1'b0;
                    u_pending <= 1'b0;
                end
            end

            // ---- ST_ENC_U/V sub-phase 1: INTT(psum) mode=1 in place (bank_t) ----
            if ((st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && u_sub == 2'd1) begin
                if (u_nloading) begin
                    if (u_ridx == 9'd256) begin
                        u_nloading <= 1'b0;
                        u_nvalid   <= 1'b0;
                    end else begin
                        u_ridx   <= u_ridx + 9'd1;
                        u_nvalid <= 1'b1;
                    end
                end
                if (ntt_vo) u_widx <= u_widx + 8'd1;   // in-place write (u_intt_we)
                if (ntt_done) begin
                    // INTT done -> ADD sub-phase (psum + e1[u_row])
                    u_sub    <= 2'd2;
                    u_aidx   <= 9'd0;
                    u_awidx  <= 8'd0;
                    u_avalid <= 1'b0;
                end
            end

            // ---- ST_ENC_U/V + ST_DEC_W sub-phase 2: ADD/SUB + writeback ----
            // U: u[u_row] = psum + e1[u_row] mod Q -> bank_se (u_add_we).
            // V: v       = psum + e2 + mu     mod Q -> bank_t[UPSUM] (u_v_we).
            // D3 w     = v' - psum            mod Q -> bank_t[UPSUM] (u_w_we).
            // Read psum (bank_t) + e1/e2/v' at u_aidx; both arrive 1 cyc later, so
            // register (valid,widx) and write next cycle. u_row_max bounds the
            // row loop (K for U, 1 for V/D3).
            if ((st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && u_sub == 2'd2) begin
                if (u_aidx < 9'd256) begin
                    u_aidx   <= u_aidx + 9'd1;
                    u_avalid <= 1'b1;            // addr presented this cyc -> write next
                    u_awidx  <= u_aidx[7:0];
                end else begin
                    u_avalid <= 1'b0;            // no more addresses
                end
                // writeback (u_add_we / u_v_we) commits index u_awidx this cycle.
                // Advance row after the 256th coeff (u_awidx==255) is written.
                if (u_avalid && u_awidx == 8'd255) begin
                    if (u_row + 3'd1 < u_row_max) begin
                        u_row     <= u_row + 3'd1;
                        u_sub     <= 2'd0;       // next row MAC
                        u_j       <= 3'd0;
                        u_ld      <= 9'd0;
                        u_oidx    <= 8'd0;
                        u_loading <= 1'b1;
                        u_pending <= 1'b0;
                        pm_valid  <= 1'b0;
                    end else begin
                        u_row <= u_row + 3'd1;   // U: ==K ->C1; V: ==1 ->DONE
                    end
                end
            end

            // Arm E5 (ST_ENC_C1) when U finishes: c1 = byteEncode_du(Compress_du(u)).
            // Walk K polys * 256 coeffs; reset bit-packer + ct write pointer.
            if (st == ST_ENC_U && st_next == ST_ENC_C1) begin
                cp_poly  <= 3'd0;
                cp_idx   <= 8'd0;
                cp_ph    <= 3'd0;
                cp_buf   <= 25'd0;
                cp_nbits <= 6'd0;
                cp_wa    <= 12'd0;
                cp_done  <= 1'b0;
                cd_valid <= 1'b0;
            end

            // ---- ST_ENC_C1/C2: Compress_d -> byteEncode_d -> ct region ----
            // C1 (E5): Compress_du(u[0..K-1]) from bank_se -> ct[0..c1_bytes).
            // C2 (E7): Compress_dv(v) from bank_t[UPSUM] -> ct[c1_bytes..ct_bytes).
            // Per coeff, 5-phase micro-sequence (read-ahead 1 cyc bram + 1 cyc
            // comp_decomp pipe), then a drain sub-phase emitting whole bytes:
            //   ph0: present coeff addr (cp_se_full / cp_bt_full by state).
            //   ph1: coeff arrives (cp_coeff_src) -> latch into cd_coeff, pulse cd_valid.
            //   ph2: drop cd_valid (1-cyc pulse); comp_decomp captures.
            //   ph3: cd_vo high -> cd_out (low cp_d bits) valid; append LSB-first.
            //   ph4: drain: while >=8 bits buffered, emit one ct byte/cycle; then advance.
            // Each poly = 256 coeffs = 32*d bytes (whole), so cp_buf empties at
            // each poly boundary (no carry across polys / regions).
            if ((st == ST_ENC_C1 || st == ST_ENC_C2) && !cp_done) begin
                case (cp_ph)
                    3'd0: cp_ph <= 3'd1;                       // addr presented; wait read
                    3'd1: begin
                        cd_coeff <= cp_coeff_src;             // u (bank_se) or v (bank_t)
                        cd_valid <= 1'b1;                     // feed comp_decomp (1-cyc pulse)
                        cp_ph    <= 3'd2;
                    end
                    3'd2: begin
                        cd_valid <= 1'b0;                     // comp_decomp captured this cyc
                        cp_ph    <= 3'd3;
                    end
                    3'd3: begin
                        // cd_out valid (cd_vo): append cp_d bits LSB-first at bit cp_nbits
                        cp_buf   <= cp_buf | (({13'd0, cd_out} & ((25'd1 << cp_d) - 25'd1)) << cp_nbits);
                        cp_nbits <= cp_nbits + {1'b0, cp_d};
                        cp_ph    <= 3'd4;
                    end
                    default: begin                            // 3'd4: drain whole bytes
                        if (cp_nbits >= 6'd8) begin
                            ct_we    <= 1'b1;
                            ct_wa    <= cp_wa;
                            ct_wd    <= cp_buf[7:0];
                            cp_wa    <= cp_wa + 12'd1;
                            cp_buf   <= cp_buf >> 8;
                            cp_nbits <= cp_nbits - 6'd8;
                        end else begin
                            // coeff fully packed; advance coeff / poly
                            if (cp_idx == 8'd255) begin
                                cp_idx <= 8'd0;
                                if (cp_poly + 3'd1 < cp_poly_max) cp_poly <= cp_poly + 3'd1;
                                else                              cp_done <= 1'b1;  // region complete
                            end else begin
                                cp_idx <= cp_idx + 8'd1;
                            end
                            cp_ph <= 3'd0;
                        end
                    end
                endcase
            end

            // Arm E7 (ST_ENC_C2) when V finishes: c2 = byteEncode_dv(Compress_dv(v)).
            // Single poly (v at bank_t UPSUM); ct write pointer cp_wa CONTINUES
            // from c1_bytes (NOT reset) so c2 lands right after c1. cp_buf/nbits
            // are 0 here (c1 ended on a poly/byte boundary) but reset for safety.
            if (st == ST_ENC_V && st_next == ST_ENC_C2) begin
                cp_poly  <= 3'd0;
                cp_idx   <= 8'd0;
                cp_ph    <= 3'd0;
                cp_buf   <= 25'd0;
                cp_nbits <= 6'd0;
                cp_done  <= 1'b0;
                cd_valid <= 1'b0;
                // cp_wa intentionally preserved (= c1_bytes_rt from C1).
            end

            // Arm D1 DECOMP when load settles: decode c1 (u') first.
            if (st == ST_DEC_LOAD && st_next == ST_DEC_DECOMP) begin
                dec_poly     <= 3'd0;
                dec_cidx     <= 8'd0;
                dec_widx     <= 8'd0;
                dec_ph       <= 2'd0;
                dec_ba       <= 12'd0;
                dec_buf      <= 25'd0;
                dec_nbits    <= 6'd0;
                dec_in_c2    <= 1'b0;
                dec_dc_valid <= 1'b0;
                dec_dc_done  <= 1'b0;
                cin_rd_addr_r<= 11'd0;     // present byte 0 address
            end

            // ---- ST_DEC_DECOMP (D1): byteDecode_d + Decompress c1/c2 ----
            // Reverse of the C1/C2 packer. Per coeff:
            //   ph0: if nbits<d, need a byte (addr already presented) -> ph1;
            //        else extract low d bits -> comp_decomp (decompress) -> ph2.
            //   ph1: append cin_rd_data (8 bits) at dec_nbits; advance dec_ba;
            //        present next byte addr; back to ph0.
            //   ph2: drop dc_valid (pulse captured) -> ph3.
            //   ph3: cd_vo -> dec_u_we/dec_v_we writes cd_out; advance coeff.
            // c1 = K polys (d=du) -> u'[poly] bank_se; then c2 = 1 poly (d=dv)
            // -> v' bank_a[DEC_VASLOT]. Byte stream is contiguous in c_in_bram.
            if (st == ST_DEC_DECOMP && !dec_dc_done) begin
                case (dec_ph)
                    2'd0: begin
                        if (dec_nbits >= {1'b0, dec_dc_d}) begin
                            // enough bits: extract low d as the decoded symbol
                            dec_dc_coeff <= dec_buf[11:0] & ((12'd1 << dec_dc_d) - 12'd1);
                            dec_dc_valid <= 1'b1;
                            dec_buf      <= dec_buf >> dec_dc_d;
                            dec_nbits    <= dec_nbits - {1'b0, dec_dc_d};
                            dec_ph       <= 2'd2;
                        end else begin
                            // need another byte (cin_rd_data reflects dec_ba)
                            dec_ph <= 2'd1;
                        end
                    end
                    2'd1: begin
                        dec_buf      <= dec_buf | ({17'd0, cin_rd_data} << dec_nbits);
                        dec_nbits    <= dec_nbits + 6'd8;
                        dec_ba       <= dec_ba + 12'd1;
                        cin_rd_addr_r<= dec_ba[10:0] + 11'd1;   // present next byte
                        dec_ph       <= 2'd0;
                    end
                    2'd2: begin
                        dec_dc_valid <= 1'b0;                   // comp_decomp captured
                        dec_ph       <= 2'd3;
                    end
                    default: begin                              // 2'd3: capture cd_vo
                        if (cd_vo) begin
                            // dec_u_we/dec_v_we commit cd_out at dec_widx this cycle.
                            dec_widx <= dec_widx + 8'd1;
                            if (dec_cidx == 8'd255) begin
                                dec_cidx <= 8'd0;
                                dec_widx <= 8'd0;
                                if (!dec_in_c2) begin
                                    if (dec_poly + 3'd1 < k_r) begin
                                        dec_poly <= dec_poly + 3'd1;   // next u' row
                                    end else begin
                                        dec_in_c2 <= 1'b1;             // switch to c2 (v')
                                        dec_poly  <= 3'd0;
                                    end
                                end else begin
                                    dec_dc_done <= 1'b1;               // v' done -> D1 complete
                                end
                            end else begin
                                dec_cidx <= dec_cidx + 8'd1;
                            end
                            dec_ph <= 2'd0;
                        end
                    end
                endcase
            end

            // Arm E stage when M finishes
            if (st == ST_M && st_next == ST_E) begin
                e_poly <= 3'd0;
                e_pair <= 8'd0;
                e_ph   <= 2'd0;     // start at fetch-c0 of first pair
                e_rho  <= 10'd0;
                e_done <= 1'b0;
            end

            // ---- ST_E: byteEncode12 t_hat -> ek, s_hat -> dk_pke ----
            // Single registered bank read per cycle; 4-cycle micro-phase per
            // coeff pair (ph0 fetch c0, ph1 write b0 + fetch c1, ph2 write b1,
            // ph3 write b2). The coeff for the addr presented last cycle is
            // e_rd_coeff (= bse_rd_data for dk-half, bt_rd_data for ek-half),
            // both registered inside their sd_bram (bse_rd_addr/bt_rd_addr =
            // e_rd_full in ST_E).
            if (st == ST_E && !e_done) begin
                if (e_poly < {1'b0, k_r, 1'b0}) begin
                    // ph0: fetch only (prime). ph1..3: write one packed byte.
                    if (e_ph != 2'd0) begin
                        if (!e_is_dk) begin
                            ek_we <= 1'b1;
                            ek_wa <= e_boff[10:0] + {9'd0, e_wb};
                            ek_wd <= e_wbyte;
                        end else begin
                            dkp_we <= 1'b1;
                            dkp_wa <= e_boff[10:0] + {9'd0, e_wb};
                            dkp_wd <= e_wbyte;
                        end
                    end
                    // save coeff high nibbles as they become available
                    if (e_ph == 2'd1) e_c0_hi <= e_rd_coeff[11:8];  // c0[11:8]
                    if (e_ph == 2'd2) e_c1_hi <= e_rd_coeff[11:4];  // c1[11:4]

                    // advance micro-phase / pair / poly
                    if (e_ph == 2'd3) begin
                        e_ph <= 2'd0;
                        if (e_pair == 8'd127) begin
                            e_pair <= 8'd0;
                            e_poly <= e_poly + 5'd1;   // next poly or -> rho phase
                        end else begin
                            e_pair <= e_pair + 8'd1;
                        end
                    end else begin
                        e_ph <= e_ph + 2'd1;
                    end
                end else begin
                    // rho copy: ek[384*K + r] = rho byte r  (r = 0..31), 1 byte/cycle
                    ek_we <= 1'b1;
                    ek_wa <= dk_bytes_rt[10:0] + {1'b0, e_rho};
                    ek_wd <= rho_r[e_rho*8 +: 8];
                    if (e_rho == 10'd31) e_done <= 1'b1;
                    else                 e_rho  <= e_rho + 10'd1;
                end
            end

            // Arm H stage when E finishes
            if (st == ST_E && st_next == ST_H) begin
                h_blk    <= 3'd0;
                h_byte   <= 8'd0;
                h_phase  <= 2'd0;     // assemble
                h_mbvalid<= 1'b0;
                h_mblast <= 1'b0;
                h_ack    <= 1'b1;     // ready to consume final digest
                h_wb_vld <= 1'b0;     // no pending writeback yet
            end

            // ---- ST_H / ST_ENC_H: H(ek) via multi-block SHA3-256 ----
            // Same logic for KeyGen (ek just written by ST_E) and Encaps (ek
            // preloaded into ek_bram by the TB). hek_r captures H(ek) either way.
            if (st == ST_H || st == ST_ENC_H) begin
                case (h_phase)
                    // assemble 136 bytes of block h_blk into h_block_r.
                    // ek BRAM read is registered: present addr for h_byte this
                    // cycle, write back the byte that arrived for the addr we
                    // presented last cycle (h_wb_*). h_byte runs 0..136 (one
                    // extra cycle to flush the final writeback).
                    2'd0: begin
                        // writeback the byte read for the previous address
                        if (h_wb_vld)
                            h_block_r[h_wb_idx*8 +: 8] <= h_wb_inek ? ek_rd_data : h_wb_pad;

                        if (h_byte <= 8'd135) begin
                            // set up writeback for the address presented this cycle
                            h_wb_vld  <= 1'b1;
                            h_wb_idx  <= h_byte;
                            h_wb_g    <= h_g_addr;
                            h_wb_inek <= (h_g_addr < ek_bytes_rt);
                            h_wb_pad  <= h_padconst(h_blk, h_byte);
                            h_byte    <= h_byte + 8'd1;
                        end else begin
                            // h_byte==136: final writeback (for byte 135) done above
                            h_wb_vld  <= 1'b0;
                            h_byte    <= 8'd0;
                            h_mbvalid <= 1'b1;
                            h_mblast  <= (h_blk == h_nblk_rt - 4'd1);
                            h_phase   <= 2'd1;        // feed
                        end
                    end
                    // feed: hold valid until accepted (mb_ready drops)
                    2'd1: begin
                        if (h_mbvalid && !h_mbready) begin
                            h_mbvalid <= 1'b0;
                            h_mblast  <= 1'b0;
                            h_phase   <= 2'd2;        // wait permute
                        end
                    end
                    // wait permute done: ready again (more blocks) or digest valid (last)
                    2'd2: begin
                        if (h_vo) begin
                            hek_r   <= h_hash[255:0];
                            h_phase <= 2'd3;          // done
                        end else if (h_mbready) begin
                            h_blk   <= h_blk + 3'd1;
                            h_phase <= 2'd0;          // assemble next block
                        end
                    end
                    default: ; // 2'd3 done: hold
                endcase
            end

            // Arm Encaps G when H(ek) finishes: fire the 64-byte single-block
            // G(m||H(ek)). hek_r now holds H(ek); enc_g_data = {hek_r, m_r}.
            if (st == ST_ENC_H && st_next == ST_ENC_G) begin
                sha3_valid <= 1'b1;
                sha3_ack   <= 1'b1;
            end
            // Capture (K, r) when Encaps G completes. ss = K = low half, r = high.
            if (st == ST_ENC_G && sha3_vo) begin
                ss_r     <= sha3_hash[255:0];     // K (shared secret) = G bytes 0..31
                r_r      <= sha3_hash[511:256];   // r = G bytes 32..63
                sha3_ack <= 1'b0;
            end

            // Arm Decaps D5 G when m' is ready (ST_DEC_MENC -> ST_DEC_G): fire the
            // 64-byte single-block G(m'||h). dec_g_data = {hek_r, m_r}.
            if (st == ST_DEC_MENC && st_next == ST_DEC_G) begin
                sha3_valid <= 1'b1;
                sha3_ack   <= 1'b1;
            end
            // Capture (K', r') when D5 G completes. K' = low half (candidate ss),
            // r' = high half (PRF seed for D6 re-encrypt). Reuse ss_r / r_r.
            if (st == ST_DEC_G && sha3_vo) begin
                ss_r     <= sha3_hash[255:0];     // K' = G bytes 0..31
                r_r      <= sha3_hash[511:256];   // r' = G bytes 32..63
                sha3_ack <= 1'b0;
            end

            // Arm D5 J(z||c) when G completes (ST_DEC_G -> ST_DEC_J): assemble the
            // first 136-byte block. dj_ack high to consume the final digest.
            if (st == ST_DEC_G && st_next == ST_DEC_J) begin
                dj_blk    <= 4'd0;
                dj_byte   <= 8'd0;
                dj_phase  <= 2'd0;
                dj_mbvalid<= 1'b0;
                dj_mblast <= 1'b0;
                dj_ack    <= 1'b1;
                dj_done   <= 1'b0;
                dj_wb_vld <= 1'b0;
                cin_rd_addr_r <= 11'd0;   // present c byte 0 (g=32 -> c_idx 0)
            end

            // ---- ST_DEC_J: K-bar = J(z||c) multi-block SHAKE256 ----
            // Mirror of the H(ek) multi-block machine. Byte source by global g:
            //   g < 32           -> z_r byte g
            //   32 <= g < msglen -> c_in_bram byte (g-32), registered read
            //   else             -> SHAKE256 pad constant
            // c_in_bram read is registered: present c_idx for dj_byte this cycle,
            // write back the byte that arrived for the addr presented last cycle.
            if (st == ST_DEC_J && !dj_done) begin
                case (dj_phase)
                    2'd0: begin
                        // writeback the byte read for the previous address
                        if (dj_wb_vld) begin
                            if (dj_wb_inz)
                                dj_block_r[dj_wb_idx*8 +: 8] <= z_r[dj_wb_zidx*8 +: 8];
                            else if (dj_wb_inc)
                                dj_block_r[dj_wb_idx*8 +: 8] <= cin_rd_data;
                            else
                                dj_block_r[dj_wb_idx*8 +: 8] <= dj_wb_pad;
                        end
                        if (dj_byte <= 8'd135) begin
                            dj_wb_vld  <= 1'b1;
                            dj_wb_idx  <= dj_byte;
                            dj_wb_inz  <= (dj_g < 12'd32);
                            dj_wb_inc  <= (dj_g >= 12'd32) && (dj_g < dj_msglen);
                            dj_wb_zidx <= dj_g[4:0];
                            dj_wb_pad  <= dj_padconst(dj_g);
                            // c_in addr is presented combinationally (cin_rd_addr
                            // mux uses dj_c_idx during ST_DEC_J); data lands next
                            // cycle, matching this byte's writeback.
                            dj_byte    <= dj_byte + 8'd1;
                        end else begin
                            dj_wb_vld  <= 1'b0;
                            dj_byte    <= 8'd0;
                            dj_mbvalid <= 1'b1;
                            dj_mblast  <= (dj_blk == dj_nblk - 4'd1);
                            dj_phase   <= 2'd1;
                        end
                    end
                    2'd1: begin
                        if (dj_mbvalid && !h_mbready) begin
                            dj_mbvalid <= 1'b0;
                            dj_mblast  <= 1'b0;
                            dj_phase   <= 2'd2;
                        end
                    end
                    2'd2: begin
                        if (sha3_vo) begin
                            kbar_r  <= sha3_hash[255:0];
                            dj_done <= 1'b1;
                            dj_phase<= 2'd3;
                        end else if (h_mbready) begin
                            dj_blk   <= dj_blk + 4'd1;
                            dj_phase <= 2'd0;
                        end
                    end
                    default: ; // done
                endcase
            end

            // Arm D7 compare when re-encrypt finishes (ST_ENC_C2 -> ST_DEC_CMP).
            if (st == ST_ENC_C2 && st_next == ST_DEC_CMP) begin
                cmp_idx  <= 11'd0;
                cmp_ph   <= 1'b0;
                cmp_neq  <= 1'b0;
                cmp_done <= 1'b0;
            end

            // ---- ST_DEC_CMP (D7): constant-time c' == c compare + ss select ----
            // ph0: present cmp_idx to both ct_bram (c') and c_in_bram (c).
            // ph1: registered reads valid -> OR any difference into cmp_neq,
            //      advance. After the last byte, latch dec_reject and ss stays
            //      K' (ss_r) or switches to K-bar (kbar_r) via the ss_o mux.
            if (st == ST_DEC_CMP && !cmp_done) begin
                if (!cmp_ph) begin
                    cmp_ph <= 1'b1;          // addr presented; wait for read
                end else begin
                    cmp_neq <= cmp_neq | cmp_byte_neq;
                    if (cmp_idx == ct_bytes_rt - 12'd1) begin
                        dec_reject <= cmp_neq | cmp_byte_neq;  // include final byte
                        cmp_done   <= 1'b1;
                    end else begin
                        cmp_idx <= cmp_idx + 11'd1;
                    end
                    cmp_ph <= 1'b0;
                end
            end

            if (st_next == ST_ENC_LOAD &&
                (st == ST_ENC_G || st == ST_DEC_J)) begin
                rl_idx  <= 6'd0;
                rl_widx <= 6'd0;
                rl_vld  <= 1'b0;
            end

            // ---- ST_ENC_LOAD: stream 32 rho bytes from ek_bram into rho_r ----
            // ek read is registered (1-cyc latency): the data arriving this cycle
            // is for the address presented last cycle (rl_idx-1). Present addr
            // rl_idx (valid 0..31); write rho_r[rl_idx-1] when rl_idx in 1..32.
            if (st == ST_ENC_LOAD) begin
                if (rl_idx >= 6'd1 && rl_idx <= 6'd32)
                    rho_r[(rl_idx-6'd1)*8 +: 8] <= ek_rd_data;
                rl_idx <= rl_idx + 6'd1;   // exit when rl_idx >= 33
            end

            // Arm ST_ENC_A (regenerate A_hat via SampleNTT, same as ST_A).
            if (st == ST_ENC_LOAD && st_next == ST_ENC_A) begin
                snt_valid <= 1'b1;
                snt_ack   <= 1'b1;
                a_pair    <= 5'd0;
                a_i       <= 3'd0;
                a_j       <= 3'd0;
                a_widx    <= 8'd0;
                a_busy    <= 1'b0;
            end

            // Arm ST_ENC_TDEC (byteDecode12: ek -> t_hat in bank_a). Now follows
            // E5 (ST_ENC_C1) since e2 occupied bank_t during C/N/U and t_hat is
            // only needed for V. (Was armed off ST_ENC_A in the E1 scaffold.)
            if (st == ST_ENC_C1 && st_next == ST_ENC_TDEC) begin
                td_poly <= 3'd0;
                td_trip <= 8'd0;
                td_ph   <= 3'd0;
                td_done <= 1'b0;
            end
            // Arm D2 s_hat decode (ST_DEC_SDEC) when DECOMP finishes: same TDEC
            // machine, byteDecode12 dk_pke -> s_hat in bank_a slot j*K.
            if (st == ST_DEC_DECOMP && st_next == ST_DEC_SDEC) begin
                td_poly <= 3'd0;
                td_trip <= 8'd0;
                td_ph   <= 3'd0;
                td_done <= 1'b0;
            end

            // Encaps: ek -> t_hat. Decaps D2: dk_pke -> s_hat. Same 5-cycle
            // micro-phase per triple; only the byte source (td_byte) differs.
            if ((st == ST_ENC_TDEC || st == ST_DEC_SDEC) && !td_done) begin
                // capture the byte that arrived for the address presented last cycle
                if (td_ph == 3'd1) td_b0 <= td_byte;
                if (td_ph == 3'd2) td_b1 <= td_byte;
                if (td_ph == 3'd3) td_b2 <= td_byte;

                // write decoded coeffs into bank_a at slot td_poly*K (so V's MAC
                // reads t_hat[j] via the same u_aslot=u_j*K addressing, u_row=0).
                if (td_ph == 3'd3) begin
                    td_we <= 1'b1;
                    td_wa <= (({4'd0,td_poly}*{9'd0,k_r})*256 + {td_trip, 1'b0}) & ((1<<PA_AW)-1);  // 2*trip
                    td_wd <= td_c0;
                end
                if (td_ph == 3'd4) begin
                    td_we <= 1'b1;
                    td_wa <= (({4'd0,td_poly}*{9'd0,k_r})*256 + {td_trip, 1'b0} + 1) & ((1<<PA_AW)-1); // 2*trip+1
                    td_wd <= td_c1;
                end

                // advance micro-phase / triple / poly
                if (td_ph == 3'd4) begin
                    td_ph <= 3'd0;
                    if (td_trip == 8'd127) begin
                        td_trip <= 8'd0;
                        if (td_poly + 3'd1 < k_r) td_poly <= td_poly + 3'd1;
                        else                      td_done <= 1'b1;
                    end else begin
                        td_trip <= td_trip + 8'd1;
                    end
                end else begin
                    td_ph <= td_ph + 3'd1;
                end
            end

            // Arm ST_ENC_E2MV when TDEC finishes: copy e2 (bank_t rel 0) into
            // bank_a[E2_ASLOT]. em_ridx is a read-ahead pointer (bank_t read is
            // registered, 1-cyc latency); em_we/em_widx trail it by 1.
            if (st == ST_ENC_TDEC && st_next == ST_ENC_E2MV) begin
                em_ridx <= 9'd0;
                em_widx <= 8'd0;
                em_we   <= 1'b0;
                em_done <= 1'b0;
            end

            // ---- ST_ENC_E2MV: e2 bank_t[0] -> bank_a[E2_ASLOT], 256 coeffs ----
            // Present read addr em_ridx (bank_t rel 0); bank_t read is registered
            // so bt_rd_data next cycle = e2[em_ridx]. Schedule the bank_a write
            // (em_we/em_widx) for that next cycle with em_widx == em_ridx, so the
            // write commits e2[R] into slot[R] (no off-by-one).
            if (st == ST_ENC_E2MV && !em_done) begin
                if (em_ridx <= 9'd255) begin     // presenting a valid e2 addr this cycle
                    em_we   <= 1'b1;             // -> write next cycle (ba_wd = bt_rd_data)
                    em_widx <= em_ridx[7:0];
                end
                if (em_ridx == 9'd256) em_done <= 1'b1;   // last write (e2[255]) issued
                else                   em_ridx <= em_ridx + 9'd1;
            end

            // Arm ST_ENC_V when e2 relocated: prime first (only) row MAC (i=0,j=0).
            // Reuses the u_* MAC/INTT/ADD machine with u_row tied to 0.
            if (st == ST_ENC_E2MV && st_next == ST_ENC_V) begin
                u_row     <= 3'd0;
                u_sub     <= 2'd0;       // MAC
                u_j       <= 3'd0;
                u_ld      <= 9'd0;
                u_oidx    <= 8'd0;
                u_loading <= 1'b1;
                u_pending <= 1'b0;
                pm_valid  <= 1'b0;
            end

            // Arm D3 (ST_DEC_W) when D2 NTT finishes: same u_* MAC/INTT machine,
            // single row (u_row 0..0). MAC s_hat[j] o u_hat[j], INTT, then SUB.
            if (st == ST_DEC_NTT && st_next == ST_DEC_W) begin
                u_row     <= 3'd0;
                u_sub     <= 2'd0;       // MAC
                u_j       <= 3'd0;
                u_ld      <= 9'd0;
                u_oidx    <= 8'd0;
                u_loading <= 1'b1;
                u_pending <= 1'b0;
                pm_valid  <= 1'b0;
            end

            // Arm D4 (ST_DEC_MENC) when D3 w done: byteEncode_1(Compress_1(w)).
            if (st == ST_DEC_W && st_next == ST_DEC_MENC) begin
                men_idx  <= 8'd0;
                men_ph   <= 2'd0;
                men_done <= 1'b0;
            end

            // ---- ST_DEC_MENC (D4): m' = byteEncode_1(Compress_1(w)) ----
            // Compress_1(w)=1 iff 832 < w <= 2496 (Q=3329). Pack 256 bits
            // LSB-first into m_r (bit men_idx). Read w from bank_t[UPSUM].
            //   ph0: present w[men_idx] addr; ph1: bt_rd_data valid -> set bit.
            if (st == ST_DEC_MENC && !men_done) begin
                case (men_ph)
                    2'd0: men_ph <= 2'd1;          // addr presented; wait read
                    default: begin                 // ph1: bt_rd_data = w[men_idx]
                        m_r[men_idx] <= men_w_bit;
                        if (men_idx == 8'd255) men_done <= 1'b1;
                        else                   men_idx  <= men_idx + 8'd1;
                        men_ph <= 2'd0;
                    end
                endcase
            end
        end
    end

endmodule
