1use poulpy_hal::{
  2    api::{ModuleN, ScratchAvailable, VecZnxAddScalarInplace, VecZnxDftBytesOf, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes},
  3    // oep::SvpPPolAllocBytesImpl,
  4    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxZero},
  5    source::Source,
  6};
  7
  8use crate::{
  9    encryption::glwe_ct::GLWEEncryptSk, layouts::GLWEInfos, ScratchTakeCore,
 10    layouts::{
 11        GGLWE, GGLWEInfos, GGLWEToMut, GLWE, GLWEPlaintext, LWEInfos,
 12        prepared::{GLWESecretPrepared, GLWESecretPreparedToRef},
 13    },
 14};
 15
 16impl GGLWE<Vec<u8>> {
 17
 18    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 19    where
 20        A: GGLWEInfos,
 21        M: GGLWEEncryptSk<BE>,
 22    {
 23        module.gglwe_encrypt_sk_tmp_bytes(infos)
 24    }
 25
 26    pub fn encrypt_pk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 27    where
 28        A: GGLWEInfos,
 29        M: GGLWEEncryptSk<BE>,
 30    {
 31        module.gglwe_encrypt_sk_tmp_bytes(infos)
 32    }
 33}
 34
 35pub trait GGLWEEncryptSk<B: Backend> {
 36
 37    fn gglwe_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 38    where
 39        A: GGLWEInfos;
 40
 41    fn gglwe_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
 42        where
 43            A: GGLWEInfos;        
 44
 45    fn gglwe_encrypt_sk<R, P, S>(
 46        &self,
 47        res: &mut R,
 48        pt: &P,
 49        sk: &S,
 50        source_xa: &mut Source,
 51        source_xe: &mut Source,
 52        scratch: &mut Scratch<B>,
 53    ) where
 54        R: GGLWEToMut,
 55        P: ScalarZnxToRef,
 56        S: GLWESecretPreparedToRef<B>;
 57}
 58
 59impl<B: Backend> GGLWEEncryptSk<B> for Module<B>
 60where
 61    Module<B>: ModuleN
 62        + GLWEEncryptSk<B>
 63        + VecZnxNormalizeTmpBytes
 64        + VecZnxDftBytesOf
 65        + VecZnxAddScalarInplace
 66        + VecZnxNormalizeInplace<B>,
 67    // + SvpPPolAllocBytesImpl<B>,
 68    Scratch<B>: ScratchAvailable + ScratchTakeCore<B>,
 69{
 70    fn gglwe_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 71    where
 72        A: GGLWEInfos,
 73    {
 74        GLWE::encrypt_sk_tmp_bytes(self, &infos.glwe_layout())
 75            + (GLWEPlaintext::bytes_of_from_infos(self, &infos.glwe_layout()) | self.vec_znx_normalize_tmp_bytes())
 76    }
 77
 78    fn gglwe_encrypt_pk_tmp_bytes<A>(&self, _infos: &A) -> usize
 79    where
 80        A: GGLWEInfos,
 81    {
 82        unimplemented!()
 83        // TODO: Is this correct?
 84        // GLWE::encrypt_pk_tmp_bytes(self, &infos.glwe_layout())
 85        //     + (GLWEPlaintext::bytes_of_from_infos(self, &infos.glwe_layout()) | self.vec_znx_normalize_tmp_bytes())
 86
 87    }
 88
 89    fn gglwe_encrypt_sk<R, P, S>(
 90        &self,
 91        res: &mut R,
 92        pt: &P,
 93        sk: &S,
 94        source_xa: &mut Source,
 95        source_xe: &mut Source,
 96        scratch: &mut Scratch<B>,
 97    ) where
 98        R: GGLWEToMut,
 99        P: ScalarZnxToRef,
100        S: GLWESecretPreparedToRef<B>,
101    {
102        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
103        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
104
105        #[cfg(debug_assertions)]
106        {
107            use poulpy_hal::layouts::ZnxInfos;
108            let sk: GLWESecretPrepared<&[u8], B> = sk.to_ref();
109
110            assert_eq!(
111                res.rank_in(),
112                pt.cols() as u32,
113                "res.rank_in(): {} != pt.cols(): {}",
114                res.rank_in(),
115                pt.cols()
116            );
117            assert_eq!(
118                res.rank_out(),
119                sk.rank(),
120                "res.rank_out(): {} != sk.rank(): {}",
121                res.rank_out(),
122                sk.rank()
123            );
124            assert_eq!(res.n(), sk.n());
125            assert_eq!(pt.n() as u32, sk.n());
126            assert!(
127                scratch.available() >= GGLWE::encrypt_sk_tmp_bytes(self, res),
128                "scratch.available: {} < GGLWECiphertext::encrypt_sk_tmp_bytes(self, res.rank()={}, res.size()={}): {}",
129                scratch.available(),
130                res.rank_out(),
131                res.size(),
132                GGLWE::encrypt_sk_tmp_bytes(self, res)
133            );
134            assert!(
135                res.dnum().0 * res.dsize().0 * res.base2k().0 <= res.k().0,
136                "res.dnum() : {} * res.dsize() : {} * res.base2k() : {} = {} >= res.k() = {}",
137                res.dnum(),
138                res.dsize(),
139                res.base2k(),
140                res.dnum().0 * res.dsize().0 * res.base2k().0,
141                res.k()
142            );
143        }
144
145        let dnum: usize = res.dnum().into();
146        let dsize: usize = res.dsize().into();
147        let base2k: usize = res.base2k().into();
148        let rank_in: usize = res.rank_in().into();
149
150        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, &res.glwe_layout());
151        // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
152        //
153        // Example for ksk rank 2 to rank 3:
154        //
155        // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
156        // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
157        //
158        // Example ksk rank 2 to rank 1
159        //
160        // (-(a*s) + s0, a)
161        // (-(b*s) + s1, b)
162
163        for col_i in 0..rank_in {
164            for row_i in 0..dnum {
165                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
166                tmp_pt.data.zero(); // zeroes for next iteration
167                self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, col_i);
168                self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
169                self.glwe_encrypt_sk(
170                    &mut res.at_mut(row_i, col_i),
171                    &tmp_pt,
172                    sk,
173                    source_xa,
174                    source_xe,
175                    scrach_1,
176                );
177            }
178        }
179    }
180}
181
182impl<DataSelf: DataMut> GGLWE<DataSelf> {
183    #[allow(clippy::too_many_arguments)]
184    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
185        &mut self,
186        module: &Module<B>,
187        pt: &ScalarZnx<DataPt>,
188        sk: &GLWESecretPrepared<DataSk, B>,
189        source_xa: &mut Source,
190        source_xe: &mut Source,
191        scratch: &mut Scratch<B>,
192    ) where
193        Module<B>: GGLWEEncryptSk<B>,
194    {
195        module.gglwe_encrypt_sk(self, pt, sk, source_xa, source_xe, scratch);
196    }
197}