ml_kem/
compress.rs

1use crate::algebra::{BaseField, Elem, Int, Polynomial, Vector};
2use array::ArraySize;
3use module_lattice::EncodingSize;
4use module_lattice::{Field, Truncate};
5
6// A convenience trait to allow us to associate some constants with a typenum
7pub(crate) trait CompressionFactor: EncodingSize {
8    const POW2_HALF: u32;
9    const MASK: Int;
10    const DIV_SHIFT: usize;
11    const DIV_MUL: u64;
12}
13
14impl<T> CompressionFactor for T
15where
16    T: EncodingSize,
17{
18    const POW2_HALF: u32 = 1 << (T::USIZE - 1);
19    const MASK: Int = (1 << T::USIZE) - 1;
20    const DIV_SHIFT: usize = 34;
21    #[allow(clippy::integer_division_remainder_used, reason = "constant")]
22    const DIV_MUL: u64 = (1 << T::DIV_SHIFT) / BaseField::QLL;
23}
24
25// Traits for objects that allow compression / decompression
26pub(crate) trait Compress {
27    fn compress<D: CompressionFactor>(&mut self) -> &Self;
28    fn decompress<D: CompressionFactor>(&mut self) -> &Self;
29}
30
31impl Compress for Elem {
32    // Equation 4.5: Compress_d(x) = round((2^d / q) x)
33    //
34    // Here and in decompression, we leverage the following facts:
35    //
36    //   round(a / b) = floor((a + b/2) / b)
37    //   a / q ~= (a * x) >> s where x >> s ~= 1/q
38    fn compress<D: CompressionFactor>(&mut self) -> &Self {
39        const Q_HALF: u64 = (BaseField::QLL + 1) >> 1;
40        let x = u64::from(self.0);
41        let y = (((x << D::USIZE) + Q_HALF) * D::DIV_MUL) >> D::DIV_SHIFT;
42        let y: u16 = Truncate::truncate(y);
43        self.0 = y & D::MASK;
44        self
45    }
46
47    // Equation 4.6: Decompress_d(x) = round((q / 2^d) x)
48    fn decompress<D: CompressionFactor>(&mut self) -> &Self {
49        let x = u32::from(self.0);
50        let y = ((x * BaseField::QL) + D::POW2_HALF) >> D::USIZE;
51        self.0 = Truncate::truncate(y);
52        self
53    }
54}
55impl Compress for Polynomial {
56    fn compress<D: CompressionFactor>(&mut self) -> &Self {
57        for x in &mut self.0 {
58            x.compress::<D>();
59        }
60
61        self
62    }
63
64    fn decompress<D: CompressionFactor>(&mut self) -> &Self {
65        for x in &mut self.0 {
66            x.decompress::<D>();
67        }
68
69        self
70    }
71}
72
73impl<K: ArraySize> Compress for Vector<K> {
74    fn compress<D: CompressionFactor>(&mut self) -> &Self {
75        for x in &mut self.0 {
76            x.compress::<D>();
77        }
78
79        self
80    }
81
82    fn decompress<D: CompressionFactor>(&mut self) -> &Self {
83        for x in &mut self.0 {
84            x.decompress::<D>();
85        }
86
87        self
88    }
89}
90
91#[cfg(test)]
92#[allow(clippy::cast_possible_truncation, reason = "tests")]
93#[allow(clippy::integer_division_remainder_used, reason = "tests")]
94pub(crate) mod tests {
95    use super::*;
96    use array::typenum::{U1, U4, U5, U6, U10, U11, U12};
97    use num_rational::Ratio;
98
99    fn rational_compress<D: CompressionFactor>(input: u16) -> u16 {
100        let fraction = Ratio::new(u32::from(input) * (1 << D::USIZE), BaseField::QL);
101        (fraction.round().to_integer() as u16) & D::MASK
102    }
103
104    fn rational_decompress<D: CompressionFactor>(input: u16) -> u16 {
105        let fraction = Ratio::new(u32::from(input) * BaseField::QL, 1 << D::USIZE);
106        fraction.round().to_integer() as u16
107    }
108
109    // Verify against inequality 4.7
110    fn compression_decompression_inequality<D: CompressionFactor>() {
111        const QI32: i32 = BaseField::Q as i32;
112        let error_threshold = i32::from(Ratio::new(BaseField::Q, 1 << D::USIZE).to_integer());
113
114        for x in 0..BaseField::Q {
115            let mut y = Elem::new(x);
116            y.compress::<D>();
117            y.decompress::<D>();
118
119            let mut error = i32::from(y.0) - i32::from(x) + QI32;
120            if error > (QI32 - 1) / 2 {
121                error -= QI32;
122            }
123
124            assert!(
125                error.abs() <= error_threshold,
126                "Inequality failed for x = {x}: error = {}, error_threshold = {error_threshold}, D = {:?}",
127                error.abs(),
128                D::USIZE
129            );
130        }
131    }
132
133    fn decompression_compression_equality<D: CompressionFactor>() {
134        for x in 0..(1 << D::USIZE) {
135            let mut y = Elem::new(x);
136            y.decompress::<D>();
137            y.compress::<D>();
138
139            assert_eq!(y.0, x, "failed for x: {}, D: {}", x, D::USIZE);
140        }
141    }
142
143    fn decompress_KAT<D: CompressionFactor>() {
144        for y in 0..(1 << D::USIZE) {
145            let x_expected = rational_decompress::<D>(y);
146            let mut x_actual = Elem::new(y);
147            x_actual.decompress::<D>();
148
149            assert_eq!(x_expected, x_actual.0);
150        }
151    }
152
153    fn compress_KAT<D: CompressionFactor>() {
154        for x in 0..BaseField::Q {
155            let y_expected = rational_compress::<D>(x);
156            let mut y_actual = Elem::new(x);
157            y_actual.compress::<D>();
158
159            assert_eq!(y_expected, y_actual.0, "for x: {}, D: {}", x, D::USIZE);
160        }
161    }
162
163    fn compress_decompress_properties<D: CompressionFactor>() {
164        compression_decompression_inequality::<D>();
165        decompression_compression_equality::<D>();
166    }
167
168    fn compress_decompress_KATs<D: CompressionFactor>() {
169        decompress_KAT::<D>();
170        compress_KAT::<D>();
171    }
172
173    #[test]
174    fn decompress_compress() {
175        compress_decompress_properties::<U1>();
176        compress_decompress_properties::<U4>();
177        compress_decompress_properties::<U5>();
178        compress_decompress_properties::<U6>();
179        compress_decompress_properties::<U10>();
180        compress_decompress_properties::<U11>();
181        // preservation under decompression first only holds for d < 12
182        compression_decompression_inequality::<U12>();
183
184        compress_decompress_KATs::<U1>();
185        compress_decompress_KATs::<U4>();
186        compress_decompress_KATs::<U5>();
187        compress_decompress_KATs::<U6>();
188        compress_decompress_KATs::<U10>();
189        compress_decompress_KATs::<U11>();
190        compress_decompress_KATs::<U12>();
191    }
192}
ml_kem/compress.rs

ml_kem/
compress.rs
