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    Scratch<B>: ScratchAvailable + ScratchTakeCore<B>,
 68{
 69    fn gglwe_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 70    where
 71        A: GGLWEInfos,
 72    {
 73        GLWE::encrypt_sk_tmp_bytes(self, &infos.glwe_layout())
 74            + (GLWEPlaintext::bytes_of_from_infos(self, &infos.glwe_layout()) | self.vec_znx_normalize_tmp_bytes())
 75    }
 76
 77    fn gglwe_encrypt_pk_tmp_bytes<A>(&self, _infos: &A) -> usize
 78    where
 79        A: GGLWEInfos,
 80    {
 81        unimplemented!()
 82        // TODO: Is this correct?
 83        // GLWE::encrypt_pk_tmp_bytes(self, &infos.glwe_layout())
 84        //     + (GLWEPlaintext::bytes_of_from_infos(self, &infos.glwe_layout()) | self.vec_znx_normalize_tmp_bytes())
 85
 86    }
 87
 88    fn gglwe_encrypt_sk<R, P, S>(
 89        &self,
 90        res: &mut R,
 91        pt: &P,
 92        sk: &S,
 93        source_xa: &mut Source,
 94        source_xe: &mut Source,
 95        scratch: &mut Scratch<B>,
 96    ) where
 97        R: GGLWEToMut,
 98        P: ScalarZnxToRef,
 99        S: GLWESecretPreparedToRef<B>,
100    {
101        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
102        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
103
104        #[cfg(debug_assertions)]
105        {
106            use poulpy_hal::layouts::ZnxInfos;
107            let sk: GLWESecretPrepared<&[u8], B> = sk.to_ref();
108
109            assert_eq!(
110                res.rank_in(),
111                pt.cols() as u32,
112                "res.rank_in(): {} != pt.cols(): {}",
113                res.rank_in(),
114                pt.cols()
115            );
116            assert_eq!(
117                res.rank_out(),
118                sk.rank(),
119                "res.rank_out(): {} != sk.rank(): {}",
120                res.rank_out(),
121                sk.rank()
122            );
123            assert_eq!(res.n(), sk.n());
124            assert_eq!(pt.n() as u32, sk.n());
125            assert!(
126                scratch.available() >= GGLWE::encrypt_sk_tmp_bytes(self, res),
127                "scratch.available: {} < GGLWECiphertext::encrypt_sk_tmp_bytes(self, res.rank()={}, res.size()={}): {}",
128                scratch.available(),
129                res.rank_out(),
130                res.size(),
131                GGLWE::encrypt_sk_tmp_bytes(self, res)
132            );
133            assert!(
134                res.dnum().0 * res.dsize().0 * res.base2k().0 <= res.k().0,
135                "res.dnum() : {} * res.dsize() : {} * res.base2k() : {} = {} >= res.k() = {}",
136                res.dnum(),
137                res.dsize(),
138                res.base2k(),
139                res.dnum().0 * res.dsize().0 * res.base2k().0,
140                res.k()
141            );
142        }
143
144        let dnum: usize = res.dnum().into();
145        let dsize: usize = res.dsize().into();
146        let base2k: usize = res.base2k().into();
147        let rank_in: usize = res.rank_in().into();
148
149        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, &res.glwe_layout());
150        // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
151        //
152        // Example for ksk rank 2 to rank 3:
153        //
154        // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
155        // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
156        //
157        // Example ksk rank 2 to rank 1
158        //
159        // (-(a*s) + s0, a)
160        // (-(b*s) + s1, b)
161
162        for col_i in 0..rank_in {
163            for row_i in 0..dnum {
164                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
165                tmp_pt.data.zero(); // zeroes for next iteration
166                self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, col_i);
167                self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
168                self.glwe_encrypt_sk(
169                    &mut res.at_mut(row_i, col_i),
170                    &tmp_pt,
171                    sk,
172                    source_xa,
173                    source_xe,
174                    scrach_1,
175                );
176            }
177        }
178    }
179}
180
181impl<DataSelf: DataMut> GGLWE<DataSelf> {
182    #[allow(clippy::too_many_arguments)]
183    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
184        &mut self,
185        module: &Module<B>,
186        pt: &ScalarZnx<DataPt>,
187        sk: &GLWESecretPrepared<DataSk, B>,
188        source_xa: &mut Source,
189        source_xe: &mut Source,
190        scratch: &mut Scratch<B>,
191    ) where
192        Module<B>: GGLWEEncryptSk<B>,
193    {
194        module.gglwe_encrypt_sk(self, pt, sk, source_xa, source_xe, scratch);
195    }
196}