1use itertools::izip;
  2use math::poly::Poly;
  3use math::ring::impl_u64::ring;
  4use math::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_u64::<NTT:false>", || {
 15        test_automorphism_u64::<false>(&ring, nth_root)
 16    });
 17    sub_test("test_automorphism_u64::<NTT:true>", || {
 18        test_automorphism_u64::<true>(&ring, nth_root)
 19    });
 20}
 21
 22fn sub_test<F: FnOnce()>(name: &str, f: F) {
 23    println!("Running {}", name);
 24    f();
 25}
 26
 27fn test_automorphism_u64<const NTT: bool>(ring: &Ring<u64>, nth_root: usize) {
 28    let n: usize = ring.n();
 29    let q: u64 = ring.modulus.q;
 30
 31    let mut p0: Poly<u64> = ring.new_poly();
 32    let mut p1: Poly<u64> = ring.new_poly();
 33
 34    for i in 0..p0.n() {
 35        p0.0[i] = i as u64
 36    }
 37
 38    if NTT {
 39        ring.ntt_inplace::<false>(&mut p0);
 40    }
 41
 42    ring.automorphism::<NTT>(&p0, 2 * n - 1, nth_root, &mut p1);
 43
 44    if NTT {
 45        ring.intt_inplace::<false>(&mut p1);
 46    }
 47
 48    p0.0[0] = 0;
 49    for i in 1..p0.n() {
 50        p0.0[i] = q - (n - i) as u64
 51    }
 52
 53    izip!(p0.0, p1.0).for_each(|(a, b)| assert_eq!(a, b));
 54}
 55
 56#[test]
 57fn packing_u64() {
 58    let n: usize = 1 << 5;
 59    let q_base: u64 = 65537u64;
 60    let q_power: usize = 1usize;
 61    let ring: Ring<u64> = Ring::new(n, q_base, q_power);
 62
 63    sub_test("test_packing_u64::<NTT:false>", || {
 64        test_packing_full_u64::<false>(&ring)
 65    });
 66    sub_test("test_packing_u64::<NTT:true>", || {
 67        test_packing_full_u64::<true>(&ring)
 68    });
 69    sub_test("test_packing_sparse_u64::<NTT:false>", || {
 70        test_packing_sparse_u64::<false>(&ring)
 71    });
 72    sub_test("test_packing_sparse_u64::<NTT:true>", || {
 73        test_packing_sparse_u64::<true>(&ring)
 74    });
 75}
 76
 77fn test_packing_full_u64<const NTT: bool>(ring: &Ring<u64>) {
 78    let n: usize = ring.n();
 79
 80    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
 81
 82    for i in 0..n {
 83        let mut poly: Poly<u64> = ring.new_poly();
 84        poly.fill(&(1 + i as u64));
 85        if NTT {
 86            ring.ntt_inplace::<false>(&mut poly);
 87        }
 88
 89        result[i] = Some(poly);
 90    }
 91
 92    ring.pack::<true, NTT>(&mut result, ring.log_n());
 93
 94    if let Some(poly) = result[0].as_mut() {
 95        if NTT {
 96            ring.intt_inplace::<false>(poly);
 97        }
 98
 99        poly.0
100            .iter()
101            .enumerate()
102            .for_each(|(i, x)| assert_eq!(*x, 1 + i as u64));
103    }
104}
105
106fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>) {
107    let n: usize = ring.n();
108
109    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
110
111    let gap: usize = 3;
112
113    for i in (0..n).step_by(gap) {
114        let mut poly: Poly<u64> = ring.new_poly();
115        poly.fill(&(1 + i as u64));
116        if NTT {
117            ring.ntt_inplace::<false>(&mut poly);
118        }
119        result[i] = Some(poly);
120    }
121
122    ring.pack::<true, NTT>(&mut result, ring.log_n());
123
124    if let Some(poly) = result[0].as_mut() {
125        if NTT {
126            ring.intt_inplace::<false>(poly);
127        }
128
129        poly.0.iter().enumerate().for_each(|(i, x)| {
130            if i % gap == 0 {
131                assert_eq!(*x, 1 + i as u64)
132            } else {
133                assert_eq!(*x, 0u64)
134            }
135        });
136    }
137}