1use backend::hal::{
  2    api::{
  3        ScratchAvailable, SvpApplyInplace, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
  4        VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform,
  5        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, ZnxZero,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
  8};
  9use sampling::source::Source;
 10
 11use crate::{
 12    TakeGLWEPt,
 13    encryption::glwe_encrypt_sk_internal,
 14    layouts::{GGSWCiphertext, Infos, compressed::GGSWCiphertextCompressed, prepared::GLWESecretPrepared},
 15};
 16
 17impl GGSWCiphertextCompressed<Vec<u8>> {
 18    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
 19    where
 20        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
 21    {
 22        GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k, rank)
 23    }
 24}
 25
 26impl<DataSelf: DataMut> GGSWCiphertextCompressed<DataSelf> {
 27    #[allow(clippy::too_many_arguments)]
 28    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
 29        &mut self,
 30        module: &Module<B>,
 31        pt: &ScalarZnx<DataPt>,
 32        sk: &GLWESecretPrepared<DataSk, B>,
 33        seed_xa: [u8; 32],
 34        source_xe: &mut Source,
 35        sigma: f64,
 36        scratch: &mut Scratch<B>,
 37    ) where
 38        Module<B>: VecZnxAddScalarInplace
 39            + VecZnxDftAllocBytes
 40            + VecZnxBigNormalize<B>
 41            + VecZnxDftFromVecZnx<B>
 42            + SvpApplyInplace<B>
 43            + VecZnxDftToVecZnxBigConsume<B>
 44            + VecZnxNormalizeTmpBytes
 45            + VecZnxFillUniform
 46            + VecZnxSubABInplace
 47            + VecZnxAddInplace
 48            + VecZnxNormalizeInplace<B>
 49            + VecZnxAddNormal
 50            + VecZnxNormalize<B>
 51            + VecZnxSub,
 52        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 53    {
 54        #[cfg(debug_assertions)]
 55        {
 56            use backend::hal::api::ZnxInfos;
 57
 58            assert_eq!(self.rank(), sk.rank());
 59            assert_eq!(self.n(), sk.n());
 60            assert_eq!(pt.n(), sk.n());
 61        }
 62
 63        let basek: usize = self.basek();
 64        let k: usize = self.k();
 65        let rank: usize = self.rank();
 66        let cols: usize = rank + 1;
 67        let digits: usize = self.digits();
 68
 69        let (mut tmp_pt, scratch_1) = scratch.take_glwe_pt(self.n(), basek, k);
 70
 71        let mut source = Source::new(seed_xa);
 72
 73        self.seed = vec![[0u8; 32]; self.rows() * cols];
 74
 75        (0..self.rows()).for_each(|row_i| {
 76            tmp_pt.data.zero();
 77
 78            // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
 79            module.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (digits - 1) + row_i * digits, pt, 0);
 80            module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scratch_1);
 81
 82            (0..rank + 1).for_each(|col_j| {
 83                // rlwe encrypt of vec_znx_pt into vec_znx_ct
 84
 85                let (seed, mut source_xa_tmp) = source.branch();
 86
 87                self.seed[row_i * cols + col_j] = seed;
 88
 89                glwe_encrypt_sk_internal(
 90                    module,
 91                    self.basek(),
 92                    self.k(),
 93                    &mut self.at_mut(row_i, col_j).data,
 94                    cols,
 95                    true,
 96                    Some((&tmp_pt, col_j)),
 97                    sk,
 98                    &mut source_xa_tmp,
 99                    source_xe,
100                    sigma,
101                    scratch_1,
102                );
103            });
104        });
105    }
106}