1use poulpy_hal::{
 2    api::VecZnxAddScalarInplace,
 3    layouts::{Backend, DataRef, Module, ScalarZnxToRef, Scratch, Stats, ZnxInfos, ZnxZero},
 4};
 5
 6use crate::{
 7    GLWENoise,
 8    layouts::{GGLWE, GGLWEInfos, GGLWEToRef, prepared::GLWESecretPreparedToRef},
 9};
10use crate::{ScratchTakeCore, layouts::GLWEPlaintext};
11
12impl<D: DataRef> GGLWE<D> {
13    pub fn noise<M, S, P, BE: Backend>(
14        &self,
15        module: &M,
16        row: usize,
17        col: usize,
18        pt_want: &P,
19        sk_prepared: &S,
20        scratch: &mut Scratch<BE>,
21    ) -> Stats
22    where
23        S: GLWESecretPreparedToRef<BE>,
24        P: ScalarZnxToRef,
25        M: GGLWENoise<BE>,
26    {
27        module.gglwe_noise(self, row, col, pt_want, sk_prepared, scratch)
28    }
29}
30
31pub trait GGLWENoise<BE: Backend> {
32    fn gglwe_noise_tmp_bytes<A>(&self, infos: &A) -> usize
33    where
34        A: GGLWEInfos;
35
36    fn gglwe_noise<R, S, P>(
37        &self,
38        res: &R,
39        res_row: usize,
40        res_col: usize,
41        pt_want: &P,
42        sk_prepared: &S,
43        scratch: &mut Scratch<BE>,
44    ) -> Stats
45    where
46        R: GGLWEToRef,
47        S: GLWESecretPreparedToRef<BE>,
48        P: ScalarZnxToRef;
49}
50
51impl<BE: Backend> GGLWENoise<BE> for Module<BE>
52where
53    Module<BE>: VecZnxAddScalarInplace + GLWENoise<BE>,
54    Scratch<BE>: ScratchTakeCore<BE>,
55{
56    fn gglwe_noise_tmp_bytes<A>(&self, infos: &A) -> usize
57    where
58        A: GGLWEInfos,
59    {
60        GLWEPlaintext::bytes_of_from_infos(infos) + self.glwe_noise_tmp_bytes(infos)
61    }
62
63    fn gglwe_noise<R, S, P>(
64        &self,
65        res: &R,
66        res_row: usize,
67        res_col: usize,
68        pt_want: &P,
69        sk_prepared: &S,
70        scratch: &mut Scratch<BE>,
71    ) -> Stats
72    where
73        R: GGLWEToRef,
74        S: GLWESecretPreparedToRef<BE>,
75        P: ScalarZnxToRef,
76    {
77        let res: &GGLWE<&[u8]> = &res.to_ref();
78        let dsize: usize = res.dsize().into();
79        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(res);
80        pt.data_mut().zero();
81        println!("col: {res_col} {}", pt_want.to_ref().cols());
82        self.vec_znx_add_scalar_inplace(
83            &mut pt.data,
84            0,
85            (dsize - 1) + res_row * dsize,
86            pt_want,
87            res_col,
88        );
89        self.glwe_noise(&res.at(res_row, res_col), &pt, sk_prepared, scratch_1)
90    }
91}