1use poulpy_hal::{
 2    api::{
 3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigAddInplace,
 4        VecZnxBigAddSmallInplace, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume,
 5        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSubInplace,
 6    },
 7    layouts::{Backend, DataRef, Module, Scratch, ScratchOwned},
 8    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl},
 9};
10
11use crate::layouts::{GLWE, GLWEPlaintext, LWEInfos, prepared::GLWESecretPrepared};
12
13impl<D: DataRef> GLWE<D> {
14    pub fn noise<B, DataSk, DataPt>(
15        &self,
16        module: &Module<B>,
17        sk_prepared: &GLWESecretPrepared<DataSk, B>,
18        pt_want: &GLWEPlaintext<DataPt>,
19        scratch: &mut Scratch<B>,
20    ) -> f64
21    where
22        DataSk: DataRef,
23        DataPt: DataRef,
24        B: Backend,
25        Module<B>: VecZnxDftApply<B>
26            + VecZnxSubInplace
27            + VecZnxNormalizeInplace<B>
28            + SvpApplyDftToDftInplace<B>
29            + VecZnxIdftApplyConsume<B>
30            + VecZnxBigAddInplace<B>
31            + VecZnxBigAddSmallInplace<B>
32            + VecZnxBigNormalize<B>,
33        Scratch<B>: TakeVecZnxDft<B> + TakeVecZnxBig<B>,
34    {
35        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(self);
36        self.decrypt(module, &mut pt_have, sk_prepared, scratch);
37        module.vec_znx_sub_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
38        module.vec_znx_normalize_inplace(self.base2k().into(), &mut pt_have.data, 0, scratch);
39        pt_have.data.std(self.base2k().into(), 0).log2()
40    }
41
42    pub fn assert_noise<B, DataSk, DataPt>(
43        &self,
44        module: &Module<B>,
45        sk_prepared: &GLWESecretPrepared<DataSk, B>,
46        pt_want: &GLWEPlaintext<DataPt>,
47        max_noise: f64,
48    ) where
49        DataSk: DataRef,
50        DataPt: DataRef,
51        Module<B>: VecZnxDftBytesOf
52            + VecZnxBigBytesOf
53            + VecZnxDftApply<B>
54            + SvpApplyDftToDftInplace<B>
55            + VecZnxIdftApplyConsume<B>
56            + VecZnxBigAddInplace<B>
57            + VecZnxBigAddSmallInplace<B>
58            + VecZnxBigNormalize<B>
59            + VecZnxNormalizeTmpBytes
60            + VecZnxSubInplace
61            + VecZnxNormalizeInplace<B>,
62        B: Backend + TakeVecZnxDftImpl<B> + TakeVecZnxBigImpl<B> + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
63    {
64        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWE::decrypt_scratch_space(module, self));
65        let noise_have: f64 = self.noise(module, sk_prepared, pt_want, scratch.borrow());
66        assert!(noise_have <= max_noise, "{noise_have} {max_noise}");
67    }
68}