1use poulpy_hal::{
  2    api::{VecZnxAddScalarInplace, VecZnxDftBytesOf, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxZero},
  4    source::Source,
  5};
  6
  7use crate::{
  8    encryption::{SIGMA, glwe_ct::GLWEEncryptSkInternal},
  9    layouts::{
 10        GGSW, GGSWInfos, GLWEInfos, LWEInfos,
 11        compressed::{GGSWCompressed, GGSWCompressedToMut},
 12        prepared::{GLWESecretPrepared, GLWESecretPreparedToRef},
 13    },
 14};
 15
 16impl GGSWCompressed<Vec<u8>> {
 17    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 18    where
 19        A: GGSWInfos,
 20        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftBytesOf,
 21    {
 22        GGSW::encrypt_sk_scratch_space(module, infos)
 23    }
 24}
 25
 26pub trait GGSWCompressedEncryptSk<B: Backend> {
 27    fn ggsw_compressed_encrypt_sk<R, P, S>(
 28        &self,
 29        res: &mut R,
 30        pt: &P,
 31        sk: &S,
 32        seed_xa: [u8; 32],
 33        source_xe: &mut Source,
 34        scratch: &mut Scratch<B>,
 35    ) where
 36        R: GGSWCompressedToMut,
 37        P: ScalarZnxToRef,
 38        S: GLWESecretPreparedToRef<B>;
 39}
 40
 41impl<B: Backend> GGSWCompressedEncryptSk<B> for Module<B>
 42where
 43    Module<B>: GLWEEncryptSkInternal<B> + VecZnxAddScalarInplace + VecZnxNormalizeInplace<B>,
 44    Scratch<B>:,
 45{
 46    fn ggsw_compressed_encrypt_sk<R, P, S>(
 47        &self,
 48        res: &mut R,
 49        pt: &P,
 50        sk: &S,
 51        seed_xa: [u8; 32],
 52        source_xe: &mut Source,
 53        scratch: &mut Scratch<B>,
 54    ) where
 55        R: GGSWCompressedToMut,
 56        P: ScalarZnxToRef,
 57        S: GLWESecretPreparedToRef<B>,
 58    {
 59        let res: &mut GGSWCompressed<&mut [u8]> = &mut res.to_mut();
 60        let sk: &GLWESecretPrepared<&[u8], B> = &sk.to_ref();
 61        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
 62
 63        #[cfg(debug_assertions)]
 64        {
 65            use poulpy_hal::layouts::ZnxInfos;
 66
 67            assert_eq!(res.rank(), sk.rank());
 68            assert_eq!(res.n(), sk.n());
 69            assert_eq!(pt.n() as u32, sk.n());
 70        }
 71
 72        let base2k: usize = res.base2k().into();
 73        let rank: usize = res.rank().into();
 74        let cols: usize = rank + 1;
 75        let dsize: usize = res.dsize().into();
 76
 77        let (mut tmp_pt, scratch_1) = scratch.take_glwe_pt(&res.glwe_layout());
 78
 79        let mut source = Source::new(seed_xa);
 80
 81        res.seed = vec![[0u8; 32]; res.dnum().0 as usize * cols];
 82
 83        for row_i in 0..res.dnum().into() {
 84            tmp_pt.data.zero();
 85
 86            // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
 87            self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, 0);
 88            self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scratch_1);
 89
 90            for col_j in 0..rank + 1 {
 91                // rlwe encrypt of vec_znx_pt into vec_znx_ct
 92
 93                let (seed, mut source_xa_tmp) = source.branch();
 94
 95                res.seed[row_i * cols + col_j] = seed;
 96
 97                self.glwe_encrypt_sk_internal(
 98                    res.base2k().into(),
 99                    res.k().into(),
100                    &mut res.at_mut(row_i, col_j).data,
101                    cols,
102                    true,
103                    Some((&tmp_pt, col_j)),
104                    sk,
105                    &mut source_xa_tmp,
106                    source_xe,
107                    SIGMA,
108                    scratch_1,
109                );
110            }
111        }
112    }
113}
114
115impl<DataSelf: DataMut> GGSWCompressed<DataSelf> {
116    #[allow(clippy::too_many_arguments)]
117    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
118        &mut self,
119        module: &Module<B>,
120        pt: &ScalarZnx<DataPt>,
121        sk: &GLWESecretPrepared<DataSk, B>,
122        seed_xa: [u8; 32],
123        source_xe: &mut Source,
124        scratch: &mut Scratch<B>,
125    ) where
126        Module<B>: GGSWCompressedEncryptSk<B>,
127    {
128        module.ggsw_compressed_encrypt_sk(self, pt, sk, seed_xa, source_xe, scratch);
129    }
130}