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