1use backend::hal::{
 2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxSubScalarInplace, ZnxZero},
 3    layouts::{Backend, DataRef, Module, ScalarZnx, ScratchOwned},
 4    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl},
 5};
 6
 7use crate::{
 8    layouts::{GGLWECiphertext, GLWECiphertext, GLWEPlaintext, Infos, prepared::GLWESecretPrepared},
 9    trait_families::GLWEDecryptFamily,
10};
11
12impl<D: DataRef> GGLWECiphertext<D> {
13    pub fn assert_noise<B: Backend, DataSk, DataWant>(
14        self,
15        module: &Module<B>,
16        sk: &GLWESecretPrepared<DataSk, B>,
17        pt_want: &ScalarZnx<DataWant>,
18        max_noise: f64,
19    ) where
20        DataSk: DataRef,
21        DataWant: DataRef,
22        Module<B>: GLWEDecryptFamily<B> + VecZnxSubScalarInplace,
23        B: TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
24    {
25        let digits: usize = self.digits();
26        let basek: usize = self.basek();
27        let k: usize = self.k();
28
29        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWECiphertext::decrypt_scratch_space(
30            module,
31            self.n(),
32            basek,
33            k,
34        ));
35        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
36
37        (0..self.rank_in()).for_each(|col_i| {
38            (0..self.rows()).for_each(|row_i| {
39                self.at(row_i, col_i)
40                    .decrypt(&module, &mut pt, &sk, scratch.borrow());
41
42                module.vec_znx_sub_scalar_inplace(
43                    &mut pt.data,
44                    0,
45                    (digits - 1) + row_i * digits,
46                    pt_want,
47                    col_i,
48                );
49
50                let noise_have: f64 = pt.data.std(basek, 0).log2();
51
52                assert!(
53                    noise_have <= max_noise,
54                    "noise_have: {} > max_noise: {}",
55                    noise_have,
56                    max_noise
57                );
58
59                pt.data.zero();
60            });
61        });
62    }
63}