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