1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, 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> + 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(self.n(), basek, k);
 39        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
 40        let mut pt_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(self.n(), 1, self.size());
 41        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(self.n(), 1, self.size());
 42
 43        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 44            GLWECiphertext::decrypt_scratch_space(module, self.n(), basek, k) | module.vec_znx_normalize_tmp_bytes(self.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}
 74
 75impl<D: DataRef> GGSWCiphertext<D> {
 76    pub fn print_noise<B: Backend, DataSk, DataScalar>(
 77        &self,
 78        module: &Module<B>,
 79        sk_exec: &GLWESecretExec<DataSk, B>,
 80        pt_want: &ScalarZnx<DataScalar>,
 81    ) where
 82        DataSk: DataRef,
 83        DataScalar: DataRef,
 84        Module<B>: GGSWAssertNoiseFamily<B> + VecZnxAddScalarInplace + VecZnxSubABInplace + VecZnxStd,
 85        B: TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
 86    {
 87        let basek: usize = self.basek();
 88        let k: usize = self.k();
 89        let digits: usize = self.digits();
 90
 91        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
 92        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
 93        let mut pt_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(self.n(), 1, self.size());
 94        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(self.n(), 1, self.size());
 95
 96        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 97            GLWECiphertext::decrypt_scratch_space(module, self.n(), basek, k) | module.vec_znx_normalize_tmp_bytes(module.n()),
 98        );
 99
100        (0..self.rank() + 1).for_each(|col_j| {
101            (0..self.rows()).for_each(|row_i| {
102                module.vec_znx_add_scalar_inplace(&mut pt.data, 0, (digits - 1) + row_i * digits, pt_want, 0);
103
104                // mul with sk[col_j-1]
105                if col_j > 0 {
106                    module.vec_znx_dft_from_vec_znx(1, 0, &mut pt_dft, 0, &pt.data, 0);
107                    module.svp_apply_inplace(&mut pt_dft, 0, &sk_exec.data, col_j - 1);
108                    module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
109                    module.vec_znx_big_normalize(basek, &mut pt.data, 0, &pt_big, 0, scratch.borrow());
110                }
111
112                self.at(row_i, col_j)
113                    .decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
114
115                module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt.data, 0);
116
117                let std_pt: f64 = module.vec_znx_std(basek, &pt_have.data, 0).log2();
118                println!("{}", std_pt);
119                pt.data.zero();
120            });
121        });
122    }
123}