feat: add sampling and ntt algorithms

Author: lonerapier

src/kem/kyber/ntt.rs

@@ -1,6 +1,6 @@
 use std::ops::Mul;
 
-use super::{MlKemField, Ntt};
+use super::{MlKemField, Ntt, PolyVec};
 use crate::{
   algebra::{field::Field, Finite},
   polynomial::{Monomial, Polynomial},
@@ -71,21 +71,38 @@ impl<const D: usize> Polynomial<Monomial, MlKemField, D> {
   }
 }
 
-impl Mul<Polynomial<Ntt, MlKemField, 256>> for Polynomial<Ntt, MlKemField, 256> {
-  type Output = Self;
+impl<const D: usize, const K: usize> PolyVec<Monomial, D, K> {
+  pub fn ntt(self) -> PolyVec<Ntt, D, K> {
+    let ntt_vec = self.vec.iter().map(|poly| poly.ntt()).collect::<Vec<_>>().try_into().unwrap();
 
-  fn mul(self, rhs: Polynomial<Ntt, MlKemField, 256>) -> Self::Output {
-    let mut res_coeffs = [MlKemField::ZERO; 256];
+    PolyVec::<Ntt, D, K> { vec: ntt_vec }
+  }
+}
+
+impl<const D: usize, const K: usize> PolyVec<Ntt, D, K> {
+  pub fn ntt_inv(self) -> PolyVec<Monomial, D, K> {
+    let ntt_inv_vec =
+      self.vec.iter().map(|poly| poly.ntt_inv()).collect::<Vec<_>>().try_into().unwrap();
+
+    PolyVec { vec: ntt_inv_vec }
+  }
+}
+
+impl<const D: usize> Mul<&Polynomial<Ntt, MlKemField, D>> for &Polynomial<Ntt, MlKemField, D> {
+  type Output = Polynomial<Ntt, MlKemField, D>;
+
+  fn mul(self, rhs: &Polynomial<Ntt, MlKemField, D>) -> Self::Output {
+    let mut res_coeffs = [MlKemField::ZERO; D];
 
-    for i in 0..128 {
+    for i in 0..D >> 1 {
       let (a0, a1) = (self.coefficients[2 * i], self.coefficients[2 * i + 1]);
       let (b0, b1) = (rhs.coefficients[2 * i], rhs.coefficients[2 * i + 1]);
       let (c0, c1) = (a0 * b0 + GAMMA[i] * a1 * b1, a0 * b1 + a1 * b0);
       res_coeffs[2 * i] = c0;
       res_coeffs[2 * i + 1] = c1;
     }
 
-    Polynomial::<Ntt, MlKemField, 256> { coefficients: res_coeffs, basis: Ntt }
+    Polynomial::<Ntt, MlKemField, D> { coefficients: res_coeffs, basis: Ntt }
   }
 }
 

src/kem/kyber/sampling.rs

@@ -1,13 +1,16 @@
-use sha3::digest::XofReader;
-
 use super::{auxiliary::Xof, MlKemField};
 use crate::algebra::{field::Field, Finite};
 
-pub fn sample_ntt(input: &[u8]) -> [MlKemField; 256] {
-  assert!(input.len() == 34);
+pub fn sample_ntt(rho: &[u8], j: u8, i: u8) -> [MlKemField; 256] {
+  assert!(rho.len() == 32);
+  let mut input = [0u8; 34];
+  input[..32].copy_from_slice(rho);
+  input[32] = j;
+  input[33] = i;
+
   let mut ntt = [MlKemField::ZERO; 256];
 
-  let mut xof = Xof::init().absorb(input);
+  let mut xof = Xof::init().absorb(&input);
   let mut j = 0;
   while j < 256 {
     let mut buf = [0u8; 3];

src/polynomial/arithmetic.rs

@@ -34,6 +34,27 @@ impl<F: FiniteField, const D: usize, const D2: usize> Add<Polynomial<Monomial, F
   }
 }
 
+impl<F: FiniteField, const D: usize, const D2: usize> Add<&Polynomial<Monomial, F, D2>>
+  for &Polynomial<Monomial, F, D>
+{
+  type Output = Polynomial<Monomial, F, D>;
+
+  /// Implements addition of two polynomials by adding their coefficients.
+  /// Note: degree of first operand > deg of second operand.
+  fn add(self, rhs: &Polynomial<Monomial, F, D2>) -> Self::Output {
+    let coefficients = self
+      .coefficients
+      .iter()
+      .zip(rhs.coefficients.iter().chain(std::iter::repeat(&F::ZERO)))
+      .map(|(&a, &b)| a + b)
+      .take(D)
+      .collect::<Vec<F>>()
+      .try_into()
+      .unwrap_or_else(|v: Vec<F>| panic!("Expected a Vec of length {} but it was {}", D, v.len()));
+    Self::Output { coefficients, basis: self.basis }
+  }
+}
+
 impl<F: FiniteField, const D: usize, const D2: usize> AddAssign<Polynomial<Monomial, F, D2>>
   for Polynomial<Monomial, F, D>
 {

src/polynomial/mod.rs

@@ -17,7 +17,7 @@
 //! - Includes Discrete Fourier Transform (DFT) for polynomials in the [`Monomial`] basis to convert
 //!   into the [`Lagrange`] basis via evaluation at the roots of unity.
 
-use std::array;
+use std::{array, fmt::Debug};
 
 use super::*;
 use crate::algebra::field::FiniteField;
@@ -45,7 +45,7 @@ pub struct Polynomial<B: Basis, F: FiniteField, const D: usize> {
 
 /// [`Basis`] trait is used to specify the basis of the polynomial.
 /// The basis can be [`Monomial`] or [`Lagrange`]. This is a type-state pattern for [`Polynomial`].
-pub trait Basis {
+pub trait Basis: Debug + Clone {
   /// The associated data type for the basis.
   type Data;
 }