1use itertools::izip;
  2use math::num_bigint::Div;
  3use math::poly::{Poly, PolyRNS};
  4use math::ring::RingRNS;
  5use num_bigint::BigInt;
  6use sampling::source::Source;
  7
  8#[test]
  9fn rescaling_rns_u64() {
 10    let n = 1 << 10;
 11    let moduli: Vec<u64> = vec![
 12        0x1fffffffffc80001u64,
 13        0x1fffffffffe00001u64,
 14        0x1fffffffffb40001,
 15        0x1fffffffff500001,
 16    ];
 17    let ring_rns: RingRNS<u64> = RingRNS::new(n, moduli);
 18
 19    sub_test("test_div_by_last_modulus::<ROUND:false, NTT:false>", || {
 20        test_div_by_last_modulus::<false, false>(&ring_rns)
 21    });
 22    sub_test("test_div_by_last_modulus::<ROUND:false, NTT:true>", || {
 23        test_div_by_last_modulus::<false, true>(&ring_rns)
 24    });
 25    sub_test("test_div_by_last_modulus::<ROUND:true, NTT:false>", || {
 26        test_div_by_last_modulus::<true, false>(&ring_rns)
 27    });
 28    sub_test("test_div_by_last_modulus::<ROUND:true, NTT:true>", || {
 29        test_div_by_last_modulus::<true, true>(&ring_rns)
 30    });
 31    sub_test(
 32        "test_div_by_last_modulus_inplace::<ROUND:false, NTT:false>",
 33        || test_div_by_last_modulus_inplace::<false, false>(&ring_rns),
 34    );
 35    sub_test(
 36        "test_div_by_last_modulus_inplace::<ROUND:false, NTT:true>",
 37        || test_div_by_last_modulus_inplace::<false, true>(&ring_rns),
 38    );
 39    sub_test(
 40        "test_div_by_last_modulus_inplace::<ROUND:true, NTT:true>",
 41        || test_div_by_last_modulus_inplace::<true, true>(&ring_rns),
 42    );
 43    sub_test(
 44        "test_div_by_last_modulus_inplace::<ROUND:true, NTT:false>",
 45        || test_div_by_last_modulus_inplace::<true, false>(&ring_rns),
 46    );
 47    sub_test("test_div_by_last_moduli::<ROUND:false, NTT:false>", || {
 48        test_div_by_last_moduli::<false, false>(&ring_rns)
 49    });
 50    sub_test("test_div_by_last_moduli::<ROUND:false, NTT:true>", || {
 51        test_div_by_last_moduli::<false, true>(&ring_rns)
 52    });
 53    sub_test("test_div_by_last_moduli::<ROUND:true, NTT:false>", || {
 54        test_div_by_last_moduli::<true, false>(&ring_rns)
 55    });
 56    sub_test("test_div_by_last_moduli::<ROUND:true, NTT:true>", || {
 57        test_div_by_last_moduli::<true, true>(&ring_rns)
 58    });
 59    sub_test(
 60        "test_div_by_last_moduli_inplace::<ROUND:false, NTT:false>",
 61        || test_div_by_last_moduli_inplace::<false, false>(&ring_rns),
 62    );
 63    sub_test(
 64        "test_div_by_last_moduli_inplace::<ROUND:false, NTT:true>",
 65        || test_div_by_last_moduli_inplace::<false, true>(&ring_rns),
 66    );
 67    sub_test(
 68        "test_div_by_last_moduli_inplace::<ROUND:true, NTT:false>",
 69        || test_div_by_last_moduli_inplace::<true, false>(&ring_rns),
 70    );
 71    sub_test(
 72        "test_div_by_last_moduli_inplace::<ROUND:true, NTT:true>",
 73        || test_div_by_last_moduli_inplace::<true, true>(&ring_rns),
 74    );
 75}
 76
 77fn sub_test<F: FnOnce()>(name: &str, f: F) {
 78    println!("Running {}", name);
 79    f();
 80}
 81
 82fn test_div_by_last_modulus<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
 83    let seed: [u8; 32] = [0; 32];
 84    let mut source: Source = Source::new(seed);
 85
 86    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
 87    let mut buf: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
 88    let mut c: PolyRNS<u64> = ring_rns.at_level(ring_rns.level() - 1).new_polyrns();
 89
 90    // Allocates a random PolyRNS
 91    ring_rns.fill_uniform(&mut source, &mut a);
 92
 93    // Maps PolyRNS to [BigInt]
 94    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
 95    ring_rns
 96        .at_level(a.level())
 97        .to_bigint_inplace(&a, 1, &mut coeffs_a);
 98
 99    // Performs c = intt(ntt(a) / q_level)
100    if NTT {
101        ring_rns.ntt_inplace::<false>(&mut a);
102    }
103
104    ring_rns.div_by_last_modulus::<ROUND, NTT>(&a, &mut buf, &mut c);
105
106    if NTT {
107        ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c);
108    }
109
110    // Exports c to coeffs_c
111    let mut coeffs_c = vec![BigInt::from(0); c.n()];
112    ring_rns
113        .at_level(c.level())
114        .to_bigint_inplace(&c, 1, &mut coeffs_c);
115
116    // Performs floor division on a
117    let scalar_big = BigInt::from(ring_rns.0[ring_rns.level()].modulus.q);
118    coeffs_a.iter_mut().for_each(|a| {
119        if ROUND {
120            *a = a.div_round(&scalar_big);
121        } else {
122            *a = a.div_floor(&scalar_big);
123        }
124    });
125
126    izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
127}
128
129fn test_div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
130    let seed: [u8; 32] = [0; 32];
131    let mut source: Source = Source::new(seed);
132
133    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
134    let mut buf: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
135
136    // Allocates a random PolyRNS
137    ring_rns.fill_uniform(&mut source, &mut a);
138
139    // Maps PolyRNS to [BigInt]
140    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
141    ring_rns
142        .at_level(a.level())
143        .to_bigint_inplace(&a, 1, &mut coeffs_a);
144
145    // Performs c = intt(ntt(a) / q_level)
146    if NTT {
147        ring_rns.ntt_inplace::<false>(&mut a);
148    }
149
150    ring_rns.div_by_last_modulus_inplace::<ROUND, NTT>(&mut buf, &mut a);
151
152    if NTT {
153        ring_rns
154            .at_level(a.level() - 1)
155            .intt_inplace::<false>(&mut a);
156    }
157
158    // Exports c to coeffs_c
159    let mut coeffs_c = vec![BigInt::from(0); a.n()];
160    ring_rns
161        .at_level(a.level() - 1)
162        .to_bigint_inplace(&a, 1, &mut coeffs_c);
163
164    // Performs floor division on a
165    let scalar_big = BigInt::from(ring_rns.0[ring_rns.level()].modulus.q);
166    coeffs_a.iter_mut().for_each(|a| {
167        if ROUND {
168            *a = a.div_round(&scalar_big);
169        } else {
170            *a = a.div_floor(&scalar_big);
171        }
172    });
173
174    izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
175}
176
177fn test_div_by_last_moduli<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
178    let seed: [u8; 32] = [0; 32];
179    let mut source: Source = Source::new(seed);
180
181    let nb_moduli_dropped: usize = ring_rns.level();
182
183    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
184    let mut buf0: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
185    let mut buf1: PolyRNS<u64> = ring_rns.new_polyrns();
186    let mut c: PolyRNS<u64> = ring_rns
187        .at_level(ring_rns.level() - nb_moduli_dropped)
188        .new_polyrns();
189
190    // Allocates a random PolyRNS
191    ring_rns.fill_uniform(&mut source, &mut a);
192
193    // Maps PolyRNS to [BigInt]
194    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
195    ring_rns
196        .at_level(a.level())
197        .to_bigint_inplace(&a, 1, &mut coeffs_a);
198
199    // Performs c = intt(ntt(a) / q_level)
200    if NTT {
201        ring_rns.ntt_inplace::<false>(&mut a);
202    }
203
204    ring_rns.div_by_last_moduli::<ROUND, NTT>(nb_moduli_dropped, &a, &mut buf0, &mut buf1, &mut c);
205
206    if NTT {
207        ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c);
208    }
209
210    // Exports c to coeffs_c
211    let mut coeffs_c = vec![BigInt::from(0); a.n()];
212    ring_rns
213        .at_level(c.level())
214        .to_bigint_inplace(&c, 1, &mut coeffs_c);
215
216    // Performs floor division on a
217    let mut scalar_big = BigInt::from(1);
218    (0..nb_moduli_dropped)
219        .for_each(|i| scalar_big *= BigInt::from(ring_rns.0[ring_rns.level() - i].modulus.q));
220    coeffs_a.iter_mut().for_each(|a| {
221        if ROUND {
222            *a = a.div_round(&scalar_big);
223        } else {
224            *a = a.div_floor(&scalar_big);
225        }
226    });
227
228    izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
229}
230
231fn test_div_by_last_moduli_inplace<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
232    let seed: [u8; 32] = [0; 32];
233    let mut source: Source = Source::new(seed);
234
235    let nb_moduli_dropped: usize = ring_rns.level();
236
237    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
238    let mut buf0: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
239    let mut buf1: PolyRNS<u64> = ring_rns.new_polyrns();
240
241    // Allocates a random PolyRNS
242    ring_rns.fill_uniform(&mut source, &mut a);
243
244    // Maps PolyRNS to [BigInt]
245    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
246    ring_rns
247        .at_level(a.level())
248        .to_bigint_inplace(&a, 1, &mut coeffs_a);
249
250    // Performs c = intt(ntt(a) / q_level)
251    if NTT {
252        ring_rns.ntt_inplace::<false>(&mut a);
253    }
254
255    ring_rns.div_by_last_moduli_inplace::<ROUND, NTT>(
256        nb_moduli_dropped,
257        &mut buf0,
258        &mut buf1,
259        &mut a,
260    );
261
262    if NTT {
263        ring_rns
264            .at_level(a.level() - nb_moduli_dropped)
265            .intt_inplace::<false>(&mut a);
266    }
267
268    // Exports c to coeffs_c
269    let mut coeffs_c = vec![BigInt::from(0); a.n()];
270    ring_rns
271        .at_level(a.level() - nb_moduli_dropped)
272        .to_bigint_inplace(&a, 1, &mut coeffs_c);
273
274    // Performs floor division on a
275    let mut scalar_big = BigInt::from(1);
276    (0..nb_moduli_dropped)
277        .for_each(|i| scalar_big *= BigInt::from(ring_rns.0[ring_rns.level() - i].modulus.q));
278    coeffs_a.iter_mut().for_each(|a| {
279        if ROUND {
280            *a = a.div_round(&scalar_big);
281        } else {
282            *a = a.div_floor(&scalar_big);
283        }
284    });
285
286    izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
287}