1use itertools::izip;
 2use math::modulus::{WordOps, ONCE};
 3use math::poly::Poly;
 4use math::ring::Ring;
 5use sampling::source::Source;
 6
 7#[test]
 8fn digit_decomposition() {
 9    let n: usize = 1 << 4;
10    let q_base: u64 = 65537u64;
11    let q_power: usize = 1usize;
12    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
13
14    sub_test("test_unsigned_digit_decomposition", || {
15        test_unsigned_digit_decomposition(&ring)
16    });
17
18    sub_test("test_signed_digit_decomposition::<BALANCED=false>", || {
19        test_signed_digit_decomposition::<false>(&ring)
20    });
21
22    sub_test("test_signed_digit_decomposition::<BALANCED=true>", || {
23        test_signed_digit_decomposition::<true>(&ring)
24    });
25}
26
27fn sub_test<F: FnOnce()>(name: &str, f: F) {
28    println!("Running {}", name);
29    f();
30}
31
32fn test_unsigned_digit_decomposition(ring: &Ring<u64>) {
33    let mut a: Poly<u64> = ring.new_poly();
34    let mut b: Poly<u64> = ring.new_poly();
35    let mut c: Poly<u64> = ring.new_poly();
36
37    let seed: [u8; 32] = [0; 32];
38    let mut source: Source = Source::new(seed);
39    ring.fill_uniform(&mut source, &mut a);
40
41    let base: usize = 8;
42    let log_q: usize = ring.modulus.q.log2();
43    let d: usize = ((log_q + base - 1) / base) as _;
44
45    (0..d).for_each(|i| {
46        ring.a_ith_digit_unsigned_base_scalar_b_into_c(i, &a, &base, &mut b);
47        ring.a_mul_b_scalar_into_a::<ONCE>(&(1 << (i * base)), &mut b);
48        ring.a_add_b_into_b::<ONCE>(&b, &mut c);
49    });
50
51    izip!(a.0, c.0).for_each(|(a, c)| assert_eq!(a, c));
52}
53
54fn test_signed_digit_decomposition<const BALANCED: bool>(ring: &Ring<u64>) {
55    let mut a: Poly<u64> = ring.new_poly();
56    let mut b: Poly<u64> = ring.new_poly();
57    let mut carry: Poly<u64> = ring.new_poly();
58    let mut c: Poly<u64> = ring.new_poly();
59
60    let seed: [u8; 32] = [0; 32];
61    let mut source: Source = Source::new(seed);
62    ring.fill_uniform(&mut source, &mut a);
63
64    let base: usize = 8;
65    let log_q: usize = ring.modulus.q.log2();
66    let d: usize = ((log_q + base - 1) / base) as _;
67
68    (0..d).for_each(|i| {
69        ring.a_ith_digit_signed_base_scalar_b_into_c::<BALANCED>(i, &a, &base, &mut carry, &mut b);
70        ring.a_mul_b_scalar_into_a::<ONCE>(&(1 << (i * base)), &mut b);
71        ring.a_add_b_into_b::<ONCE>(&b, &mut c);
72    });
73
74    izip!(a.0, c.0).for_each(|(a, c)| assert_eq!(a, c));
75}