1use poulpy_hal::{
  2    api::{
  3        ScratchTakeBasic, SvpApplyDftToDftInplace, VecZnxAddScalarInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
  4        VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume,
  5    },
  6    layouts::{Backend, DataRef, Module, ScalarZnxToRef, Scratch, Stats, ZnxZero},
  7};
  8
  9use crate::layouts::{GGSW, GGSWInfos, GGSWToRef, LWEInfos, prepared::GLWESecretPrepared};
 10use crate::{GLWENoise, layouts::prepared::GLWESecretPreparedToRef};
 11use crate::{ScratchTakeCore, layouts::GLWEPlaintext};
 12
 13impl<D: DataRef> GGSW<D> {
 14    pub fn noise<M, BE: Backend, P, S>(
 15        &self,
 16        module: &M,
 17        row: usize,
 18        col: usize,
 19        pt_want: &P,
 20        sk_prepared: &S,
 21        scratch: &mut Scratch<BE>,
 22    ) -> Stats
 23    where
 24        S: GLWESecretPreparedToRef<BE>,
 25        P: ScalarZnxToRef,
 26        M: GGSWNoise<BE>,
 27        Scratch<BE>: ScratchTakeCore<BE>,
 28    {
 29        module.ggsw_noise(self, row, col, pt_want, sk_prepared, scratch)
 30    }
 31}
 32
 33pub trait GGSWNoise<BE: Backend> {
 34    fn ggsw_noise_tmp_bytes<A>(&self, infos: &A) -> usize
 35    where
 36        A: GGSWInfos;
 37
 38    fn ggsw_noise<R, S, P>(
 39        &self,
 40        res: &R,
 41        res_row: usize,
 42        res_col: usize,
 43        pt_want: &P,
 44        sk_prepared: &S,
 45        scratch: &mut Scratch<BE>,
 46    ) -> Stats
 47    where
 48        R: GGSWToRef,
 49        S: GLWESecretPreparedToRef<BE>,
 50        P: ScalarZnxToRef;
 51}
 52
 53impl<BE: Backend> GGSWNoise<BE> for Module<BE>
 54where
 55    Module<BE>: VecZnxAddScalarInplace
 56        + VecZnxDftApply<BE>
 57        + SvpApplyDftToDftInplace<BE>
 58        + VecZnxIdftApplyConsume<BE>
 59        + VecZnxDftBytesOf
 60        + VecZnxBigNormalize<BE>
 61        + VecZnxBigNormalizeTmpBytes
 62        + GLWENoise<BE>,
 63    Scratch<BE>: ScratchTakeCore<BE>,
 64{
 65    fn ggsw_noise_tmp_bytes<A>(&self, infos: &A) -> usize
 66    where
 67        A: GGSWInfos,
 68    {
 69        GLWEPlaintext::bytes_of_from_infos(infos)
 70            + (self.bytes_of_vec_znx_dft(1, infos.size()) + self.vec_znx_big_normalize_tmp_bytes())
 71                .max(self.glwe_noise_tmp_bytes(infos))
 72    }
 73
 74    fn ggsw_noise<R, S, P>(
 75        &self,
 76        res: &R,
 77        res_row: usize,
 78        res_col: usize,
 79        pt_want: &P,
 80        sk_prepared: &S,
 81        scratch: &mut Scratch<BE>,
 82    ) -> Stats
 83    where
 84        R: GGSWToRef,
 85        S: GLWESecretPreparedToRef<BE>,
 86        P: ScalarZnxToRef,
 87        Scratch<BE>: ScratchTakeCore<BE>,
 88    {
 89        let res: &GGSW<&[u8]> = &res.to_ref();
 90        let sk_prepared: &GLWESecretPrepared<&[u8], BE> = &sk_prepared.to_ref();
 91
 92        let base2k: usize = res.base2k().into();
 93        let dsize: usize = res.dsize().into();
 94
 95        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(res);
 96        pt.data_mut().zero();
 97        self.vec_znx_add_scalar_inplace(&mut pt.data, 0, (dsize - 1) + res_row * dsize, pt_want, 0);
 98
 99        // mul with sk[col_j-1]
100        if res_col > 0 {
101            let (mut pt_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, 1, res.size());
102            self.vec_znx_dft_apply(1, 0, &mut pt_dft, 0, &pt.data, 0);
103            self.svp_apply_dft_to_dft_inplace(&mut pt_dft, 0, &sk_prepared.data, res_col - 1);
104            let pt_big = self.vec_znx_idft_apply_consume(pt_dft);
105            self.vec_znx_big_normalize(&mut pt.data, base2k, 0, 0, &pt_big, base2k, 0, scratch_2);
106        }
107
108        self.glwe_noise(&res.at(res_row, res_col), &pt, sk_prepared, scratch_1)
109    }
110}