1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxBigAlloc, VecZnxBigNormalize, VecZnxDftAlloc,
  4        VecZnxDftToVecZnxBigTmpA, VecZnxNormalizeTmpBytes, VecZnxStd, VecZnxSubABInplace, ZnxZero,
  5    },
  6    layouts::{Backend, DataRef, Module, ScalarZnx, ScratchOwned, VecZnxBig, VecZnxDft},
  7    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl},
  8};
  9
 10use crate::{
 11    layouts::{GGSWCiphertext, GLWECiphertext, GLWEPlaintext, Infos, prepared::GLWESecretExec},
 12    trait_families::GGSWAssertNoiseFamily,
 13};
 14
 15impl<D: DataRef> GGSWCiphertext<D> {
 16    pub fn assert_noise<B: Backend, DataSk, DataScalar, F>(
 17        &self,
 18        module: &Module<B>,
 19        sk_exec: &GLWESecretExec<DataSk, B>,
 20        pt_want: &ScalarZnx<DataScalar>,
 21        max_noise: F,
 22    ) where
 23        DataSk: DataRef,
 24        DataScalar: DataRef,
 25        Module<B>: GGSWAssertNoiseFamily<B> + VecZnxAddScalarInplace + VecZnxSubABInplace + VecZnxStd,
 26        B: TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
 27        F: Fn(usize) -> f64,
 28    {
 29        let basek: usize = self.basek();
 30        let k: usize = self.k();
 31        let digits: usize = self.digits();
 32
 33        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
 34        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
 35        let mut pt_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(self.n(), 1, self.size());
 36        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(self.n(), 1, self.size());
 37
 38        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 39            GLWECiphertext::decrypt_scratch_space(module, self.n(), basek, k) | module.vec_znx_normalize_tmp_bytes(self.n()),
 40        );
 41
 42        (0..self.rank() + 1).for_each(|col_j| {
 43            (0..self.rows()).for_each(|row_i| {
 44                module.vec_znx_add_scalar_inplace(&mut pt.data, 0, (digits - 1) + row_i * digits, pt_want, 0);
 45
 46                // mul with sk[col_j-1]
 47                if col_j > 0 {
 48                    module.vec_znx_dft_from_vec_znx(1, 0, &mut pt_dft, 0, &pt.data, 0);
 49                    module.svp_apply_inplace(&mut pt_dft, 0, &sk_exec.data, col_j - 1);
 50                    module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
 51                    module.vec_znx_big_normalize(basek, &mut pt.data, 0, &pt_big, 0, scratch.borrow());
 52                }
 53
 54                self.at(row_i, col_j)
 55                    .decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
 56
 57                module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt.data, 0);
 58
 59                let std_pt: f64 = module.vec_znx_std(basek, &pt_have.data, 0).log2();
 60                let noise: f64 = max_noise(col_j);
 61                println!("{} {}", std_pt, noise);
 62                assert!(std_pt <= noise, "{} > {}", std_pt, noise);
 63
 64                pt.data.zero();
 65            });
 66        });
 67    }
 68}
 69
 70impl<D: DataRef> GGSWCiphertext<D> {
 71    pub fn print_noise<B: Backend, DataSk, DataScalar>(
 72        &self,
 73        module: &Module<B>,
 74        sk_exec: &GLWESecretExec<DataSk, B>,
 75        pt_want: &ScalarZnx<DataScalar>,
 76    ) where
 77        DataSk: DataRef,
 78        DataScalar: DataRef,
 79        Module<B>: GGSWAssertNoiseFamily<B> + VecZnxAddScalarInplace + VecZnxSubABInplace + VecZnxStd,
 80        B: TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
 81    {
 82        let basek: usize = self.basek();
 83        let k: usize = self.k();
 84        let digits: usize = self.digits();
 85
 86        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
 87        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
 88        let mut pt_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(self.n(), 1, self.size());
 89        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(self.n(), 1, self.size());
 90
 91        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 92            GLWECiphertext::decrypt_scratch_space(module, self.n(), basek, k) | module.vec_znx_normalize_tmp_bytes(module.n()),
 93        );
 94
 95        (0..self.rank() + 1).for_each(|col_j| {
 96            (0..self.rows()).for_each(|row_i| {
 97                module.vec_znx_add_scalar_inplace(&mut pt.data, 0, (digits - 1) + row_i * digits, pt_want, 0);
 98
 99                // mul with sk[col_j-1]
100                if col_j > 0 {
101                    module.vec_znx_dft_from_vec_znx(1, 0, &mut pt_dft, 0, &pt.data, 0);
102                    module.svp_apply_inplace(&mut pt_dft, 0, &sk_exec.data, col_j - 1);
103                    module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
104                    module.vec_znx_big_normalize(basek, &mut pt.data, 0, &pt_big, 0, scratch.borrow());
105                }
106
107                self.at(row_i, col_j)
108                    .decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
109
110                module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt.data, 0);
111
112                let std_pt: f64 = module.vec_znx_std(basek, &pt_have.data, 0).log2();
113                println!("{}", std_pt);
114                pt.data.zero();
115            });
116        });
117    }
118}