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