1use backend::hal::{
 2    api::{
 3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAlloc, VecZnxBigAlloc, VecZnxBigNormalize,
 4        VecZnxBigNormalizeTmpBytes, VecZnxDftAlloc, VecZnxDftToVecZnxBigTmpA, VecZnxNormalizeTmpBytes, VecZnxStd,
 5        VecZnxSubABInplace, ZnxZero,
 6    },
 7    layouts::{Backend, DataRef, Module, ScalarZnx, ScratchOwned, VecZnxBig, VecZnxDft},
 8    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl},
 9};
10
11use crate::{GGSWCiphertext, GLWECiphertext, GLWEDecryptFamily, GLWEPlaintext, GLWESecretExec, Infos};
12
13pub trait GGSWAssertNoiseFamily<B: Backend> = GLWEDecryptFamily<B>
14    + VecZnxBigAlloc<B>
15    + VecZnxDftAlloc<B>
16    + VecZnxBigNormalizeTmpBytes
17    + VecZnxBigNormalize<B>
18    + VecZnxDftToVecZnxBigTmpA<B>;
19
20impl<D: DataRef> GGSWCiphertext<D> {
21    pub fn assert_noise<B: Backend, DataSk, DataScalar, F>(
22        &self,
23        module: &Module<B>,
24        sk_exec: &GLWESecretExec<DataSk, B>,
25        pt_want: &ScalarZnx<DataScalar>,
26        max_noise: F,
27    ) where
28        DataSk: DataRef,
29        DataScalar: DataRef,
30        Module<B>: GGSWAssertNoiseFamily<B> + VecZnxAlloc + VecZnxAddScalarInplace + VecZnxSubABInplace + VecZnxStd,
31        B: TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
32        F: Fn(usize) -> f64,
33    {
34        let basek: usize = self.basek();
35        let k: usize = self.k();
36        let digits: usize = self.digits();
37
38        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k);
39        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k);
40        let mut pt_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(1, self.size());
41        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(1, self.size());
42
43        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
44            GLWECiphertext::decrypt_scratch_space(module, basek, k) | module.vec_znx_normalize_tmp_bytes(module.n()),
45        );
46
47        (0..self.rank() + 1).for_each(|col_j| {
48            (0..self.rows()).for_each(|row_i| {
49                module.vec_znx_add_scalar_inplace(&mut pt.data, 0, (digits - 1) + row_i * digits, pt_want, 0);
50
51                // mul with sk[col_j-1]
52                if col_j > 0 {
53                    module.vec_znx_dft_from_vec_znx(1, 0, &mut pt_dft, 0, &pt.data, 0);
54                    module.svp_apply_inplace(&mut pt_dft, 0, &sk_exec.data, col_j - 1);
55                    module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
56                    module.vec_znx_big_normalize(basek, &mut pt.data, 0, &pt_big, 0, scratch.borrow());
57                }
58
59                self.at(row_i, col_j)
60                    .decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
61
62                module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt.data, 0);
63
64                let std_pt: f64 = module.vec_znx_std(basek, &pt_have.data, 0).log2();
65                let noise: f64 = max_noise(col_j);
66                println!("{} {}", std_pt, noise);
67                assert!(std_pt <= noise, "{} > {}", std_pt, noise);
68
69                pt.data.zero();
70            });
71        });
72    }
73}