From b818911af618d647bdb2f66fcd1f4953adfa6857 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 1 Oct 2026 17:32:53 +0000 Subject: [PATCH 1/2] Spec: ML-DSA's RejNTTPoly four times (vg_mldsa_rej_ntt_poly4) The contract of a function that samples four elements of the matrix A at once, as ML-KEM's vg_mlkem_sample_ntt4 does: for each of the four 34-byte seeds, RejNTTPoly (FIPS 204 Algorithm 30) of it, reduced, or 0 if the loop does not finish within Appendix C's least bound for one of them. Its scratch is that of vg_mldsa_rej_ntt_poly (256 u64s), so callers can pass the same working space. No implementation yet: an x86-64 one, with four SHAKE128 instances in AVX2 registers, follows in its own PR. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr --- lean/VerifiedGarbage/Spec/MlDsa/Poly.lean | 46 +++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean b/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean index c3610ab84..ed63c8ff9 100644 --- a/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean +++ b/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean @@ -182,6 +182,34 @@ def rejNTTContract {M : ISA} (A : Abi M) (stack : Nat := 0) : Contract M := (stack := stack) (leak := some fun seed _a _scratch m => leakBytes (bytesAt m seed 34)) +/-- `vg_mldsa_rej_ntt_poly4(seeds: *const [u8; 136], a: *mut [u32; 1024], scratch: *mut [u64; 256]) -> u32`. -/ +def rejNTT4Sig : Sig where + params := [("seeds", .array false .u8 136), ("a", .array true .u32 1024), + ("scratch", .array true .u64 256)] + ret := some .u32 + +/-- Seed `k` of four at `seeds`: the 34 bytes from byte `34 k`. -/ +def seed4 (m : Mem) (seeds : Addr) (k : Nat) : List Byte := bytesAt m (seeds + BitVec.ofNat 64 (34 * k)) 34 + +/-- Polynomial `k` of four at `a`: from byte `1024 k`. -/ +def poly4 (a : Addr) (k : Nat) : Addr := a + BitVec.ofNat 64 (1024 * k) + +/-- `RejNTTPoly` four times: with the four 34-byte seeds `ρ₀, …, ρ₃` at +`seeds` (`seed4`), writes `RejNTTPoly(ρₖ)` (Algorithm 30) to the polynomial +at `a + 1024 k` (`poly4`), reduced, for each `k`, and returns 1; or returns +0 if the loop of `RejNTTPoly` does not finish within the least bound +Appendix C allows (`minBounds`) for one of them, and `a` is unspecified. +May leak the seeds. -/ +def rejNTT4Contract {M : ISA} (A : Abi M) (stack : Nat := 0) : Contract M := + rejNTT4Sig.contract A + (post := fun seeds a _scratch m m' r => + (r = 1 → ∀ k < 4, Reduced m' (poly4 a k)) ∧ + ((r = 1 ∧ ∀ k < 4, ∃ b : Bounds, rejNTTPoly b.rejNTT (seed4 m seeds k) = some (polyAt m' (poly4 a k))) ∨ + (r = 0 ∧ ∃ k < 4, rejNTTPoly minBounds.rejNTT (seed4 m seeds k) = none))) + (writeArgs := true) + (stack := stack) + (leak := some fun seeds _a _scratch m => leakBytes (bytesAt m seeds 136)) + /-- Whether `RejBoundedPoly` accepts the half-byte `b`: 1 if `CoeffFromHalfByte(b)` (Algorithm 15) is not `⊥`, 0 if it is. -/ def halfByteOk (η b : Nat) : Nat := if (coeffFromHalfByte η b).isSome then 1 else 0 @@ -591,6 +619,24 @@ def rejNTTApi : Api where indices), but not on anything else." safety := [scratchSafety] +/-- `vg_mldsa_rej_ntt_poly4` on every target. -/ +def rejNTT4Api : Api where + module := "mldsa" + name := "vg_mldsa_rej_ntt_poly4" + sig := rejNTT4Sig + writeArgs := true + contracts := some fun A stack => rejNTT4Contract A stack + summary := "`RejNTTPoly` (FIPS 204 Algorithm 30) four times: for each `k` < 4, writes the \ + element of `T_q` sampled from the SHAKE128 output of the 34 bytes of `*seeds` from byte \ + `34 k` to the 256 coefficients of `*a` from coefficient `256 k` (each less than `q` = \ + 8380417), and returns 1. " ++ boundDoc "a" "894 bytes of SHAKE128 output for each" ++ " The \ + four are independent, so an implementation may compute them together (e.g. four SHAKE128 \ + instances at once in vector registers).\n\n\ + Contract: `VG.Spec.MlDsa.rejNTT4Contract`. Not constant time in the seeds: timing may \ + depend on the pointers and on `*seeds` (public in ML-DSA: the seed `ρ` of the matrix and \ + indices), but not on anything else." + safety := [scratchSafety] + /-- `vg_mldsa_rej_bounded_poly` on every target. -/ def rejBoundedApi : Api where module := "mldsa" From 779ca16ebf070c12dfad8f8b5d9538b22483bb49 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 1 Oct 2026 17:33:52 +0000 Subject: [PATCH 2/2] Spec: give vg_mldsa_rej_ntt_poly4 8 KiB of scratch Four interleaved Keccak states, the second buffer of the 4-way permutation and its table of round constants already take 2368 bytes, before the squeezed output; ML-KEM's vg_mlkem_sample_ntt4 has 1024 u64s for the same. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ddof3szoTi7HB8iCsM2MCr --- lean/VerifiedGarbage/Spec/MlDsa/Poly.lean | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean b/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean index ed63c8ff9..002339b9b 100644 --- a/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean +++ b/lean/VerifiedGarbage/Spec/MlDsa/Poly.lean @@ -182,10 +182,10 @@ def rejNTTContract {M : ISA} (A : Abi M) (stack : Nat := 0) : Contract M := (stack := stack) (leak := some fun seed _a _scratch m => leakBytes (bytesAt m seed 34)) -/-- `vg_mldsa_rej_ntt_poly4(seeds: *const [u8; 136], a: *mut [u32; 1024], scratch: *mut [u64; 256]) -> u32`. -/ +/-- `vg_mldsa_rej_ntt_poly4(seeds: *const [u8; 136], a: *mut [u32; 1024], scratch: *mut [u64; 1024]) -> u32`. -/ def rejNTT4Sig : Sig where params := [("seeds", .array false .u8 136), ("a", .array true .u32 1024), - ("scratch", .array true .u64 256)] + ("scratch", .array true .u64 1024)] ret := some .u32 /-- Seed `k` of four at `seeds`: the 34 bytes from byte `34 k`. -/