1use itertools::izip;
 2use rns::automorphism::AutoPerm;
 3use rns::poly::Poly;
 4use rns::ring::Ring;
 5
 6#[test]
 7fn automorphism_u64() {
 8    let n: usize = 1 << 4;
 9    let nth_root: usize = n << 1;
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_automorphism_native_u64::<NTT:false>", || {
15        test_automorphism_native_u64::<false>(&ring, nth_root)
16    });
17    sub_test("test_automorphism_native_u64::<NTT:true>", || {
18        test_automorphism_native_u64::<true>(&ring, nth_root)
19    });
20
21    sub_test("test_automorphism_from_perm_u64::<NTT:false>", || {
22        test_automorphism_from_perm_u64::<false>(&ring)
23    });
24    sub_test("test_automorphism_from_perm_u64::<NTT:true>", || {
25        test_automorphism_from_perm_u64::<true>(&ring)
26    });
27}
28
29fn sub_test<F: FnOnce()>(name: &str, f: F) {
30    println!("Running {}", name);
31    f();
32}
33
34fn test_automorphism_native_u64<const NTT: bool>(ring: &Ring<u64>, nth_root: usize) {
35    let n: usize = ring.n();
36    let q: u64 = ring.modulus.q;
37
38    let mut p0: Poly<u64> = ring.new_poly();
39    let mut p1: Poly<u64> = ring.new_poly();
40
41    for i in 0..p0.n() {
42        p0.0[i] = i as u64
43    }
44
45    if NTT {
46        ring.ntt_inplace::<false>(&mut p0);
47    }
48
49    let gal_el: usize = 2 * nth_root - 1;
50
51    ring.a_apply_automorphism_native_into_b::<NTT>(&p0, gal_el, nth_root, &mut p1);
52
53    if NTT {
54        ring.intt_inplace::<false>(&mut p1);
55    }
56
57    p0.0[0] = 0;
58    for i in 1..p0.n() {
59        p0.0[i] = q - (n - i) as u64
60    }
61
62    izip!(p0.0, p1.0).for_each(|(a, b)| assert_eq!(a, b));
63}
64
65fn test_automorphism_from_perm_u64<const NTT: bool>(ring: &Ring<u64>) {
66    let n: usize = ring.n();
67    let q: u64 = ring.modulus.q;
68
69    let mut p0: Poly<u64> = ring.new_poly();
70    let mut p1: Poly<u64> = ring.new_poly();
71
72    for i in 0..p0.n() {
73        p0.0[i] = i as u64
74    }
75
76    if NTT {
77        ring.ntt_inplace::<false>(&mut p0);
78    }
79
80    let gen_1 = 0;
81    let gen_2 = true;
82
83    let auto_perm = AutoPerm::new::<_, NTT>(&ring, gen_1, gen_2);
84
85    ring.a_apply_automorphism_from_perm_into_b::<NTT>(&p0, &auto_perm, &mut p1);
86
87    if NTT {
88        ring.intt_inplace::<false>(&mut p1);
89    }
90
91    p0.0[0] = 0;
92    for i in 1..p0.n() {
93        p0.0[i] = q - (n - i) as u64
94    }
95
96    izip!(p0.0, p1.0).for_each(|(a, b)| assert_eq!(a, b));
97}