1use itertools::izip;
  2use math::modulus::WordOps;
  3use math::poly::Poly;
  4use math::ring::impl_u64::packing::StreamRepacker;
  5use math::ring::Ring;
  6
  7#[test]
  8fn automorphism_u64() {
  9    let n: usize = 1 << 4;
 10    let nth_root: usize = n << 1;
 11    let q_base: u64 = 65537u64;
 12    let q_power: usize = 1usize;
 13    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
 14
 15    sub_test("test_automorphism_u64::<NTT:false>", || {
 16        test_automorphism_u64::<false>(&ring, nth_root)
 17    });
 18    sub_test("test_automorphism_u64::<NTT:true>", || {
 19        test_automorphism_u64::<true>(&ring, nth_root)
 20    });
 21}
 22
 23fn sub_test<F: FnOnce()>(name: &str, f: F) {
 24    println!("Running {}", name);
 25    f();
 26}
 27
 28fn test_automorphism_u64<const NTT: bool>(ring: &Ring<u64>, nth_root: usize) {
 29    let n: usize = ring.n();
 30    let q: u64 = ring.modulus.q;
 31
 32    let mut p0: Poly<u64> = ring.new_poly();
 33    let mut p1: Poly<u64> = ring.new_poly();
 34
 35    for i in 0..p0.n() {
 36        p0.0[i] = i as u64
 37    }
 38
 39    if NTT {
 40        ring.ntt_inplace::<false>(&mut p0);
 41    }
 42
 43    ring.a_apply_automorphism_into_b::<NTT>(&p0, 2 * n - 1, nth_root, &mut p1);
 44
 45    if NTT {
 46        ring.intt_inplace::<false>(&mut p1);
 47    }
 48
 49    p0.0[0] = 0;
 50    for i in 1..p0.n() {
 51        p0.0[i] = q - (n - i) as u64
 52    }
 53
 54    izip!(p0.0, p1.0).for_each(|(a, b)| assert_eq!(a, b));
 55}
 56
 57#[test]
 58fn packing_u64() {
 59    let n: usize = 1 << 5;
 60    let q_base: u64 = 65537u64;
 61    let q_power: usize = 1usize;
 62    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
 63
 64    sub_test("test_packing_u64::<NTT:false>", || {
 65        test_packing_full_u64::<false>(&ring)
 66    });
 67    sub_test("test_packing_u64::<NTT:true>", || {
 68        test_packing_full_u64::<true>(&ring)
 69    });
 70    sub_test("test_packing_sparse_u64::<NTT:false>", || {
 71        test_packing_sparse_u64::<false>(&ring)
 72    });
 73    sub_test("test_packing_sparse_u64::<NTT:true>", || {
 74        test_packing_sparse_u64::<true>(&ring)
 75    });
 76}
 77
 78fn test_packing_full_u64<const NTT: bool>(ring: &Ring<u64>) {
 79    let n: usize = ring.n();
 80
 81    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
 82
 83    for i in 0..n {
 84        let mut poly: Poly<u64> = ring.new_poly();
 85        poly.fill(&(1 + i as u64));
 86        if NTT {
 87            ring.ntt_inplace::<false>(&mut poly);
 88        }
 89
 90        result[i] = Some(poly);
 91    }
 92
 93    ring.pack::<true, NTT>(&mut result, ring.log_n());
 94
 95    if let Some(poly) = result[0].as_mut() {
 96        if NTT {
 97            ring.intt_inplace::<false>(poly);
 98        }
 99
100        poly.0
101            .iter()
102            .enumerate()
103            .for_each(|(i, x)| assert_eq!(*x, 1 + i as u64));
104    }
105}
106
107fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>) {
108    let n: usize = ring.n();
109
110    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
111
112    let gap: usize = 3;
113
114    for i in (0..n).step_by(gap) {
115        let mut poly: Poly<u64> = ring.new_poly();
116        poly.fill(&(1 + i as u64));
117        if NTT {
118            ring.ntt_inplace::<false>(&mut poly);
119        }
120        result[i] = Some(poly);
121    }
122
123    ring.pack::<true, NTT>(&mut result, ring.log_n());
124
125    if let Some(poly) = result[0].as_mut() {
126        if NTT {
127            ring.intt_inplace::<false>(poly);
128        }
129
130        poly.0.iter().enumerate().for_each(|(i, x)| {
131            if i % gap == 0 {
132                assert_eq!(*x, 1 + i as u64)
133            } else {
134                assert_eq!(*x, 0u64)
135            }
136        });
137    }
138}
139
140#[test]
141fn packing_streaming_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_packing_streaming_dense_u64::<NTT:true>", || {
148        test_packing_streaming_dense_u64::<true>(&ring)
149    });
150}
151
152fn test_packing_streaming_dense_u64<const NTT: bool>(ring: &Ring<u64>) {
153    let n: usize = ring.n();
154
155    let mut values: Vec<u64> = vec![0; n];
156    values
157        .iter_mut()
158        .enumerate()
159        .for_each(|(i, x)| *x = (i + 1) as u64);
160
161    let gap: usize = 3;
162
163    let mut packer = StreamRepacker::new(ring);
164
165    let mut poly: Poly<u64> = ring.new_poly();
166    for i in 0..n {
167        let i_rev: usize = i.reverse_bits_msb(ring.log_n() as u32);
168
169        if i_rev % gap == 0 {
170            poly.fill(&values[i_rev]);
171            if NTT {
172                ring.ntt_inplace::<false>(&mut poly);
173            }
174            packer.add::<NTT>(ring, Some(&poly))
175        } else {
176            packer.add::<NTT>(ring, None)
177        }
178    }
179
180    packer.flush::<NTT>(ring);
181
182    let result: &mut Poly<u64> = &mut packer.results[0];
183
184    if NTT {
185        ring.intt_inplace::<false>(result);
186    }
187
188    result.0.iter().enumerate().for_each(|(i, x)| {
189        if i % gap == 0 {
190            assert_eq!(*x, values[i] as u64)
191        } else {
192            assert_eq!(*x, 0u64)
193        }
194    });
195}
196
197#[test]
198fn trace_u64() {
199    let n: usize = 1 << 5;
200    let q_base: u64 = 65537u64;
201    let q_power: usize = 1usize;
202    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
203
204    sub_test("test_trace::<NTT:false>", || test_trace_u64::<false>(&ring));
205    sub_test("test_trace::<NTT:true>", || test_trace_u64::<true>(&ring));
206}
207
208fn test_trace_u64<const NTT: bool>(ring: &Ring<u64>) {
209    let n: usize = ring.n();
210
211    let mut poly: Poly<u64> = ring.new_poly();
212
213    poly.0
214        .iter_mut()
215        .enumerate()
216        .for_each(|(i, x)| *x = (i + 1) as u64);
217
218    if NTT {
219        ring.ntt_inplace::<false>(&mut poly);
220    }
221
222    let step_start: usize = 2;
223
224    ring.trace_inplace::<NTT>(step_start, &mut poly);
225
226    if NTT {
227        ring.intt_inplace::<false>(&mut poly);
228    }
229
230    let gap: usize = 1 << (ring.log_n() - step_start);
231
232    poly.0.iter().enumerate().for_each(|(i, x)| {
233        if i % gap == 0 {
234            assert_eq!(*x, 1 + i as u64)
235        } else {
236            assert_eq!(*x, 0u64)
237        }
238    });
239}