feat: add auxialiary, compress and encode functions

Author: lonerapier

src/kem/kyber/auxiliary.rs

@@ -1,14 +1,19 @@
+//! Contains auxiliary cryptographic functions for Kyber KEM.
+//! - [`prf`] - Pseudorandom function
+//! - [`h`], [`g`] - Hash function
+//! - [`Xof`] - Extendable output function
 use sha3::{
   digest::{ExtendableOutput, Update, XofReader},
   Digest, Shake128, Shake256,
 };
 
-pub fn prf<const eta: usize>(s: [u8; 32], b: u8) -> [u8; 64 * eta] {
-  // concat s and b
+pub fn prf<const ETA: usize>(s: &[u8], b: u8) -> [u8; 64 * ETA] {
+  assert!(s.len() == 32);
+
   let mut hasher = Shake256::default();
   hasher.update(&s);
   hasher.update(&[b]);
-  let mut res = [0u8; 64 * eta];
+  let mut res = [0u8; 64 * ETA];
   XofReader::read(&mut hasher.finalize_xof(), &mut res);
   res
 }

src/kem/kyber/compress.rs

@@ -1,4 +1,4 @@
-use super::MlKemField;
+use super::{MlKemField, PolyVec};
 use crate::{
   algebra::Finite,
   polynomial::{Monomial, Polynomial},
@@ -7,57 +7,69 @@ use crate::{
 /// Compresses a number x to a number in the range [0, 2^d) using the formula round((2^d / q) * x)
 /// mod 2^d.
 /// round(a / b) = floor((a + b/2) / b)
-pub fn compress_fieldelement<const D: usize>(x: &MlKemField) -> MlKemField {
+pub fn compress_fieldelement<const d: usize>(x: &MlKemField) -> MlKemField {
   // TODO: Implement using barrett reduction
   let q_half = (MlKemField::ORDER + 1) >> 1;
-  MlKemField::new((((x.value << D) + q_half) / MlKemField::ORDER) % (1 << D))
+  MlKemField::new((((x.value << d) + q_half) / MlKemField::ORDER) % (1 << d))
 }
 
 /// Decompresses a number y to a number in the range [0, q) using the formula round((q / 2^d)) * y.
-pub fn decompress_fieldelement<const D: usize>(y: &MlKemField) -> MlKemField {
-  let d_pow_half = 1 << (D - 1);
+pub fn decompress_fieldelement<const d: usize>(y: &MlKemField) -> MlKemField {
+  let d_pow_half = 1 << (d - 1);
   let quotient = MlKemField::ORDER * y.value + d_pow_half;
-  MlKemField::new(quotient >> D)
+  MlKemField::new(quotient >> d)
 }
 
-pub fn poly_compress<const D: usize>(
+pub fn poly_compress<const D: usize, const d: usize>(
   poly: &Polynomial<Monomial, MlKemField, D>,
-) -> [MlKemField; D] {
+) -> Polynomial<Monomial, MlKemField, D> {
   // TODO: remove unwrap
-  poly
+  let coeffs = poly
     .coefficients
     .iter()
-    .map(compress_fieldelement::<8>)
+    .map(compress_fieldelement::<d>)
     .collect::<Vec<MlKemField>>()
     .try_into()
-    .unwrap()
+    .unwrap();
+
+  Polynomial::<Monomial, MlKemField, D>::new(coeffs)
 }
 
-pub fn poly_decompress<const D: usize>(
+pub fn poly_decompress<const D: usize, const d: usize>(
   poly: &[MlKemField; D],
 ) -> Polynomial<Monomial, MlKemField, D> {
   let mut coefficients = [MlKemField::default(); D];
   for (i, x) in poly.iter().enumerate() {
-    coefficients[i] = decompress_fieldelement::<8>(x);
+    coefficients[i] = decompress_fieldelement::<d>(x);
   }
   Polynomial::<Monomial, MlKemField, D>::new(coefficients)
 }
 
-// pub fn polyvec_compress<const D: usize, const K: usize>(
-//   poly_vec: &PolyVec<Monomial, D, K>,
-// ) -> [[MlKemField; D]; K] {
-//   let mut res = [[MlKemField::default(); D]; K];
+pub fn polyvec_compress<const D: usize, const d: usize, const K: usize>(
+  poly_vec: &PolyVec<Monomial, D, K>,
+) -> PolyVec<Monomial, D, K> {
+  let mut res = Vec::with_capacity(K);
 
-//   for (i, poly) in poly_vec.vec.iter().enumerate() {
-//     res[i] = poly_compress(poly);
-//   }
+  for poly in poly_vec.vec.iter() {
+    res.push(poly_compress::<D, d>(poly));
+  }
 
-//   res
-// }
+  let res = res.try_into().unwrap();
+  PolyVec::new(res)
+}
 
-// pub fn polyvec_decompress<const D: usize, const K: usize>(
+pub fn polyvec_decompress<const D: usize, const d: usize, const K: usize>(
+  poly_vec: &PolyVec<Monomial, D, K>,
+) -> PolyVec<Monomial, D, K> {
+  let mut res = Vec::with_capacity(K);
 
-// )
+  for poly in poly_vec.vec.iter() {
+    res.push(poly_decompress::<D, d>(&poly.coefficients));
+  }
+
+  let res = res.try_into().unwrap();
+  PolyVec::new(res)
+}
 
 #[test]
 fn test_compress_decompress() {

src/kem/kyber/encode.rs

@@ -1,28 +1,42 @@
-use super::MlKemField;
-use crate::algebra::field::Field;
+use super::{MlKemField, PolyVec};
+use crate::{
+  algebra::field::Field,
+  polynomial::{Basis, Polynomial},
+};
 
-/// Encodes a field element into a byte array where each field element is represented by D bits.
+/// Encodes a field element into a byte array where each field element is represented by d bits.
 /// Converts the field element into a binary representation and then packs the bits into bytes.
-fn byte_encode<const D: usize>(f: [MlKemField; 256]) -> [u8; 32 * D]
-where [(); 256 * D]: {
-  let mut encoded_bits = [0u8; 256 * D];
+pub fn byte_encode<const d: usize, const D: usize>(f: &[MlKemField; D]) -> Vec<u8> {
+  let mut encoded_bits = Vec::with_capacity(D * d);
 
   for (i, x) in f.iter().enumerate() {
     let mut val = x.value;
-    for j in 0..D {
-      encoded_bits[i * D + j] = (val & 1) as u8;
+    for j in 0..d {
+      encoded_bits[i * d + j] = (val & 1) as u8;
       val >>= 1;
     }
   }
 
-  let mut encoded_bytes = [0u8; 32 * D];
+  let mut encoded_bytes = Vec::with_capacity(D / 8 * d);
   for (i, chunk) in encoded_bits.chunks(8).enumerate() {
     encoded_bytes[i] = chunk.iter().enumerate().fold(0, |acc, (j, &b)| acc | (b << j));
   }
 
   encoded_bytes
 }
 
+pub fn byte_encode_polyvec<B: Basis, const D: usize, const K: usize, const d: usize>(
+  f: PolyVec<B, D, K>,
+) -> Vec<u8> {
+  let mut encoded_bytes = Vec::with_capacity(D / 8 * d * K);
+  for (i, poly) in f.vec.iter().enumerate() {
+    let encoded = byte_encode::<d, D>(&poly.coefficients);
+    encoded_bytes[i * D / 8 * d..(i + 1) * D / 8 * d].copy_from_slice(&encoded);
+  }
+
+  encoded_bytes
+}
+
 /// Encodes a field element into a byte array where each field element is represented by D bits.
 /// Converts the field element into a binary representation and then packs the bits into bytes.
 fn byte_encode_optimized<const D: usize>(f: [MlKemField; 256]) -> [u8; 32 * D]
@@ -51,27 +65,42 @@ where [(); 256 * D]: {
 
 /// Decodes a byte array into a field element where each field element is represented by D bits.
 /// Unpacks the bytes into bits and then converts the bits into a field element.
-fn byte_decode<const D: usize>(encoded_bytes: [u8; 32 * D]) -> [MlKemField; 256]
-where [(); 256 * D]: {
-  let mut encoded_bits = [0u8; 256 * D];
+pub fn byte_decode<const d: usize, const D: usize>(encoded_bytes: &[u8]) -> [MlKemField; D] {
+  let mut encoded_bits = Vec::with_capacity(256 * d);
   for (i, &byte) in encoded_bytes.iter().enumerate() {
     for j in 0..8 {
-      encoded_bits[i * 8 + j] = (byte >> j) & 1;
+      encoded_bits.push((byte >> j) & 1);
     }
   }
 
-  let mut f = [MlKemField::ZERO; 256];
+  let mask: usize = (1 << d) - 1;
+  let mut f = [MlKemField::ZERO; D];
   for (i, chunk) in encoded_bits.chunks(D).enumerate() {
     let mut val = 0;
     for (j, &bit) in chunk.iter().enumerate() {
-      val |= (bit as usize) << j;
+      val |= ((bit as usize) << j) & mask;
     }
     f[i].value = val;
   }
 
   f
 }
 
+pub fn byte_decode_polyvec<B: Basis, const D: usize, const K: usize, const d: usize>(
+  encoded_bytes: &[u8],
+  basis: B,
+) -> PolyVec<B, D, K> {
+  let mut f = Vec::with_capacity(K);
+
+  for bytes in encoded_bytes.chunks(32 * d) {
+    let coeffs = byte_decode::<d, D>(bytes.try_into().unwrap());
+    f.push(Polynomial { coefficients: coeffs, basis: basis.clone() })
+  }
+
+  let f = f.try_into().unwrap();
+  PolyVec::new(f)
+}
+
 #[cfg(test)]
 mod tests {
   use super::*;
@@ -88,15 +117,15 @@ mod tests {
   #[test]
   fn test_byte_encode() {
     let f = generate_test_data();
-    let encoded = byte_encode::<8>(f);
+    let encoded = byte_encode::<8, 256>(&f);
     let encoded_optimized = byte_encode_optimized::<8>(f);
     assert_eq!(encoded, encoded_optimized);
   }
 
   #[bench]
   fn bench_byte_encode(b: &mut test::Bencher) {
     let f = generate_test_data();
-    b.iter(|| byte_encode::<8>(f));
+    b.iter(|| byte_encode::<8, 256>(&f));
   }
 
   #[bench]