rustpq 0.3.0

Pure Rust post-quantum cryptography suite - ML-KEM, ML-DSA, and more
Documentation
use crate::ml_dsa::params::N;
use crate::ml_dsa::poly::Poly;

pub fn polyt1_pack(r: &mut [u8], a: &Poly) {
    for i in 0..N / 4 {
        r[5 * i] = a.coeffs[4 * i] as u8;
        r[5 * i + 1] = ((a.coeffs[4 * i] >> 8) | (a.coeffs[4 * i + 1] << 2)) as u8;
        r[5 * i + 2] = ((a.coeffs[4 * i + 1] >> 6) | (a.coeffs[4 * i + 2] << 4)) as u8;
        r[5 * i + 3] = ((a.coeffs[4 * i + 2] >> 4) | (a.coeffs[4 * i + 3] << 6)) as u8;
        r[5 * i + 4] = (a.coeffs[4 * i + 3] >> 2) as u8;
    }
}

pub fn polyt1_unpack(r: &mut Poly, a: &[u8]) {
    for i in 0..N / 4 {
        r.coeffs[4 * i] = ((a[5 * i] as i32) | ((a[5 * i + 1] as i32) << 8)) & 0x3FF;
        r.coeffs[4 * i + 1] = (((a[5 * i + 1] as i32) >> 2) | ((a[5 * i + 2] as i32) << 6)) & 0x3FF;
        r.coeffs[4 * i + 2] = (((a[5 * i + 2] as i32) >> 4) | ((a[5 * i + 3] as i32) << 4)) & 0x3FF;
        r.coeffs[4 * i + 3] = (((a[5 * i + 3] as i32) >> 6) | ((a[5 * i + 4] as i32) << 2)) & 0x3FF;
    }
}

pub fn polyt0_pack(r: &mut [u8], a: &Poly) {
    const D: i32 = 1 << 12;
    let mut t = [0i32; 8];

    for i in 0..N / 8 {
        t[0] = D - a.coeffs[8 * i];
        t[1] = D - a.coeffs[8 * i + 1];
        t[2] = D - a.coeffs[8 * i + 2];
        t[3] = D - a.coeffs[8 * i + 3];
        t[4] = D - a.coeffs[8 * i + 4];
        t[5] = D - a.coeffs[8 * i + 5];
        t[6] = D - a.coeffs[8 * i + 6];
        t[7] = D - a.coeffs[8 * i + 7];

        r[13 * i] = t[0] as u8;
        r[13 * i + 1] = (t[0] >> 8) as u8;
        r[13 * i + 1] |= (t[1] << 5) as u8;
        r[13 * i + 2] = (t[1] >> 3) as u8;
        r[13 * i + 3] = (t[1] >> 11) as u8;
        r[13 * i + 3] |= (t[2] << 2) as u8;
        r[13 * i + 4] = (t[2] >> 6) as u8;
        r[13 * i + 4] |= (t[3] << 7) as u8;
        r[13 * i + 5] = (t[3] >> 1) as u8;
        r[13 * i + 6] = (t[3] >> 9) as u8;
        r[13 * i + 6] |= (t[4] << 4) as u8;
        r[13 * i + 7] = (t[4] >> 4) as u8;
        r[13 * i + 8] = (t[4] >> 12) as u8;
        r[13 * i + 8] |= (t[5] << 1) as u8;
        r[13 * i + 9] = (t[5] >> 7) as u8;
        r[13 * i + 9] |= (t[6] << 6) as u8;
        r[13 * i + 10] = (t[6] >> 2) as u8;
        r[13 * i + 11] = (t[6] >> 10) as u8;
        r[13 * i + 11] |= (t[7] << 3) as u8;
        r[13 * i + 12] = (t[7] >> 5) as u8;
    }
}

pub fn polyt0_unpack(r: &mut Poly, a: &[u8]) {
    const D: i32 = 1 << 12;

    for i in 0..N / 8 {
        r.coeffs[8 * i] = a[13 * i] as i32;
        r.coeffs[8 * i] |= (a[13 * i + 1] as i32) << 8;
        r.coeffs[8 * i] &= 0x1FFF;

        r.coeffs[8 * i + 1] = (a[13 * i + 1] as i32) >> 5;
        r.coeffs[8 * i + 1] |= (a[13 * i + 2] as i32) << 3;
        r.coeffs[8 * i + 1] |= (a[13 * i + 3] as i32) << 11;
        r.coeffs[8 * i + 1] &= 0x1FFF;

        r.coeffs[8 * i + 2] = (a[13 * i + 3] as i32) >> 2;
        r.coeffs[8 * i + 2] |= (a[13 * i + 4] as i32) << 6;
        r.coeffs[8 * i + 2] &= 0x1FFF;

        r.coeffs[8 * i + 3] = (a[13 * i + 4] as i32) >> 7;
        r.coeffs[8 * i + 3] |= (a[13 * i + 5] as i32) << 1;
        r.coeffs[8 * i + 3] |= (a[13 * i + 6] as i32) << 9;
        r.coeffs[8 * i + 3] &= 0x1FFF;

        r.coeffs[8 * i + 4] = (a[13 * i + 6] as i32) >> 4;
        r.coeffs[8 * i + 4] |= (a[13 * i + 7] as i32) << 4;
        r.coeffs[8 * i + 4] |= (a[13 * i + 8] as i32) << 12;
        r.coeffs[8 * i + 4] &= 0x1FFF;

        r.coeffs[8 * i + 5] = (a[13 * i + 8] as i32) >> 1;
        r.coeffs[8 * i + 5] |= (a[13 * i + 9] as i32) << 7;
        r.coeffs[8 * i + 5] &= 0x1FFF;

        r.coeffs[8 * i + 6] = (a[13 * i + 9] as i32) >> 6;
        r.coeffs[8 * i + 6] |= (a[13 * i + 10] as i32) << 2;
        r.coeffs[8 * i + 6] |= (a[13 * i + 11] as i32) << 10;
        r.coeffs[8 * i + 6] &= 0x1FFF;

        r.coeffs[8 * i + 7] = (a[13 * i + 11] as i32) >> 3;
        r.coeffs[8 * i + 7] |= (a[13 * i + 12] as i32) << 5;
        r.coeffs[8 * i + 7] &= 0x1FFF;

        r.coeffs[8 * i] = D - r.coeffs[8 * i];
        r.coeffs[8 * i + 1] = D - r.coeffs[8 * i + 1];
        r.coeffs[8 * i + 2] = D - r.coeffs[8 * i + 2];
        r.coeffs[8 * i + 3] = D - r.coeffs[8 * i + 3];
        r.coeffs[8 * i + 4] = D - r.coeffs[8 * i + 4];
        r.coeffs[8 * i + 5] = D - r.coeffs[8 * i + 5];
        r.coeffs[8 * i + 6] = D - r.coeffs[8 * i + 6];
        r.coeffs[8 * i + 7] = D - r.coeffs[8 * i + 7];
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_polyt1_pack_unpack() {
        let mut a = Poly::new();
        for i in 0..N {
            a.coeffs[i] = (i % 1024) as i32;
        }

        let mut packed = [0u8; 320];
        polyt1_pack(&mut packed, &a);

        let mut b = Poly::new();
        polyt1_unpack(&mut b, &packed);

        for i in 0..N {
            assert_eq!(a.coeffs[i], b.coeffs[i]);
        }
    }

    #[test]
    fn test_polyt0_pack_unpack() {
        let mut a = Poly::new();
        for i in 0..N {
            a.coeffs[i] = (i as i32 % 8192) - 4095;
        }

        let mut packed = [0u8; 416];
        polyt0_pack(&mut packed, &a);

        let mut b = Poly::new();
        polyt0_unpack(&mut b, &packed);

        for i in 0..N {
            assert_eq!(a.coeffs[i], b.coeffs[i]);
        }
    }
}