ml_dsa/
encode.rs

1use core::ops::Add;
2use hybrid_array::{
3    Array,
4    typenum::{Len, Length, Sum, Unsigned},
5};
6use module_lattice::{ArraySize, Encode, EncodingSize, VectorEncodingSize};
7
8use crate::algebra::{Elem, Polynomial, Vector};
9
10/// A pair of integers that describes a range
11pub trait RangeEncodingSize {
12    type Min: Unsigned;
13    type Max: Unsigned;
14    type EncodingSize: EncodingSize;
15}
16
17impl<A, B> RangeEncodingSize for (A, B)
18where
19    A: Unsigned + Add<B>,
20    B: Unsigned,
21    Sum<A, B>: Len,
22    Length<Sum<A, B>>: EncodingSize,
23{
24    type Min = A;
25    type Max = B;
26    type EncodingSize = Length<Sum<A, B>>;
27}
28
29pub(crate) type RangeMin<A, B> = <(A, B) as RangeEncodingSize>::Min;
30pub(crate) type RangeMax<A, B> = <(A, B) as RangeEncodingSize>::Max;
31pub(crate) type RangeEncodingBits<A, B> = <(A, B) as RangeEncodingSize>::EncodingSize;
32pub(crate) type RangeEncodedPolynomialSize<A, B> =
33    <RangeEncodingBits<A, B> as EncodingSize>::EncodedPolynomialSize;
34pub(crate) type RangeEncodedPolynomial<A, B> = Array<u8, RangeEncodedPolynomialSize<A, B>>;
35pub(crate) type RangeEncodedVectorSize<A, B, K> =
36    <RangeEncodingBits<A, B> as VectorEncodingSize<K>>::EncodedVectorSize;
37pub(crate) type RangeEncodedVector<A, B, K> = Array<u8, RangeEncodedVectorSize<A, B, K>>;
38
39/// `BitPack` represents range-encoding logic
40pub(crate) trait BitPack<A, B> {
41    type PackedSize: ArraySize;
42    fn pack(&self) -> Array<u8, Self::PackedSize>;
43    fn unpack(enc: &Array<u8, Self::PackedSize>) -> Self;
44}
45
46impl<A, B> BitPack<A, B> for Polynomial
47where
48    (A, B): RangeEncodingSize,
49{
50    type PackedSize = RangeEncodedPolynomialSize<A, B>;
51
52    // Algorithm 17 BitPack
53    fn pack(&self) -> RangeEncodedPolynomial<A, B> {
54        let a = Elem::new(RangeMin::<A, B>::U32);
55        let b = Elem::new(RangeMax::<A, B>::U32);
56
57        let to_encode = Self::new(
58            self.0
59                .iter()
60                .map(|w| {
61                    assert!(w.0 <= b.0 || w.0 >= (-a).0);
62                    b - *w
63                })
64                .collect(),
65        );
66        Encode::<RangeEncodingBits<A, B>>::encode(&to_encode)
67    }
68
69    // Algorithm 17 BitUnPack
70    fn unpack(enc: &RangeEncodedPolynomial<A, B>) -> Self {
71        let a = Elem::new(RangeMin::<A, B>::U32);
72        let b = Elem::new(RangeMax::<A, B>::U32);
73        let mut decoded: Self = Encode::<RangeEncodingBits<A, B>>::decode(enc);
74
75        for z in &mut decoded.0 {
76            assert!(z.0 <= (a + b).0);
77            *z = b - *z;
78        }
79
80        decoded
81    }
82}
83
84impl<K, A, B> BitPack<A, B> for Vector<K>
85where
86    K: ArraySize,
87    (A, B): RangeEncodingSize,
88    RangeEncodingBits<A, B>: VectorEncodingSize<K>,
89{
90    type PackedSize = RangeEncodedVectorSize<A, B, K>;
91
92    fn pack(&self) -> RangeEncodedVector<A, B, K> {
93        let polys = self.0.iter().map(|x| BitPack::<A, B>::pack(x)).collect();
94        RangeEncodingBits::<A, B>::flatten(polys)
95    }
96
97    fn unpack(enc: &RangeEncodedVector<A, B, K>) -> Self {
98        let unfold = RangeEncodingBits::<A, B>::unflatten(enc);
99        Self(
100            unfold
101                .into_iter()
102                .map(|x| <Polynomial as BitPack<A, B>>::unpack(x))
103                .collect(),
104        )
105    }
106}
107
108#[cfg(test)]
109#[allow(clippy::integer_division_remainder_used, reason = "tests")]
110pub(crate) mod tests {
111    use super::*;
112    use crate::algebra::*;
113    use core::ops::Rem;
114    use getrandom::{
115        SysRng,
116        rand_core::{Rng, UnwrapErr},
117    };
118    use hybrid_array::typenum::{
119        U1, U2, U3, U4, U6, U7, U8, U9, U10, U13, U17, U19,
120        marker_traits::Zero,
121        operator_aliases::{Diff, Mod, Shleft},
122    };
123    use module_lattice::{EncodedPolynomial, Field};
124
125    // A helper trait to construct larger arrays by repeating smaller ones
126    trait Repeat<T: Clone, D: ArraySize> {
127        fn repeat(&self) -> Array<T, D>;
128    }
129
130    impl<T, N, D> Repeat<T, D> for Array<T, N>
131    where
132        N: ArraySize,
133        T: Clone,
134        D: ArraySize + Rem<N>,
135        Mod<D, N>: Zero,
136    {
137        fn repeat(&self) -> Array<T, D> {
138            Array::from_fn(|i| self[i % N::USIZE].clone())
139        }
140    }
141
142    fn simple_bit_pack_test<D>(b: u32, decoded: &Polynomial, encoded: &EncodedPolynomial<D>)
143    where
144        D: EncodingSize,
145    {
146        // Test known answer
147        let actual_encoded = Encode::<D>::encode(decoded);
148        assert_eq!(actual_encoded, *encoded);
149
150        let actual_decoded: Polynomial = Encode::<D>::decode(encoded);
151        assert_eq!(actual_decoded, *decoded);
152
153        // Test random decode/encode and encode/decode round trips
154        let mut rng = UnwrapErr(SysRng);
155        let decoded = Polynomial::new(Array::from_fn(|_| {
156            let x = rng.next_u32();
157            Elem::new(x % (b + 1))
158        }));
159
160        let actual_encoded = Encode::<D>::encode(&decoded);
161        let actual_decoded: Polynomial = Encode::<D>::decode(&actual_encoded);
162        assert_eq!(actual_decoded, decoded);
163
164        let actual_reencoded = Encode::<D>::encode(&decoded);
165        assert_eq!(actual_reencoded, actual_encoded);
166    }
167
168    #[test]
169    fn simple_bit_pack() {
170        // Use a standard test pattern across all the cases
171        let decoded = Polynomial::new(
172            Array::<_, U8>([
173                Elem::new(0),
174                Elem::new(1),
175                Elem::new(2),
176                Elem::new(3),
177                Elem::new(4),
178                Elem::new(5),
179                Elem::new(6),
180                Elem::new(7),
181            ])
182            .repeat(),
183        );
184
185        // 10 bits
186        // <-> b = 2^{bitlen(q-1) - d} - 1 = 2^10 - 1
187        let b = (1 << 10) - 1;
188        let encoded: EncodedPolynomial<U10> =
189            Array::<_, U10>([0x00, 0x04, 0x20, 0xc0, 0x00, 0x04, 0x14, 0x60, 0xc0, 0x01]).repeat();
190        simple_bit_pack_test::<U10>(b, &decoded, &encoded);
191
192        // 8 bits
193        // gamma2 = (q - 1) / 88
194        // b = (q - 1) / (2 gamma2) - 1 = 175 = 2^8 - 81
195        let b = (1 << 8) - 81;
196        let encoded: EncodedPolynomial<U8> =
197            Array::<_, U8>([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]).repeat();
198        simple_bit_pack_test::<U8>(b, &decoded, &encoded);
199
200        // 6 bits
201        // gamma2 = (q - 1) / 32
202        // b = (q - 1) / (2 gamma2) - 1 = 63 = 2^6 - 1
203        let b = (1 << 6) - 1;
204        let encoded: EncodedPolynomial<U6> =
205            Array::<_, U6>([0x40, 0x20, 0x0c, 0x44, 0x61, 0x1c]).repeat();
206        simple_bit_pack_test::<U6>(b, &decoded, &encoded);
207    }
208
209    fn bit_pack_test<A, B>(decoded: &Polynomial, encoded: &RangeEncodedPolynomial<A, B>)
210    where
211        A: Unsigned,
212        B: Unsigned,
213        (A, B): RangeEncodingSize,
214    {
215        let a = Elem::new(A::U32);
216        let b = Elem::new(B::U32);
217
218        // Test known answer
219        let actual_encoded = BitPack::<A, B>::pack(decoded);
220        assert_eq!(actual_encoded, *encoded);
221
222        let actual_decoded: Polynomial = BitPack::<A, B>::unpack(encoded);
223        assert_eq!(actual_decoded, *decoded);
224
225        // Test random decode/encode and encode/decode round trips
226        let mut rng = UnwrapErr(SysRng);
227        let decoded = Polynomial::new(Array::from_fn(|_| {
228            let mut x = rng.next_u32();
229            x %= a.0 + b.0;
230            b - Elem::new(x)
231        }));
232
233        let actual_encoded = BitPack::<A, B>::pack(&decoded);
234        let actual_decoded: Polynomial = BitPack::<A, B>::unpack(&actual_encoded);
235        assert_eq!(actual_decoded, decoded);
236
237        let actual_reencoded = BitPack::<A, B>::pack(&decoded);
238        assert_eq!(actual_reencoded, actual_encoded);
239    }
240
241    #[test]
242    fn bit_pack() {
243        type D = U13;
244        type Pow2D = Shleft<U1, D>;
245        type Pow2DMin = Diff<Pow2D, U1>;
246
247        type Gamma1Lo = Shleft<U1, U17>;
248        type Gamma1LoMin = Diff<Gamma1Lo, U1>;
249
250        type Gamma1Hi = Shleft<U1, U19>;
251        type Gamma1HiMin = Diff<Gamma1Hi, U1>;
252
253        // Use a standard test pattern across all the cases
254        // (We can't use -2 because the eta=2 case doesn't actually cover -2)
255        let decoded = Polynomial::new(
256            Array::<_, U4>([
257                Elem::new(BaseField::Q - 1),
258                Elem::new(0),
259                Elem::new(1),
260                Elem::new(2),
261            ])
262            .repeat(),
263        );
264
265        // BitPack(_, eta, eta), eta = 2, 4
266        let encoded: RangeEncodedPolynomial<U2, U2> = Array::<_, U3>([0x53, 0x30, 0x05]).repeat();
267        bit_pack_test::<U2, U2>(&decoded, &encoded);
268
269        let encoded: RangeEncodedPolynomial<U4, U4> = Array::<_, U2>([0x45, 0x23]).repeat();
270        bit_pack_test::<U4, U4>(&decoded, &encoded);
271
272        // BitPack(_, 2^d - 1, 2^d), d = 13
273        let encoded: RangeEncodedPolynomial<Pow2DMin, Pow2D> =
274            Array::<_, U7>([0x01, 0x20, 0x00, 0xf8, 0xff, 0xf9, 0x7f]).repeat();
275        bit_pack_test::<Pow2DMin, Pow2D>(&decoded, &encoded);
276
277        // BitPack(_, gamma1 - 1, gamma1), gamma1 = 2^17, 2^19
278        let encoded: RangeEncodedPolynomial<Gamma1LoMin, Gamma1Lo> =
279            Array::<_, U9>([0x01, 0x00, 0x02, 0x00, 0xf8, 0xff, 0x9f, 0xff, 0x7f]).repeat();
280        bit_pack_test::<Gamma1LoMin, Gamma1Lo>(&decoded, &encoded);
281
282        let encoded: RangeEncodedPolynomial<Gamma1HiMin, Gamma1Hi> =
283            Array::<_, U10>([0x00, 0x00, 0xf8, 0xff, 0x7f, 0xfe, 0xff, 0xd7, 0xff, 0x7f]).repeat();
284        bit_pack_test::<Gamma1Hi, Gamma1HiMin>(&decoded, &encoded);
285    }
286}
ml_dsa/encode.rs

ml_dsa/
encode.rs
