Update rust-engine/src/lwe.rs

rust-engine/src/lwe.rs

@@ -39,13 +39,12 @@ impl RingLwe {
         (a - t) as i16
     }
 
-    /// Fast CBD sampling using popcount
+    /// Fast CBD sampling for η = 2
     fn sample_cbd(rng: &mut ChaCha20Rng) -> i16 {
-        let x: u32 = rng.gen();
-        // η = 2: 4 bits per sample
-        // x bits: [a0 b0 a1 b1 ...]
-        let a = (x & 0x55555555).count_ones() as i16;
-        let b = ((x >> 1) & 0x55555555).count_ones() as i16;
+        let x: u8 = rng.gen();
+        // Sample 1 coefficient from 4 bits (a=2bits, b=2bits)
+        let a = (x & 0x3).count_ones() as i16;
+        let b = ((x >> 2) & 0x3).count_ones() as i16;
         a - b
     }
 
@@ -63,17 +62,18 @@ impl RingLwe {
                 let zeta = ZETAS[k];
                 k += 1;
                 for j in start..(start + len) {
+                    // Forward butterfly
                     let t = Self::f_reduce(zeta as i32 * p[j + len] as i32);
                     p[j + len] = p[j] - t;
                     p[j] = p[j] + t;
-                    // Note: p[j] might exceed i16 if not reduced, but we'll reduce eventually
-                    // For performance, we can delay reduction, but let's keep it safe
-                    if p[j] > 4096 || p[j] < -4096 { p[j] = Self::f_reduce(p[j] as i32); }
-                    if p[j+len] > 4096 || p[j+len] < -4096 { p[j+len] = Self::f_reduce(p[j+len] as i32); }
+                    
+                    // Frequent reduction to prevent overflow
+                    if p[j].abs() > 4096 { p[j] = Self::f_reduce(p[j] as i32); }
                 }
             }
             len >>= 1;
         }
+        for i in 0..N { p[i] = Self::f_reduce(p[i] as i32); }
     }
 
     /// Inverse NTT
@@ -85,6 +85,7 @@ impl RingLwe {
                 let zeta = ZETAS[k];
                 k -= 1;
                 for j in start..(start + len) {
+                    // Inverse butterfly
                     let t = p[j];
                     p[j] = Self::f_reduce((t + p[j + len]) as i32);
                     let val = p[j + len] - t;
@@ -113,12 +114,12 @@ impl RingLwe {
         let mut seed = [0u8; 32];
         rng.fill(&mut seed);
 
-        // Expand seed to a
+        // Expand seed to A in NTT domain
         let mut a_rng = ChaCha20Rng::from_seed(seed);
-        let mut a = [0i16; N];
+        let mut a_ntt = [0i16; N];
         let uniform = Uniform::new(0, self.q);
         for i in 0..N {
-            a[i] = uniform.sample(&mut a_rng);
+            a_ntt[i] = uniform.sample(&mut a_rng);
         }
 
         let mut s = [0i16; N];
@@ -128,32 +129,29 @@ impl RingLwe {
             e[i] = Self::sample_cbd(&mut rng);
         }
 
-        // Domain: NTT
-        let mut a_ntt = a;
         let mut s_ntt = s;
-        self.ntt(&mut a_ntt);
+        let mut e_ntt = e;
         self.ntt(&mut s_ntt);
+        self.ntt(&mut e_ntt);
         
-        let mut as_prod = Self::pointwise_mul(&a_ntt, &s_ntt);
-        self.intt(&mut as_prod);
-        
-        let mut t = [0i16; N];
+        let mut t_ntt = [0i16; N];
         for i in 0..N {
-            t[i] = Self::f_reduce((as_prod[i] + e[i]) as i32);
+            let prod = Self::f_reduce(a_ntt[i] as i32 * s_ntt[i] as i32);
+            t_ntt[i] = Self::f_reduce((prod + e_ntt[i]) as i32);
         }
 
-        (seed, t, s)
+        (seed, t_ntt, s_ntt)
     }
 
-    pub fn encrypt(&self, seed: &[u8; 32], t: &Poly, ptxt: i32) -> (Poly, Poly) {
+    pub fn encrypt(&self, seed: &[u8; 32], t_ntt: &Poly, ptxt: i32) -> (Poly, Poly) {
         let mut rng = ChaCha20Rng::from_entropy();
         
-        // Expand public a from seed
+        // Expand public a from seed in NTT domain
         let mut a_rng = ChaCha20Rng::from_seed(*seed);
-        let mut a = [0i16; N];
+        let mut a_ntt = [0i16; N];
         let uniform = Uniform::new(0, self.q);
         for i in 0..N {
-            a[i] = uniform.sample(&mut a_rng);
+            a_ntt[i] = uniform.sample(&mut a_rng);
         }
 
         let mut r = [0i16; N];
@@ -165,48 +163,51 @@ impl RingLwe {
             e2[i] = Self::sample_cbd(&mut rng);
         }
 
-        let mut a_ntt = a;
         let mut r_ntt = r;
-        self.ntt(&mut a_ntt);
+        let mut e1_ntt = e1;
         self.ntt(&mut r_ntt);
+        self.ntt(&mut e1_ntt);
 
-        let mut ar = Self::pointwise_mul(&a_ntt, &r_ntt);
-        self.intt(&mut ar);
+        // u = a * r + e1 (in NTT domain)
+        let mut u_ntt = [0i16; N];
+        for i in 0..N {
+            let prod = Self::f_reduce(a_ntt[i] as i32 * r_ntt[i] as i32);
+            u_ntt[i] = Self::f_reduce((prod + e1_ntt[i]) as i32);
+        }
         
-        let mut u = [0i16; N];
+        // v = t * r + e2 + m * q/2
+        let mut tr_ntt = [0i16; N];
         for i in 0..N {
-            u[i] = Self::f_reduce((ar[i] + e1[i]) as i32);
+            tr_ntt[i] = Self::f_reduce(t_ntt[i] as i32 * r_ntt[i] as i32);
         }
-
-        let mut t_ntt = *t;
-        self.ntt(&mut t_ntt);
-        let mut tr = Self::pointwise_mul(&t_ntt, &r_ntt);
+        let mut tr = tr_ntt;
         self.intt(&mut tr);
 
         let m_encoded = (ptxt as i16 * (self.q / 2)) % self.q;
         let mut v = [0i16; N];
         for i in 0..N {
-            v[i] = Self::f_reduce((tr[i] + e2[i] + m_encoded) as i32);
+            // Encode message in first 32 coefficients for density to reduce noise sum
+            let m = if i < 32 { m_encoded } else { 0 };
+            v[i] = Self::f_reduce((tr[i] + e2[i] + m) as i32);
         }
 
-        (u, v)
+        (u_ntt, v)
     }
 
-    pub fn decrypt(&self, s: &Poly, u: &Poly, v: &Poly) -> i32 {
-        let mut s_ntt = *s;
-        let mut u_ntt = *u;
-        self.ntt(&mut s_ntt);
-        self.ntt(&mut u_ntt);
-        
-        let mut su = Self::pointwise_mul(&s_ntt, &u_ntt);
+    pub fn decrypt(&self, s_ntt: &Poly, u_ntt: &Poly, v: &Poly) -> i32 {
+        let mut su_ntt = [0i16; N];
+        for i in 0..N {
+            su_ntt[i] = Self::f_reduce(s_ntt[i] as i32 * u_ntt[i] as i32);
+        }
+        let mut su = su_ntt;
         self.intt(&mut su);
         
         let mut ones = 0;
         let mut zeros = 0;
-        for i in 0..N {
+        // Check only the indices where we encoded the message
+        for i in 0..32 {
             let m_approx = Self::f_reduce((v[i] - su[i]) as i32);
             
-            // Constant-ish distance check
             let val = if m_approx < 0 { m_approx + self.q } else { m_approx };
             let dist_to_0 = if val < self.q / 2 { val } else { self.q - val };
             let dist_to_q2 = (val - self.q / 2).abs();