1use itertools::izip;
 2use math::automorphism::AutoPerm;
 3use math::poly::Poly;
 4use math::ring::Ring;
 5
 6#[test]
 7fn ring_switch_u64() {
 8    let n: usize = 1 << 4;
 9    let q_base: u64 = 65537u64;
10    let q_power: usize = 1usize;
11    let ring_small: Ring<u64> = Ring::new(n, q_base, q_power);
12    let ring_large = Ring::new(2 * n, q_base, q_power);
13
14    sub_test("test_ring_switch_small_to_large_u64::<NTT:false>", || {
15        test_ring_switch_small_to_large_u64::<false>(&ring_small, &ring_large)
16    });
17    sub_test("test_ring_switch_small_to_large_u64::<NTT:true>", || {
18        test_ring_switch_small_to_large_u64::<true>(&ring_small, &ring_large)
19    });
20    sub_test("test_ring_switch_large_to_small_u64::<NTT:false>", || {
21        test_ring_switch_large_to_small_u64::<false>(&ring_small, &ring_large)
22    });
23    sub_test("test_ring_switch_large_to_small_u64::<NTT:true>", || {
24        test_ring_switch_large_to_small_u64::<true>(&ring_small, &ring_large)
25    });
26}
27
28fn sub_test<F: FnOnce()>(name: &str, f: F) {
29    println!("Running {}", name);
30    f();
31}
32
33fn test_ring_switch_small_to_large_u64<const NTT: bool>(
34    ring_small: &Ring<u64>,
35    ring_large: &Ring<u64>,
36) {
37    let mut a: Poly<u64> = ring_small.new_poly();
38    let mut buf: Poly<u64> = ring_small.new_poly();
39    let mut b: Poly<u64> = ring_large.new_poly();
40
41    a.0.iter_mut().enumerate().for_each(|(i, x)| *x = i as u64);
42
43    if NTT {
44        ring_small.ntt_inplace::<false>(&mut a);
45    }
46
47    ring_large.switch_degree::<NTT>(&a, &mut buf, &mut b);
48
49    if NTT {
50        ring_small.intt_inplace::<false>(&mut a);
51        ring_large.intt_inplace::<false>(&mut b);
52    }
53
54    let gap: usize = ring_large.n() / ring_small.n();
55
56    b.0.iter()
57        .step_by(gap)
58        .zip(a.0.iter())
59        .for_each(|(x_out, x_in)| assert_eq!(x_out, x_in));
60}
61
62fn test_ring_switch_large_to_small_u64<const NTT: bool>(
63    ring_small: &Ring<u64>,
64    ring_large: &Ring<u64>,
65) {
66    let mut a: Poly<u64> = ring_large.new_poly();
67    let mut buf: Poly<u64> = ring_large.new_poly();
68    let mut b: Poly<u64> = ring_small.new_poly();
69
70    a.0.iter_mut().enumerate().for_each(|(i, x)| *x = i as u64);
71
72    if NTT {
73        ring_large.ntt_inplace::<false>(&mut a);
74    }
75
76    ring_large.switch_degree::<NTT>(&a, &mut buf, &mut b);
77
78    if NTT {
79        ring_large.intt_inplace::<false>(&mut a);
80        ring_small.intt_inplace::<false>(&mut b);
81    }
82
83    let gap: usize = ring_large.n() / ring_small.n();
84
85    a.0.iter()
86        .step_by(gap)
87        .zip(b.0.iter())
88        .for_each(|(x_out, x_in)| assert_eq!(x_out, x_in));
89}