From a734eb2cad8c10fb94e722ab727aa2d70b966f29 Mon Sep 17 00:00:00 2001
From: FallenSigh <fallensigh@gmail.com>
Date: Mon, 29 Jun 2026 21:19:38 +0800
Subject: [PATCH] feat(dec): Decaps D6 - c' = K-PKE.Encrypt(ek_pke, m', r')

Re-encryption step of the FO transform (FIPS 203 Alg 17 step 8), done by
reusing the ENTIRE Encaps E1-E7 pipeline rather than duplicating it:

- FSM: ST_DEC_J (D5) -> ST_ENC_LOAD, then the existing Encaps chain
  LOAD->A->C->N->U->C1->TDEC->E2MV->V->C2 runs unchanged and writes c' to
  ct_bram. The reuse preconditions are all in place: rho loads from ek_bram's
  ek_pke region (same 384k offset Encaps uses; populated at D0 load via
  dk_ld_ekpke), the CBD seed is r_r (r' from D5), and ek_pke is in ek_bram.
- D4 now packs the recovered message directly into m_r (dropping the separate
  mprime_r register): Encaps V's mu reads m_r[idx] and dbg_mprime_o now aliases
  m_r, so the re-encrypt sees m' with no extra plumbing.
- ST_ENC_LOAD arming generalized to fire when entered from ST_ENC_G (Encaps)
  or ST_DEC_J (Decaps re-encrypt).

The re-encrypt overwrites bank_a/bank_se/bank_t, so the bank-based stage checks
(D1 v', D2 s_hat/u_hat, D3 w) are no longer valid at end-of-run. The dec TB now
verifies the surviving register/BRAM artifacts: dk parse (D0), m' (D4, in m_r),
K'/r'/K-bar (D5), and the 768/1088/1568-byte c' against golden (D6). Earlier
stages remain proven by their per-stage builds and transitively by c'.

Verified: dec D6 K=2/3/4 all cases PASS (c' == golden == valid ciphertext c);
KeyGen + Encaps unregressed.
---
 .claude/plans/decaps_plan.md             |  2 +-
 sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v | 46 +++++++++++++++++++-----
 sync_rtl/top/mlkem_top.v                 | 26 ++++++++------
 3 files changed, 55 insertions(+), 19 deletions(-)

diff --git a/.claude/plans/decaps_plan.md b/.claude/plans/decaps_plan.md
index 2029012..d375e6f 100644
--- a/.claude/plans/decaps_plan.md
+++ b/.claude/plans/decaps_plan.md
@@ -67,7 +67,7 @@
 - **D3 — w = v' − INTT(Σ s∘u_hat)** ✅:复用 Encaps V 机(ST_DEC_W,u_row=0 单多项式)。MAC s_hat[j](bank_a slot j*K)∘ u_hat[j](bank_se rel j)→psum bank_t[UPSUM],与 V MAC 完全同址,免改。INTT 原地。SUB:w = v' − psum,(v'−psum) 负则 +Q。**关键:v'/psum 读口冲突 → D1 把 v' 落到 bank_a slot DEC_VASLOT=1(s_hat 在 j*K,slot 1 恒空 K≥2),SUB 时 psum(bank_t)+v'(bank_a)并行读,正如 V-ADD 并读 psum+e2。** K=2/3/4 w 全过。
 - **D4 — m' = byteEncode₁(Compress₁(w))** ✅:ST_DEC_MENC,逐系数读 bank_t[UPSUM] 的 w,Compress₁(w)=1 iff 832<w≤2496(Q=3329),LSB-first 打包进 mprime_r[255:0],经 dbg_mprime_o 暴露。TB verify_d4 对 32 字节 golden(== KAT 解密的 m'==原始 m)。K=2/3/4 全过,KG/Enc 回归通过。K-PKE.Decrypt 硬件全链路打通。
 - **D5 — G(m'‖h) → (K',r') + J(z‖c) → K̄** ✅:G 复用单块 SHA3-512 路径(ST_DEC_G,mode=11,dec_g_data={hek_r,mprime_r}),输出 K'→ss_r(候选 ss,ss_o),r'→r_r(D6 PRF 种子,dbg_r_o)。J 复用多块吸收口(ST_DEC_J,mb_en=1),仿照 H(ek) 机器组装 136B/块,字节源 g<32 取 z_r、32≤g<msglen 取 c_in_bram、否则 SHAKE256 pad(0x1F 后缀,末字节|0x80,与 H 的 0x06 唯一区别)。mb_* 端口按 ST_DEC_J 在 H/J 间多路选择。K̄→kbar_r(dbg_kbar_o)。**踩坑:c_in_bram 读经 cin_rd_addr_r 寄存器 + BRAM 寄存 = 2 拍延迟,而组装/写回流水只容 1 拍 → z→c 边界首个 c 字节(byte32)读到 X,毒化整个 keccak。改为 ST_DEC_J 时 cin_rd_addr 组合直接取 dj_c_idx(去掉寄存器级),数据次拍到达正好对齐写回。** K=2/3/4(密文 768/1088/1568B → 6/9/12 块)全过;K'==KAT ss(有效密文)验证 FO 正确;KG/Enc 回归通过。
-- **D6 — c' = Encrypt(ek_pke,m',r')**:复用 Encaps E1–E7 写 ct_bram。dbg 对 c'==KAT.ct(有效 ct 时)。
+- **D6 — c' = Encrypt(ek_pke,m',r')** ✅:整条 Encaps E1–E7 流水直接复用——ST_DEC_J 完成后跳 ST_ENC_LOAD(rho 从 ek_bram 的 ek_pke 区载入,与 Encaps 同 offset 384k),经 A→C→N→U→C1→TDEC→E2MV→V→C2 跑完写 ct_bram。复用前提全部就位:r'(CBD 种子)在 r_r、ek_pke 在 ek_bram(D0 load 时 dk_ld_ekpke 写入)。**关键改动:D4 把 m' 直接打包进 m_r(而非独立 mprime_r),因为 Encaps V 的 mu=m_r[idx]、dbg_mprime_o 也改指 m_r;省一个寄存器且让重加密天然读到 m'。** ST_ENC_LOAD 的 arming 扩展为从 ST_ENC_G 或 ST_DEC_J 进入皆触发。**注意:重加密覆盖 bank_a/se/t,故 D1–D3 的 bank 检查在 run 末已失效;TB 改为只校验存活的寄存器/BRAM 工件(D0 解析、m'、K'/r'/K̄、c'),bank 阶段正确性由早期分阶段构建 + c' 传递性保证。** K=2/3/4 c'==golden(==有效密文 c)全过,KG/Enc 回归通过。
 - **D7 — 比较 + 拒绝 mux + 端到端 KAT**:c'==c 常量时间比较,ss=mux。干净 TB:
   - 有效 ct(KAT.ct):ss==KAT.ss(c'==c → K')。
   - 损坏 ct(KAT ct_n / ss_n):ss==KAT.ss_n(c'≠c → K̄)。
diff --git a/sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v b/sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v
index 7106083..4a1a502 100644
--- a/sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v
+++ b/sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v
@@ -88,16 +88,19 @@ module tb_mlkem_dec_katK_xsim;
         start_i=1; @(posedge clk); start_i=0;
         c=0; while(!done_o && c<2000000) begin @(posedge clk); c=c+1; end
         if(!done_o) begin $display("FAIL K=%0d case %0d: timeout", KP, casenum); $finish; end
-        $display("=== Decaps D5 done in %0d cyc ===", c);
+        $display("=== Decaps D6 done in %0d cyc ===", c);
 
-        verify_d0;
-        verify_d1;
-        verify_d2;
-        verify_d3;
-        verify_d4;
+        // D6 re-encrypt clobbers bank_a/bank_se/bank_t, so the bank-based stage
+        // checks (D1 v', D2 s_hat/u_hat, D3 w) are no longer valid at end-of-run;
+        // their correctness was proven on earlier per-stage builds and transitively
+        // by c'. Here we check the surviving register/BRAM artifacts: dk parse (D0),
+        // m'/K'/r'/K-bar (D5), and the re-encrypted ciphertext c' (D6).
+        verify_d0;        // also initializes errors = 0
+        verify_d4;        // m' (now in m_r, survives the re-encrypt)
         verify_d5;
-        if (errors == 0) $display("K=%0d CASE %0d PASS (D5): K'/r' = G(m'||h), K-bar = J(z||c) OK", KP, casenum);
-        else             $display("K=%0d CASE %0d FAIL (D5): %0d total errors", KP, casenum, errors);
+        verify_d6;
+        if (errors == 0) $display("K=%0d CASE %0d PASS (D6): c' = Encrypt(ek,m',r') OK", KP, casenum);
+        else             $display("K=%0d CASE %0d FAIL (D6): %0d total errors", KP, casenum, errors);
         $finish;
     end
 
@@ -334,4 +337,31 @@ module tb_mlkem_dec_katK_xsim;
             errors = errors + be;
         end
     endtask
+
+    // D6: verify c' = K-PKE.Encrypt(ek_pke, m', r') in ct_bram == golden.
+    // ct length = 32*(du*K + dv): K2=768, K3=1088, K4=1568.
+    reg [7:0] cp_b [0:1567];
+    task verify_d6;
+        integer i, be, ctlen;
+        reg [8*100-1:0] fn;
+        reg [7:0] got;
+        begin
+            ctlen = (KP == 2) ? 768 : (KP == 3) ? 1088 : 1568;
+            $sformat(fn, "sync_rtl/top/TB/vectors/decgold/dc_k%0d_c%0d_cprime.hex", KP, casenum);
+            $readmemh(fn, cp_b);
+            be = 0;
+            for (i = 0; i < ctlen; i = i + 1) begin
+                dbg_ct_idx_i = i[10:0];
+                @(posedge clk); @(posedge clk); @(posedge clk);
+                got = dbg_ct_o;
+                if (got !== cp_b[i]) begin
+                    if (be < 6) $display("  c'[%0d] got=%02x exp=%02x", i, got, cp_b[i]);
+                    be = be + 1;
+                end
+            end
+            if (be == 0) $display("  PASS: c' == golden (%0d bytes)", ctlen);
+            else         $display("  FAIL: c' %0d byte mismatches", be);
+            errors = errors + be;
+        end
+    endtask
 endmodule
diff --git a/sync_rtl/top/mlkem_top.v b/sync_rtl/top/mlkem_top.v
index 8cd00c9..dbd63cf 100644
--- a/sync_rtl/top/mlkem_top.v
+++ b/sync_rtl/top/mlkem_top.v
@@ -516,8 +516,9 @@ module mlkem_top #(
     assign dbg_sigma_o = sigma_r;
     assign dbg_r_o     = r_r;
     assign dbg_hek_o   = hek_r;
-    // Decaps taps: m' from D4 (mprime_r); z/h parsed from dk at load.
-    assign dbg_mprime_o = mprime_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
@@ -530,7 +531,7 @@ module mlkem_top #(
     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, mprime_r};          // Decaps D5 G: m' || h, 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;
 
@@ -967,7 +968,8 @@ module mlkem_top #(
     reg  [7:0]  men_idx;        // coeff 0..255
     reg  [1:0]  men_ph;         // micro-phase
     reg         men_done;
-    reg  [255:0] mprime_r;      // recovered message m' (32 bytes, bit-packed LSB-first)
+    // D4 packs the recovered message bits directly into m_r (reused by D6's
+    // V/mu re-encrypt and exposed on dbg_mprime_o).
     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
 
@@ -1126,7 +1128,10 @@ module mlkem_top #(
             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;
-            ST_DEC_J:    if (dj_done) st_next = ST_DONE;
+            // 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;
@@ -1270,7 +1275,6 @@ module mlkem_top #(
             men_idx    <= 8'd0;
             men_ph     <= 2'd0;
             men_done   <= 1'b0;
-            mprime_r   <= 256'd0;
             dj_blk     <= 4'd0;
             dj_byte    <= 8'd0;
             dj_phase   <= 2'd0;
@@ -1946,7 +1950,7 @@ module mlkem_top #(
             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, mprime_r}.
+            // 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;
@@ -2033,7 +2037,9 @@ module mlkem_top #(
             end
 
             // Arm rho-load when entering ST_ENC_LOAD. rho = ek[384k .. 384k+31].
-            if (st == ST_ENC_G && st_next == ST_ENC_LOAD) begin
+            // Entered from Encaps G (ST_ENC_G) or Decaps D6 re-encrypt (ST_DEC_J).
+            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;
@@ -2173,13 +2179,13 @@ module mlkem_top #(
 
             // ---- 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 mprime_r (bit men_idx). Read w from bank_t[UPSUM].
+            // 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]
-                        mprime_r[men_idx] <= men_w_bit;
+                        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;