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_sparse_u64::<false>(&ring, 1)
 66    });
 67    sub_test("test_packing_u64::<NTT:true>", || {
 68        test_packing_sparse_u64::<true>(&ring, 1)
 69    });
 70    sub_test("test_packing_sparse_u64::<NTT:false>", || {
 71        test_packing_sparse_u64::<false>(&ring, 3)
 72    });
 73    sub_test("test_packing_sparse_u64::<NTT:true>", || {
 74        test_packing_sparse_u64::<true>(&ring, 3)
 75    });
 76}
 77
 78fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>, gap: usize) {
 79    let n: usize = ring.n();
 80
 81    let mut result: Vec<Option<&mut Poly<u64>>> = Vec::with_capacity(n);
 82    result.resize_with(n, || None);
 83
 84    let mut polys: Vec<Poly<u64>> = vec![ring.new_poly(); (n+gap-1)/gap];
 85
 86    polys.iter_mut().enumerate().for_each(|(i , poly)|{
 87        poly.fill(&((1 + i*gap) as u64));
 88        if NTT {
 89            ring.ntt_inplace::<false>(poly);
 90        }
 91        result[i*gap] = Some(poly);
 92    });
 93
 94    ring.pack::<true, NTT>(&mut result, ring.log_n());
 95
 96    if let Some(poly) = result[0].as_mut() {
 97
 98        if NTT {
 99            ring.intt_inplace::<false>(poly);
100        }
101
102        poly.0.iter().enumerate().for_each(|(i, x)| {
103            if i % gap == 0 {
104                assert_eq!(*x, (1+i) as u64)
105            } else {
106                assert_eq!(*x, 0u64)
107            }
108        });
109    }
110}
111
112#[test]
113fn packing_streaming_u64() {
114    let n: usize = 1 << 5;
115    let q_base: u64 = 65537u64;
116    let q_power: usize = 1usize;
117    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
118
119    sub_test("test_packing_streaming_dense_u64::<NTT:true>", || {
120        test_packing_streaming_dense_u64::<true>(&ring)
121    });
122}
123
124fn test_packing_streaming_dense_u64<const NTT: bool>(ring: &Ring<u64>) {
125    let n: usize = ring.n();
126
127    let mut values: Vec<u64> = vec![0; n];
128    values
129        .iter_mut()
130        .enumerate()
131        .for_each(|(i, x)| *x = (i + 1) as u64);
132
133    let gap: usize = 3;
134
135    let mut packer = StreamRepacker::new(ring);
136
137    let mut poly: Poly<u64> = ring.new_poly();
138    for i in 0..n {
139        let i_rev: usize = i.reverse_bits_msb(ring.log_n() as u32);
140
141        if i_rev % gap == 0 {
142            poly.fill(&values[i_rev]);
143            if NTT {
144                ring.ntt_inplace::<false>(&mut poly);
145            }
146            packer.add::<NTT>(ring, Some(&poly))
147        } else {
148            packer.add::<NTT>(ring, None)
149        }
150    }
151
152    packer.flush::<NTT>(ring);
153
154    let result: &mut Poly<u64> = &mut packer.results[0];
155
156    if NTT {
157        ring.intt_inplace::<false>(result);
158    }
159
160    result.0.iter().enumerate().for_each(|(i, x)| {
161        if i % gap == 0 {
162            assert_eq!(*x, values[i] as u64)
163        } else {
164            assert_eq!(*x, 0u64)
165        }
166    });
167}
168
169#[test]
170fn trace_u64() {
171    let n: usize = 1 << 5;
172    let q_base: u64 = 65537u64;
173    let q_power: usize = 1usize;
174    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
175
176    sub_test("test_trace::<NTT:false>", || test_trace_u64::<false>(&ring));
177    sub_test("test_trace::<NTT:true>", || test_trace_u64::<true>(&ring));
178}
179
180fn test_trace_u64<const NTT: bool>(ring: &Ring<u64>) {
181    let n: usize = ring.n();
182
183    let mut poly: Poly<u64> = ring.new_poly();
184
185    poly.0
186        .iter_mut()
187        .enumerate()
188        .for_each(|(i, x)| *x = (i + 1) as u64);
189
190    if NTT {
191        ring.ntt_inplace::<false>(&mut poly);
192    }
193
194    let step_start: usize = 2;
195
196    ring.trace_inplace::<NTT>(step_start, &mut poly);
197
198    if NTT {
199        ring.intt_inplace::<false>(&mut poly);
200    }
201
202    let gap: usize = 1 << (ring.log_n() - step_start);
203
204    poly.0.iter().enumerate().for_each(|(i, x)| {
205        if i % gap == 0 {
206            assert_eq!(*x, 1 + i as u64)
207        } else {
208            assert_eq!(*x, 0u64)
209        }
210    });
211}