1use poulpy_backend::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(self.n(), 1, self.size());
 49        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(self.n(), 1, self.size());
 50
 51        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 52            GLWECiphertext::decrypt_scratch_space(module, self.n(), basek, k) | module.vec_znx_normalize_tmp_bytes(self.n()),
 53        );
 54
 55        (0..self.rank() + 1).for_each(|col_j| {
 56            (0..self.rows()).for_each(|row_i| {
 57                module.vec_znx_add_scalar_inplace(&mut pt.data, 0, (digits - 1) + row_i * digits, pt_want, 0);
 58
 59                // mul with sk[col_j-1]
 60                if col_j > 0 {
 61                    module.vec_znx_dft_from_vec_znx(1, 0, &mut pt_dft, 0, &pt.data, 0);
 62                    module.svp_apply_inplace(&mut pt_dft, 0, &sk_prepared.data, col_j - 1);
 63                    module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
 64                    module.vec_znx_big_normalize(basek, &mut pt.data, 0, &pt_big, 0, scratch.borrow());
 65                }
 66
 67                self.at(row_i, col_j)
 68                    .decrypt(module, &mut pt_have, sk_prepared, scratch.borrow());
 69
 70                module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt.data, 0);
 71
 72                let std_pt: f64 = pt_have.data.std(basek, 0).log2();
 73                let noise: f64 = max_noise(col_j);
 74                println!("{} {}", std_pt, noise);
 75                assert!(std_pt <= noise, "{} > {}", std_pt, noise);
 76
 77                pt.data.zero();
 78            });
 79        });
 80    }
 81}
 82
 83impl<D: DataRef> GGSWCiphertext<D> {
 84    pub fn print_noise<B, DataSk, DataScalar>(
 85        &self,
 86        module: &Module<B>,
 87        sk_prepared: &GLWESecretPrepared<DataSk, B>,
 88        pt_want: &ScalarZnx<DataScalar>,
 89    ) where
 90        DataSk: DataRef,
 91        DataScalar: DataRef,
 92        Module<B>: VecZnxDftAllocBytes
 93            + VecZnxBigAllocBytes
 94            + VecZnxDftFromVecZnx<B>
 95            + SvpApplyInplace<B>
 96            + VecZnxDftToVecZnxBigConsume<B>
 97            + VecZnxBigAddInplace<B>
 98            + VecZnxBigAddSmallInplace<B>
 99            + VecZnxBigNormalize<B>
100            + VecZnxNormalizeTmpBytes
101            + VecZnxBigAlloc<B>
102            + VecZnxDftAlloc<B>
103            + VecZnxBigNormalizeTmpBytes
104            + VecZnxDftToVecZnxBigTmpA<B>
105            + VecZnxAddScalarInplace
106            + VecZnxSubABInplace,
107        B: Backend + TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
108    {
109        let basek: usize = self.basek();
110        let k: usize = self.k();
111        let digits: usize = self.digits();
112
113        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
114        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(self.n(), basek, k);
115        let mut pt_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(self.n(), 1, self.size());
116        let mut pt_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(self.n(), 1, self.size());
117
118        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
119            GLWECiphertext::decrypt_scratch_space(module, self.n(), basek, k) | module.vec_znx_normalize_tmp_bytes(module.n()),
120        );
121
122        (0..self.rank() + 1).for_each(|col_j| {
123            (0..self.rows()).for_each(|row_i| {
124                module.vec_znx_add_scalar_inplace(&mut pt.data, 0, (digits - 1) + row_i * digits, pt_want, 0);
125
126                // mul with sk[col_j-1]
127                if col_j > 0 {
128                    module.vec_znx_dft_from_vec_znx(1, 0, &mut pt_dft, 0, &pt.data, 0);
129                    module.svp_apply_inplace(&mut pt_dft, 0, &sk_prepared.data, col_j - 1);
130                    module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
131                    module.vec_znx_big_normalize(basek, &mut pt.data, 0, &pt_big, 0, scratch.borrow());
132                }
133
134                self.at(row_i, col_j)
135                    .decrypt(module, &mut pt_have, sk_prepared, scratch.borrow());
136
137                module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt.data, 0);
138
139                let std_pt: f64 = pt_have.data.std(basek, 0).log2();
140                println!("col: {} row: {}: {}", col_j, row_i, std_pt);
141                pt.data.zero();
142            });
143        });
144    }
145}