1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, VecZnxAddScalarInplace, VecZnxBigAddInplace,
  4        VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
  5        VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA,
  6        VecZnxNormalizeTmpBytes, VecZnxSubABInplace,
  7    },
  8    layouts::{Backend, DataRef, Module, ScalarZnx, ScratchOwned, VecZnxBig, VecZnxDft, ZnxZero},
  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            + VecZnxDftApply<B>
 27            + SvpApplyDftToDftInplace<B>
 28            + VecZnxIdftApplyConsume<B>
 29            + VecZnxBigAddInplace<B>
 30            + VecZnxBigAddSmallInplace<B>
 31            + VecZnxBigNormalize<B>
 32            + VecZnxNormalizeTmpBytes
 33            + VecZnxBigAlloc<B>
 34            + VecZnxDftAlloc<B>
 35            + VecZnxBigNormalizeTmpBytes
 36            + VecZnxIdftApplyTmpA<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_apply(1, 0, &mut pt_dft, 0, &pt.data, 0);
 61                    module.svp_apply_dft_to_dft_inplace(&mut pt_dft, 0, &sk_prepared.data, col_j - 1);
 62                    module.vec_znx_idft_apply_tmpa(&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                assert!(std_pt <= noise, "{} > {}", std_pt, noise);
 74
 75                pt.data.zero();
 76            });
 77        });
 78    }
 79}
 80
 81impl<D: DataRef> GGSWCiphertext<D> {
 82    pub fn print_noise<B, DataSk, DataScalar>(
 83        &self,
 84        module: &Module<B>,
 85        sk_prepared: &GLWESecretPrepared<DataSk, B>,
 86        pt_want: &ScalarZnx<DataScalar>,
 87    ) where
 88        DataSk: DataRef,
 89        DataScalar: DataRef,
 90        Module<B>: VecZnxDftAllocBytes
 91            + VecZnxBigAllocBytes
 92            + VecZnxDftApply<B>
 93            + SvpApplyDftToDftInplace<B>
 94            + VecZnxIdftApplyConsume<B>
 95            + VecZnxBigAddInplace<B>
 96            + VecZnxBigAddSmallInplace<B>
 97            + VecZnxBigNormalize<B>
 98            + VecZnxNormalizeTmpBytes
 99            + VecZnxBigAlloc<B>
100            + VecZnxDftAlloc<B>
101            + VecZnxBigNormalizeTmpBytes
102            + VecZnxIdftApplyTmpA<B>
103            + VecZnxAddScalarInplace
104            + VecZnxSubABInplace,
105        B: Backend + TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
106    {
107        let basek: usize = self.basek();
108        let k: usize = self.k();
109        let digits: usize = self.digits();
110
111        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
112        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
113        let mut pt_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(1, self.size());
114        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(1, self.size());
115
116        let mut scratch: ScratchOwned<B> =
117            ScratchOwned::alloc(GLWECiphertext::decrypt_scratch_space(module, basek, k) | module.vec_znx_normalize_tmp_bytes());
118
119        (0..self.rank() + 1).for_each(|col_j| {
120            (0..self.rows()).for_each(|row_i| {
121                module.vec_znx_add_scalar_inplace(&mut pt.data, 0, (digits - 1) + row_i * digits, pt_want, 0);
122
123                // mul with sk[col_j-1]
124                if col_j > 0 {
125                    module.vec_znx_dft_apply(1, 0, &mut pt_dft, 0, &pt.data, 0);
126                    module.svp_apply_dft_to_dft_inplace(&mut pt_dft, 0, &sk_prepared.data, col_j - 1);
127                    module.vec_znx_idft_apply_tmpa(&mut pt_big, 0, &mut pt_dft, 0);
128                    module.vec_znx_big_normalize(basek, &mut pt.data, 0, &pt_big, 0, scratch.borrow());
129                }
130
131                self.at(row_i, col_j)
132                    .decrypt(module, &mut pt_have, sk_prepared, scratch.borrow());
133
134                module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt.data, 0);
135
136                let std_pt: f64 = pt_have.data.std(basek, 0).log2();
137                println!("col: {} row: {}: {}", col_j, row_i, std_pt);
138                pt.data.zero();
139            });
140        });
141    }
142}