1use sampling::source::Source;
 2
 3use crate::hal::{
 4    api::{ZnxInfos, ZnxViewMut},
 5    layouts::VecZnx,
 6};
 7
 8pub fn test_vec_znx_encode_vec_i64_lo_norm() {
 9    let n: usize = 32;
10    let basek: usize = 17;
11    let size: usize = 5;
12    let k: usize = size * basek - 5;
13    let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, 2, size);
14    let mut source: Source = Source::new([0u8; 32]);
15    let raw: &mut [i64] = a.raw_mut();
16    raw.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
17    (0..a.cols()).for_each(|col_i| {
18        let mut have: Vec<i64> = vec![i64::default(); n];
19        have.iter_mut()
20            .for_each(|x| *x = (source.next_i64() << 56) >> 56);
21        a.encode_vec_i64(basek, col_i, k, &have, 10);
22        let mut want: Vec<i64> = vec![i64::default(); n];
23        a.decode_vec_i64(basek, col_i, k, &mut want);
24        assert_eq!(have, want, "{:?} != {:?}", &have, &want);
25    });
26}
27
28pub fn test_vec_znx_encode_vec_i64_hi_norm() {
29    let n: usize = 32;
30    let basek: usize = 17;
31    let size: usize = 5;
32    for k in [1, basek / 2, size * basek - 5] {
33        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, 2, size);
34        let mut source = Source::new([0u8; 32]);
35        let raw: &mut [i64] = a.raw_mut();
36        raw.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
37        (0..a.cols()).for_each(|col_i| {
38            let mut have: Vec<i64> = vec![i64::default(); n];
39            have.iter_mut().for_each(|x| {
40                if k < 64 {
41                    *x = source.next_u64n(1 << k, (1 << k) - 1) as i64;
42                } else {
43                    *x = source.next_i64();
44                }
45            });
46            a.encode_vec_i64(basek, col_i, k, &have, 63);
47            let mut want: Vec<i64> = vec![i64::default(); n];
48            a.decode_vec_i64(basek, col_i, k, &mut want);
49            assert_eq!(have, want, "{:?} != {:?}", &have, &want);
50        })
51    }
52}