1use poulpy_backend::hal::{
  2    api::{
  3        ScratchAvailable, SvpApplyInplace, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxBigNormalize,
  4        VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform, VecZnxNormalize,
  5        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  8    source::Source,
  9};
 10
 11use crate::{
 12    encryption::glwe_ct::glwe_encrypt_sk_internal,
 13    layouts::{GLWECiphertext, GLWEPlaintext, Infos, compressed::GLWECiphertextCompressed, prepared::GLWESecretPrepared},
 14};
 15
 16impl GLWECiphertextCompressed<Vec<u8>> {
 17    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
 18    where
 19        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
 20    {
 21        GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
 22    }
 23}
 24
 25impl<D: DataMut> GLWECiphertextCompressed<D> {
 26    #[allow(clippy::too_many_arguments)]
 27    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
 28        &mut self,
 29        module: &Module<B>,
 30        pt: &GLWEPlaintext<DataPt>,
 31        sk: &GLWESecretPrepared<DataSk, B>,
 32        seed_xa: [u8; 32],
 33        source_xe: &mut Source,
 34        sigma: f64,
 35        scratch: &mut Scratch<B>,
 36    ) where
 37        Module<B>: VecZnxDftAllocBytes
 38            + VecZnxBigNormalize<B>
 39            + VecZnxDftFromVecZnx<B>
 40            + SvpApplyInplace<B>
 41            + VecZnxDftToVecZnxBigConsume<B>
 42            + VecZnxNormalizeTmpBytes
 43            + VecZnxFillUniform
 44            + VecZnxSubABInplace
 45            + VecZnxAddInplace
 46            + VecZnxNormalizeInplace<B>
 47            + VecZnxAddNormal
 48            + VecZnxNormalize<B>
 49            + VecZnxSub,
 50        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 51    {
 52        self.encrypt_sk_internal(
 53            module,
 54            Some((pt, 0)),
 55            sk,
 56            seed_xa,
 57            source_xe,
 58            sigma,
 59            scratch,
 60        );
 61    }
 62
 63    #[allow(clippy::too_many_arguments)]
 64    pub(crate) fn encrypt_sk_internal<DataPt: DataRef, DataSk: DataRef, B: Backend>(
 65        &mut self,
 66        module: &Module<B>,
 67        pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
 68        sk: &GLWESecretPrepared<DataSk, B>,
 69        seed_xa: [u8; 32],
 70        source_xe: &mut Source,
 71        sigma: f64,
 72        scratch: &mut Scratch<B>,
 73    ) where
 74        Module<B>: VecZnxDftAllocBytes
 75            + VecZnxBigNormalize<B>
 76            + VecZnxDftFromVecZnx<B>
 77            + SvpApplyInplace<B>
 78            + VecZnxDftToVecZnxBigConsume<B>
 79            + VecZnxNormalizeTmpBytes
 80            + VecZnxFillUniform
 81            + VecZnxSubABInplace
 82            + VecZnxAddInplace
 83            + VecZnxNormalizeInplace<B>
 84            + VecZnxAddNormal
 85            + VecZnxNormalize<B>
 86            + VecZnxSub,
 87        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 88    {
 89        let mut source_xa = Source::new(seed_xa);
 90        let cols: usize = self.rank() + 1;
 91        glwe_encrypt_sk_internal(
 92            module,
 93            self.basek(),
 94            self.k(),
 95            &mut self.data,
 96            cols,
 97            true,
 98            pt,
 99            sk,
100            &mut source_xa,
101            source_xe,
102            sigma,
103            scratch,
104        );
105        self.seed = seed_xa;
106    }
107}