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