1use itertools::izip;
  2use math::num_bigint::Div;
  3use math::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
 48    //sub_test("test_div_by_last_moduli::<ROUND:false, NTT:false>", ||{test_div_by_last_moduli::<false, false>(&ring_rns)});
 49}
 50
 51fn sub_test<F: FnOnce()>(name: &str, f: F) {
 52    println!("Running {}", name);
 53    f();
 54}
 55
 56fn test_div_by_last_modulus<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
 57    let seed: [u8; 32] = [0; 32];
 58    let mut source: Source = Source::new(seed);
 59
 60    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
 61    let mut b: PolyRNS<u64> = ring_rns.new_polyrns();
 62    let mut c: PolyRNS<u64> = ring_rns.at_level(ring_rns.level() - 1).new_polyrns();
 63
 64    // Allocates a random PolyRNS
 65    ring_rns.fill_uniform(&mut source, &mut a);
 66
 67    // Maps PolyRNS to [BigInt]
 68    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
 69    ring_rns
 70        .at_level(a.level())
 71        .to_bigint_inplace(&a, 1, &mut coeffs_a);
 72
 73    // Performs c = intt(ntt(a) / q_level)
 74    if NTT {
 75        ring_rns.ntt_inplace::<false>(&mut a);
 76    }
 77
 78    ring_rns.div_by_last_modulus::<ROUND, NTT>(&a, &mut b, &mut c);
 79
 80    if NTT {
 81        ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c);
 82    }
 83
 84    // Exports c to coeffs_c
 85    let mut coeffs_c = vec![BigInt::from(0); c.n()];
 86    ring_rns
 87        .at_level(c.level())
 88        .to_bigint_inplace(&c, 1, &mut coeffs_c);
 89
 90    // Performs floor division on a
 91    let scalar_big = BigInt::from(ring_rns.0[ring_rns.level()].modulus.q);
 92    coeffs_a.iter_mut().for_each(|a| {
 93        if ROUND {
 94            *a = a.div_round(&scalar_big);
 95        } else {
 96            *a = a.div_floor(&scalar_big);
 97        }
 98    });
 99
100    izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
101}
102
103fn test_div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
104    let seed: [u8; 32] = [0; 32];
105    let mut source: Source = Source::new(seed);
106
107    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
108    let mut buf: PolyRNS<u64> = ring_rns.new_polyrns();
109
110    // Allocates a random PolyRNS
111    ring_rns.fill_uniform(&mut source, &mut a);
112
113    // Maps PolyRNS to [BigInt]
114    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
115    ring_rns
116        .at_level(a.level())
117        .to_bigint_inplace(&a, 1, &mut coeffs_a);
118
119    // Performs c = intt(ntt(a) / q_level)
120    if NTT {
121        ring_rns.ntt_inplace::<false>(&mut a);
122    }
123
124    ring_rns.div_by_last_modulus_inplace::<ROUND, NTT>(&mut buf, &mut a);
125
126    if NTT {
127        ring_rns
128            .at_level(a.level() - 1)
129            .intt_inplace::<false>(&mut a);
130    }
131
132    // Exports c to coeffs_c
133    let mut coeffs_c = vec![BigInt::from(0); a.n()];
134    ring_rns
135        .at_level(a.level() - 1)
136        .to_bigint_inplace(&a, 1, &mut coeffs_c);
137
138    // Performs floor division on a
139    let scalar_big = BigInt::from(ring_rns.0[ring_rns.level()].modulus.q);
140    coeffs_a.iter_mut().for_each(|a| {
141        if ROUND {
142            *a = a.div_round(&scalar_big);
143        } else {
144            *a = a.div_floor(&scalar_big);
145        }
146    });
147
148    izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
149}
150
151fn test_div_by_last_moduli<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
152    let seed: [u8; 32] = [0; 32];
153    let mut source: Source = Source::new(seed);
154
155    let nb_moduli: usize = ring_rns.level();
156
157    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
158    let mut buf: PolyRNS<u64> = ring_rns.new_polyrns();
159    let mut c: PolyRNS<u64> = ring_rns
160        .at_level(ring_rns.level() - nb_moduli)
161        .new_polyrns();
162
163    // Allocates a random PolyRNS
164    ring_rns.fill_uniform(&mut source, &mut a);
165
166    // Maps PolyRNS to [BigInt]
167    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
168    ring_rns
169        .at_level(a.level())
170        .to_bigint_inplace(&a, 1, &mut coeffs_a);
171
172    // Performs c = intt(ntt(a) / q_level)
173    if NTT {
174        ring_rns.ntt_inplace::<false>(&mut a);
175    }
176
177    ring_rns.div_by_last_moduli::<ROUND, NTT>(nb_moduli, &a, &mut buf, &mut c);
178
179    if NTT {
180        ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c);
181    }
182
183    // Exports c to coeffs_c
184    let mut coeffs_c = vec![BigInt::from(0); a.n()];
185    ring_rns
186        .at_level(c.level())
187        .to_bigint_inplace(&c, 1, &mut coeffs_c);
188
189    // Performs floor division on a
190    let mut scalar_big = BigInt::from(1);
191    (0..nb_moduli)
192        .for_each(|i| scalar_big *= BigInt::from(ring_rns.0[ring_rns.level() - i].modulus.q));
193    coeffs_a.iter_mut().for_each(|a| {
194        if ROUND {
195            *a = a.div_round(&scalar_big);
196        } else {
197            *a = a.div_floor(&scalar_big);
198        }
199    });
200
201    izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
202}
203
204/*
205fn test_div_floor_by_last_moduli_inplace<const NTT: bool>(ring_rns: &RingRNS<u64>) {
206    let seed: [u8; 32] = [0; 32];
207    let mut source: Source = Source::new(seed);
208
209    let nb_moduli: usize = ring_rns.level();
210
211    let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
212    let mut b: PolyRNS<u64> = ring_rns.new_polyrns();
213
214    // Allocates a random PolyRNS
215    ring_rns.fill_uniform(&mut source, &mut a);
216
217    // Maps PolyRNS to [BigInt]
218    let mut coeffs_a: Vec<BigInt> = (0..a.n()).map(|i| BigInt::from(i)).collect();
219    ring_rns
220        .at_level(a.level())
221        .to_bigint_inplace(&a, 1, &mut coeffs_a);
222
223    // Performs c = intt(ntt(a) / q_level)
224    if NTT {
225        ring_rns.ntt_inplace::<false>(&mut a);
226    }
227
228    ring_rns.div_floor_by_last_moduli_inplace::<NTT>(nb_moduli, &mut b, &mut a);
229
230    if NTT {
231        ring_rns.at_level(a.level()-nb_moduli).intt_inplace::<false>(&mut a);
232    }
233
234    // Exports c to coeffs_c
235    let mut coeffs_c = vec![BigInt::from(0); a.n()];
236    ring_rns
237        .at_level(a.level()-nb_moduli)
238        .to_bigint_inplace(&a, 1, &mut coeffs_c);
239
240    // Performs floor division on a
241    let mut scalar_big = BigInt::from(1);
242    (0..nb_moduli).for_each(|i|{scalar_big *= BigInt::from(ring_rns.0[ring_rns.level()-i].modulus.q)});
243    coeffs_a.iter_mut().for_each(|a| {a.div_floor(&scalar_big)});
244
245    assert!(coeffs_a == coeffs_c, "test_div_floor_by_last_moduli_inplace");
246}
247*/