1use itertools::izip;
  2use sampling::source::Source;
  3
  4use crate::hal::{
  5    api::{VecZnxAddNormal, VecZnxDecodeVeci64, VecZnxEncodeVeci64, VecZnxFillUniform, VecZnxStd, ZnxInfos, ZnxView, ZnxViewMut},
  6    layouts::{Backend, Module, VecZnx},
  7};
  8
  9pub fn test_vec_znx_fill_uniform<B: Backend>(module: &Module<B>)
 10where
 11    Module<B>: VecZnxFillUniform + VecZnxStd,
 12{
 13    let n: usize = module.n();
 14    let basek: usize = 17;
 15    let size: usize = 5;
 16    let mut source: Source = Source::new([0u8; 32]);
 17    let cols: usize = 2;
 18    let zero: Vec<i64> = vec![0; n];
 19    let one_12_sqrt: f64 = 0.28867513459481287;
 20    (0..cols).for_each(|col_i| {
 21        let mut a: VecZnx<_> = VecZnx::alloc(n, cols, size);
 22        module.vec_znx_fill_uniform(basek, &mut a, col_i, size * basek, &mut source);
 23        (0..cols).for_each(|col_j| {
 24            if col_j != col_i {
 25                (0..size).for_each(|limb_i| {
 26                    assert_eq!(a.at(col_j, limb_i), zero);
 27                })
 28            } else {
 29                let std: f64 = module.vec_znx_std(basek, &a, col_i);
 30                assert!(
 31                    (std - one_12_sqrt).abs() < 0.01,
 32                    "std={} ~!= {}",
 33                    std,
 34                    one_12_sqrt
 35                );
 36            }
 37        })
 38    });
 39}
 40
 41pub fn test_vec_znx_add_normal<B: Backend>(module: &Module<B>)
 42where
 43    Module<B>: VecZnxAddNormal + VecZnxStd,
 44{
 45    let n: usize = module.n();
 46    let basek: usize = 17;
 47    let k: usize = 2 * 17;
 48    let size: usize = 5;
 49    let sigma: f64 = 3.2;
 50    let bound: f64 = 6.0 * sigma;
 51    let mut source: Source = Source::new([0u8; 32]);
 52    let cols: usize = 2;
 53    let zero: Vec<i64> = vec![0; n];
 54    let k_f64: f64 = (1u64 << k as u64) as f64;
 55    (0..cols).for_each(|col_i| {
 56        let mut a: VecZnx<_> = VecZnx::alloc(n, cols, size);
 57        module.vec_znx_add_normal(basek, &mut a, col_i, k, &mut source, sigma, bound);
 58        (0..cols).for_each(|col_j| {
 59            if col_j != col_i {
 60                (0..size).for_each(|limb_i| {
 61                    assert_eq!(a.at(col_j, limb_i), zero);
 62                })
 63            } else {
 64                let std: f64 = module.vec_znx_std(basek, &a, col_i) * k_f64;
 65                assert!((std - sigma).abs() < 0.1, "std={} ~!= {}", std, sigma);
 66            }
 67        })
 68    });
 69}
 70
 71pub fn test_vec_znx_encode_vec_i64_lo_norm<B: Backend>(module: &Module<B>)
 72where
 73    Module<B>: VecZnxEncodeVeci64 + VecZnxDecodeVeci64,
 74{
 75    let n: usize = module.n();
 76    let basek: usize = 17;
 77    let size: usize = 5;
 78    let k: usize = size * basek - 5;
 79    let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, 2, size);
 80    let mut source: Source = Source::new([0u8; 32]);
 81    let raw: &mut [i64] = a.raw_mut();
 82    raw.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
 83    (0..a.cols()).for_each(|col_i| {
 84        let mut have: Vec<i64> = vec![i64::default(); n];
 85        have.iter_mut()
 86            .for_each(|x| *x = (source.next_i64() << 56) >> 56);
 87        module.encode_vec_i64(basek, &mut a, col_i, k, &have, 10);
 88        let mut want: Vec<i64> = vec![i64::default(); n];
 89        module.decode_vec_i64(basek, &a, col_i, k, &mut want);
 90        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
 91    });
 92}
 93
 94pub fn test_vec_znx_encode_vec_i64_hi_norm<B: Backend>(module: &Module<B>)
 95where
 96    Module<B>: VecZnxEncodeVeci64 + VecZnxDecodeVeci64,
 97{
 98    let n: usize = module.n();
 99    let basek: usize = 17;
100    let size: usize = 5;
101    for k in [1, basek / 2, size * basek - 5] {
102        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, 2, size);
103        let mut source = Source::new([0u8; 32]);
104        let raw: &mut [i64] = a.raw_mut();
105        raw.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
106        (0..a.cols()).for_each(|col_i| {
107            let mut have: Vec<i64> = vec![i64::default(); n];
108            have.iter_mut().for_each(|x| {
109                if k < 64 {
110                    *x = source.next_u64n(1 << k, (1 << k) - 1) as i64;
111                } else {
112                    *x = source.next_i64();
113                }
114            });
115            module.encode_vec_i64(basek, &mut a, col_i, k, &have, 63);
116            let mut want: Vec<i64> = vec![i64::default(); n];
117            module.decode_vec_i64(basek, &a, col_i, k, &mut want);
118            izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
119        })
120    }
121}