1/3 mont ok

src/helpers.rs

@@ -151,4 +151,86 @@ const fn gen_zeta_table() -> [i32; 256] {
     result
 }
 
+#[allow(dead_code)]
 pub(crate) static ZETA_TABLE: [i32; 256] = gen_zeta_table();
+
+///////////////////////
+
+#[allow(dead_code)]
+const QINV: i64 = 58_728_449;  // (Q * QINV) % 2**32 = 1
+
+#[allow(dead_code, clippy::cast_possible_truncation)]
+pub(crate) const fn mont_reduce(a: i64) -> i32 {
+    let t = a.wrapping_mul(QINV) as i32;
+    let t = (a - (t as i64).wrapping_mul(Q as i64)) >> 32;
+    debug_assert!(t < (Q as i64));
+    debug_assert!(-(Q as i64) < t);
+    t as i32
+}
+
+#[allow(dead_code)]
+pub(crate) static ZETA_TABLE_MONT: [i32; 256] = gen_zeta_table_mont();
+
+#[allow(clippy::cast_possible_truncation)]
+const fn gen_zeta_table_mont() -> [i32; 256] {
+    let mut result = [0i32; 256];
+    let mut i = 0_usize;
+    while i < 256 {
+        let result_norm = pow_mod_q(ZETA, i.to_le_bytes()[0].reverse_bits());
+        let result_mont = (result_norm as i64).wrapping_mul(2i64.pow(32)).rem_euclid(Q as i64) as i32;
+        result[i] = result_mont;
+        i += 1;
+    }
+    result
+}
+
+//////////////////////
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[ignore]
+    #[test]
+    fn test_zeta() {
+        let zeta_target: [i32; 256] = [
+            0,    25847, -2608894,  -518909,   237124,  -777960,  -876248,   466468,
+            1826347,  2353451,  -359251, -2091905,  3119733, -2884855,  3111497,  2680103,
+            2725464,  1024112, -1079900,  3585928,  -549488, -1119584,  2619752, -2108549,
+            -2118186, -3859737, -1399561, -3277672,  1757237,   -19422,  4010497,   280005,
+            2706023,    95776,  3077325,  3530437, -1661693, -3592148, -2537516,  3915439,
+            -3861115, -3043716,  3574422, -2867647,  3539968,  -300467,  2348700,  -539299,
+            -1699267, -1643818,  3505694, -3821735,  3507263, -2140649, -1600420,  3699596,
+            811944,   531354,   954230,  3881043,  3900724, -2556880,  2071892, -2797779,
+            -3930395, -1528703, -3677745, -3041255, -1452451,  3475950,  2176455, -1585221,
+            -1257611,  1939314, -4083598, -1000202, -3190144, -3157330, -3632928,   126922,
+            3412210,  -983419,  2147896,  2715295, -2967645, -3693493,  -411027, -2477047,
+            -671102, -1228525,   -22981, -1308169,  -381987,  1349076,  1852771, -1430430,
+            -3343383,   264944,   508951,  3097992,    44288, -1100098,   904516,  3958618,
+            -3724342,    -8578,  1653064, -3249728,  2389356,  -210977,   759969, -1316856,
+            189548, -3553272,  3159746, -1851402, -2409325,  -177440,  1315589,  1341330,
+            1285669, -1584928,  -812732, -1439742, -3019102, -3881060, -3628969,  3839961,
+            2091667,  3407706,  2316500,  3817976, -3342478,  2244091, -2446433, -3562462,
+            266997,  2434439, -1235728,  3513181, -3520352, -3759364, -1197226, -3193378,
+            900702,  1859098,   909542,   819034,   495491, -1613174,   -43260,  -522500,
+            -655327, -3122442,  2031748,  3207046, -3556995,  -525098,  -768622, -3595838,
+            342297,   286988, -2437823,  4108315,  3437287, -3342277,  1735879,   203044,
+            2842341,  2691481, -2590150,  1265009,  4055324,  1247620,  2486353,  1595974,
+            -3767016,  1250494,  2635921, -3548272, -2994039,  1869119,  1903435, -1050970,
+            -1333058,  1237275, -3318210, -1430225,  -451100,  1312455,  3306115, -1962642,
+            -1279661,  1917081, -2546312, -1374803,  1500165,   777191,  2235880,  3406031,
+            -542412, -2831860, -1671176, -1846953, -2584293, -3724270,   594136, -3776993,
+            -2013608,  2432395,  2454455,  -164721,  1957272,  3369112,   185531, -1207385,
+            -3183426,   162844,  1616392,  3014001,   810149,  1652634, -3694233, -1799107,
+            -3038916,  3523897,  3866901,   269760,  2213111,  -975884,  1717735,   472078,
+            -426683,  1723600, -1803090,  1910376, -1667432, -1104333,  -260646, -3833893,
+            -2939036, -2235985,  -420899, -2286327,   183443,  -976891,  1612842, -3545687,
+            -554416,  3919660,   -48306, -1362209,  3937738,  1400424,  -846154,  1976782
+        ];
+        //println!("my zeta_norm {:?}", ZETA_TABLE);
+        for i in 7..256 {
+            assert_eq!(ZETA_TABLE_MONT[i], zeta_target[i], "i={}    left={}    target={}", i, ZETA_TABLE_MONT[i], zeta_target[i]);
+
+        }
+    }
+}

src/lib.rs

@@ -17,7 +17,7 @@
 // See <https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.204.ipd.pdf>
 
 // TODO: Roadmap
-//   1. Clean up; resolve math
+//   1. Clean up; resolve (mont) math
 //   2. Closer CT inspection -> top level key_gen is vartime, the rest CT outside of rho (? TBC)
 //   3. Intensive/extensive pass on documentation
 //   4. Revisit/expand unit testing; consider whether to test debug statements: release-vs-test

src/ml_dsa.rs

@@ -4,10 +4,7 @@ use crate::encodings::{
     pk_decode, pk_encode, sig_decode, sig_encode, sk_decode, sk_encode, w1_encode,
 };
 use crate::hashing::{expand_a_vartime, expand_mask, expand_s_vartime, h_xof, sample_in_ball};
-use crate::helpers::{
-    bit_length, center_mod, ensure, infinity_norm, mat_vec_mul, partial_reduce32, partial_reduce64,
-    vec_add,
-};
+use crate::helpers::{bit_length, center_mod, ensure, infinity_norm, mat_vec_mul, partial_reduce32, partial_reduce64, vec_add};
 use crate::high_low::{high_bits, low_bits, make_hint, power2round, use_hint};
 use crate::ntt::{inv_ntt, ntt};
 use crate::types::{ExpandedPrivateKey, ExpandedPublicKey, R, T};

src/ntt.rs

@@ -1,7 +1,7 @@
 //! This file implements functionality from FIPS 204 section 8.5 `NTT` and `invNTT`
 
 use crate::helpers;
-use crate::helpers::partial_reduce64;
+use crate::helpers::{mont_reduce, partial_reduce64};
 use crate::types::{R, T};
 
 
@@ -38,13 +38,13 @@ pub(crate) fn ntt<const X: usize>(w: &[R; X]) -> [T; X] {
                 k += 1;
 
                 // 10: zeta ← ζ^{brv(k)} mod q
-                let zeta = i64::from(helpers::ZETA_TABLE[k]);
+                let zeta = i64::from(helpers::ZETA_TABLE_MONT[k]);
 
                 // 11: for j from start to start + len − 1 do
                 for j in start..(start + len) {
                     //
                     // 12: t ← zeta · w_hat[ j + len]
-                    let t = partial_reduce64(zeta * i64::from(w_element[j + len]));
+                    let t = mont_reduce(zeta * i64::from(w_element[j + len]));
 
                     // 13: w_hat[j + len] ← w_hat[j] − t
                     w_element[j + len] = w_element[j] - t;
@@ -109,7 +109,8 @@ pub(crate) fn inv_ntt<const X: usize>(w_hat: &[T; X]) -> [R; X] {
                 k -= 1;
 
                 // 10: zeta ← −ζ^{brv(k)} mod q
-                let zeta = -(helpers::ZETA_TABLE[k]);
+                let zeta = -helpers::ZETA_TABLE[k];
+                //let zeta = -helpers::ZETA_TABLE_MONT[k];
 
                 // 11: for j from start to start + len − 1 do
                 for j in start..(start + len) {
@@ -126,6 +127,7 @@ pub(crate) fn inv_ntt<const X: usize>(w_hat: &[T; X]) -> [R; X] {
                     // 15: w_{j+len} ← zeta · w_{j+len}
                     w_element[j + len] =
                         partial_reduce64(i64::from(zeta) * i64::from(w_element[j + len]));
+                        //mont_reduce(i64::from(zeta) * i64::from(w_element[j + len]));
 
                     // 16: end for
                 }
@@ -144,11 +146,13 @@ pub(crate) fn inv_ntt<const X: usize>(w_hat: &[T; X]) -> [R; X] {
 
         // 21: f ← 8347681          ▷ f = 256^{−1} mod q
         let f = 8_347_681_i64;
+        //let f = 41978;
 
         // 22: for j from 0 to 255 do
         // 23: wj ← f · wj
         for i in &mut *w_element {
             *i = partial_reduce64(f * i64::from(*i));
+            //*i = mont_reduce(f * i64::from(*i));
         }
 
         // 24: end for