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)
 89            + GLWEPlaintext::bytes_of_from_infos(self, infos).max(self.vec_znx_normalize_tmp_bytes())
 90    }
 91
 92    fn gglwe_encrypt_sk<R, P, S>(
 93        &self,
 94        res: &mut R,
 95        pt: &P,
 96        sk: &S,
 97        source_xa: &mut Source,
 98        source_xe: &mut Source,
 99        scratch: &mut Scratch<BE>,
100    ) where
101        R: GGLWEToMut,
102        P: ScalarZnxToRef,
103        S: GLWESecretPreparedToRef<BE>,
104    {
105        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
106        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
107        let sk: &GLWESecretPrepared<&[u8], BE> = &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() >= self.gglwe_encrypt_sk_tmp_bytes(res),
127            "scratch.available: {} < GGLWE::encrypt_sk_tmp_bytes(self, res.rank()={}, res.size()={}): {}",
128            scratch.available(),
129            res.rank_out(),
130            res.size(),
131            self.gglwe_encrypt_sk_tmp_bytes(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        let dnum: usize = res.dnum().into();
144        let dsize: usize = res.dsize().into();
145        let base2k: usize = res.base2k().into();
146        let rank_in: usize = res.rank_in().into();
147
148        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, &res.glwe_layout());
149        // For each input column (i.e. rank) produces a GGLWE of rank_out+1 columns
150        //
151        // Example for ksk rank 2 to rank 3:
152        //
153        // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
154        // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
155        //
156        // Example ksk rank 2 to rank 1
157        //
158        // (-(a*s) + s0, a)
159        // (-(b*s) + s1, b)
160
161        for col_i in 0..rank_in {
162            for row_i in 0..dnum {
163                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
164                tmp_pt.data.zero(); // zeroes for next iteration
165                self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, col_i);
166                self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
167                self.glwe_encrypt_sk(
168                    &mut res.at_mut(row_i, col_i),
169                    &tmp_pt,
170                    sk,
171                    source_xa,
172                    source_xe,
173                    scrach_1,
174                );
175            }
176        }
177    }
178}