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