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