1use itertools::izip;
  2use math::poly::Poly;
  3use math::ring::Ring;
  4
  5#[test]
  6fn automorphism_u64() {
  7    let n: usize = 1 << 4;
  8    let nth_root: usize = n << 1;
  9    let q_base: u64 = 65537u64;
 10    let q_power: usize = 1usize;
 11    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
 12
 13    sub_test("test_automorphism_u64::<NTT:false>", || {
 14        test_automorphism_u64::<false>(&ring, nth_root)
 15    });
 16    sub_test("test_automorphism_u64::<NTT:true>", || {
 17        test_automorphism_u64::<true>(&ring, nth_root)
 18    });
 19}
 20
 21fn sub_test<F: FnOnce()>(name: &str, f: F) {
 22    println!("Running {}", name);
 23    f();
 24}
 25
 26fn test_automorphism_u64<const NTT: bool>(ring: &Ring<u64>, nth_root: usize) {
 27    let n: usize = ring.n();
 28    let q: u64 = ring.modulus.q;
 29
 30    let mut p0: Poly<u64> = ring.new_poly();
 31    let mut p1: Poly<u64> = ring.new_poly();
 32
 33    for i in 0..p0.n() {
 34        p0.0[i] = i as u64
 35    }
 36
 37    if NTT {
 38        ring.ntt_inplace::<false>(&mut p0);
 39    }
 40
 41    ring.a_apply_automorphism_into_b::<NTT>(&p0, 2 * n - 1, nth_root, &mut p1);
 42
 43    if NTT {
 44        ring.intt_inplace::<false>(&mut p1);
 45    }
 46
 47    p0.0[0] = 0;
 48    for i in 1..p0.n() {
 49        p0.0[i] = q - (n - i) as u64
 50    }
 51
 52    izip!(p0.0, p1.0).for_each(|(a, b)| assert_eq!(a, b));
 53}
 54
 55#[test]
 56fn packing_u64() {
 57    let n: usize = 1 << 5;
 58    let q_base: u64 = 65537u64;
 59    let q_power: usize = 1usize;
 60    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
 61
 62    sub_test("test_packing_u64::<NTT:false>", || {
 63        test_packing_full_u64::<false>(&ring)
 64    });
 65    sub_test("test_packing_u64::<NTT:true>", || {
 66        test_packing_full_u64::<true>(&ring)
 67    });
 68    sub_test("test_packing_sparse_u64::<NTT:false>", || {
 69        test_packing_sparse_u64::<false>(&ring)
 70    });
 71    sub_test("test_packing_sparse_u64::<NTT:true>", || {
 72        test_packing_sparse_u64::<true>(&ring)
 73    });
 74}
 75
 76fn test_packing_full_u64<const NTT: bool>(ring: &Ring<u64>) {
 77    let n: usize = ring.n();
 78
 79    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
 80
 81    for i in 0..n {
 82        let mut poly: Poly<u64> = ring.new_poly();
 83        poly.fill(&(1 + i as u64));
 84        if NTT {
 85            ring.ntt_inplace::<false>(&mut poly);
 86        }
 87
 88        result[i] = Some(poly);
 89    }
 90
 91    ring.pack::<true, NTT>(&mut result, ring.log_n());
 92
 93    if let Some(poly) = result[0].as_mut() {
 94        if NTT {
 95            ring.intt_inplace::<false>(poly);
 96        }
 97
 98        poly.0
 99            .iter()
100            .enumerate()
101            .for_each(|(i, x)| assert_eq!(*x, 1 + i as u64));
102    }
103}
104
105fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>) {
106    let n: usize = ring.n();
107
108    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
109
110    let gap: usize = 3;
111
112    for i in (0..n).step_by(gap) {
113        let mut poly: Poly<u64> = ring.new_poly();
114        poly.fill(&(1 + i as u64));
115        if NTT {
116            ring.ntt_inplace::<false>(&mut poly);
117        }
118        result[i] = Some(poly);
119    }
120
121    ring.pack::<true, NTT>(&mut result, ring.log_n());
122
123    if let Some(poly) = result[0].as_mut() {
124        if NTT {
125            ring.intt_inplace::<false>(poly);
126        }
127
128        poly.0.iter().enumerate().for_each(|(i, x)| {
129            if i % gap == 0 {
130                assert_eq!(*x, 1 + i as u64)
131            } else {
132                assert_eq!(*x, 0u64)
133            }
134        });
135    }
136}
137
138
139
140#[test]
141fn trace_u64() {
142    let n: usize = 1 << 5;
143    let q_base: u64 = 65537u64;
144    let q_power: usize = 1usize;
145    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
146
147    sub_test("test_trace::<NTT:false>", || {
148        test_trace_u64::<false>(&ring)
149    });
150    sub_test("test_trace::<NTT:true>", || {
151        test_trace_u64::<true>(&ring)
152    });
153}
154
155fn test_trace_u64<const NTT: bool>(ring: &Ring<u64>) {
156    let n: usize = ring.n();
157
158    let mut poly: Poly<u64> = ring.new_poly();
159    
160    poly.0.iter_mut().enumerate().for_each(|(i, x)|{
161        *x = (i+1) as u64
162    });
163
164    if NTT{
165        ring.ntt_inplace::<false>(&mut poly);
166    }
167
168    let step_start: usize = 2;
169
170    ring.trace_inplace::<NTT>(step_start, &mut poly);
171
172    if NTT{
173        ring.intt_inplace::<false>(&mut poly);
174    }
175
176    let gap: usize = 1<<(ring.log_n() - step_start);
177
178    poly.0.iter().enumerate().for_each(|(i, x)| {
179        if i % gap == 0 {
180            assert_eq!(*x, 1 + i as u64)
181        } else {
182            assert_eq!(*x, 0u64)
183        }
184    });
185}