1use crate::{
 2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow, ZnNormalizeInplace, ZnNormalizeTmpBytes},
 3    layouts::{Backend, Module, ScratchOwned, Zn, ZnToMut, ZnxInfos, ZnxView, ZnxViewMut},
 4    reference::znx::{ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStepInplace, ZnxNormalizeMiddleStepInplace, ZnxRef},
 5    source::Source,
 6};
 7
 8pub fn zn_normalize_tmp_bytes(n: usize) -> usize {
 9    n * size_of::<i64>()
10}
11
12pub fn zn_normalize_inplace<R, ARI>(n: usize, basek: usize, res: &mut R, res_col: usize, carry: &mut [i64])
13where
14    R: ZnToMut,
15    ARI: ZnxNormalizeFirstStepInplace + ZnxNormalizeFinalStepInplace + ZnxNormalizeMiddleStepInplace,
16{
17    let mut res: Zn<&mut [u8]> = res.to_mut();
18
19    #[cfg(debug_assertions)]
20    {
21        assert_eq!(carry.len(), res.n());
22    }
23
24    let res_size: usize = res.size();
25
26    for j in (0..res_size).rev() {
27        let out = &mut res.at_mut(res_col, j)[..n];
28
29        if j == res_size - 1 {
30            ARI::znx_normalize_first_step_inplace(basek, 0, out, carry);
31        } else if j == 0 {
32            ARI::znx_normalize_final_step_inplace(basek, 0, out, carry);
33        } else {
34            ARI::znx_normalize_middle_step_inplace(basek, 0, out, carry);
35        }
36    }
37}
38
39pub fn test_zn_normalize_inplace<B: Backend>(module: &Module<B>)
40where
41    Module<B>: ZnNormalizeInplace<B> + ZnNormalizeTmpBytes,
42    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
43{
44    let mut source: Source = Source::new([0u8; 32]);
45    let cols: usize = 2;
46    let basek: usize = 12;
47
48    let n = 33;
49
50    let mut carry: Vec<i64> = vec![0i64; zn_normalize_tmp_bytes(n)];
51
52    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.zn_normalize_tmp_bytes(module.n()));
53
54    for res_size in [1, 2, 6, 11] {
55        let mut res_0: Zn<Vec<u8>> = Zn::alloc(n, cols, res_size);
56        let mut res_1: Zn<Vec<u8>> = Zn::alloc(n, cols, res_size);
57
58        res_0
59            .raw_mut()
60            .iter_mut()
61            .for_each(|x| *x = source.next_i32() as i64);
62        res_1.raw_mut().copy_from_slice(res_0.raw());
63
64        // Reference
65        for i in 0..cols {
66            zn_normalize_inplace::<_, ZnxRef>(n, basek, &mut res_0, i, &mut carry);
67            module.zn_normalize_inplace(n, basek, &mut res_1, i, scratch.borrow());
68        }
69
70        assert_eq!(res_0.raw(), res_1.raw());
71    }
72}