1use 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};
  9use sampling::source::Source;
 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    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
 27        &mut self,
 28        module: &Module<B>,
 29        pt: &GLWEPlaintext<DataPt>,
 30        sk: &GLWESecretPrepared<DataSk, B>,
 31        seed_xa: [u8; 32],
 32        source_xe: &mut Source,
 33        sigma: f64,
 34        scratch: &mut Scratch<B>,
 35    ) where
 36        Module<B>: VecZnxDftAllocBytes
 37            + VecZnxBigNormalize<B>
 38            + VecZnxDftFromVecZnx<B>
 39            + SvpApplyInplace<B>
 40            + VecZnxDftToVecZnxBigConsume<B>
 41            + VecZnxNormalizeTmpBytes
 42            + VecZnxFillUniform
 43            + VecZnxSubABInplace
 44            + VecZnxAddInplace
 45            + VecZnxNormalizeInplace<B>
 46            + VecZnxAddNormal
 47            + VecZnxNormalize<B>
 48            + VecZnxSub,
 49        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 50    {
 51        self.encrypt_sk_internal(
 52            module,
 53            Some((pt, 0)),
 54            sk,
 55            seed_xa,
 56            source_xe,
 57            sigma,
 58            scratch,
 59        );
 60    }
 61
 62    pub(crate) fn encrypt_sk_internal<DataPt: DataRef, DataSk: DataRef, B: Backend>(
 63        &mut self,
 64        module: &Module<B>,
 65        pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
 66        sk: &GLWESecretPrepared<DataSk, B>,
 67        seed_xa: [u8; 32],
 68        source_xe: &mut Source,
 69        sigma: f64,
 70        scratch: &mut Scratch<B>,
 71    ) where
 72        Module<B>: VecZnxDftAllocBytes
 73            + VecZnxBigNormalize<B>
 74            + VecZnxDftFromVecZnx<B>
 75            + SvpApplyInplace<B>
 76            + VecZnxDftToVecZnxBigConsume<B>
 77            + VecZnxNormalizeTmpBytes
 78            + VecZnxFillUniform
 79            + VecZnxSubABInplace
 80            + VecZnxAddInplace
 81            + VecZnxNormalizeInplace<B>
 82            + VecZnxAddNormal
 83            + VecZnxNormalize<B>
 84            + VecZnxSub,
 85        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 86    {
 87        let mut source_xa = Source::new(seed_xa);
 88        let cols: usize = self.rank() + 1;
 89        glwe_encrypt_sk_internal(
 90            module,
 91            self.basek(),
 92            self.k(),
 93            &mut self.data,
 94            cols,
 95            true,
 96            pt,
 97            sk,
 98            &mut source_xa,
 99            source_xe,
100            sigma,
101            scratch,
102        );
103        self.seed = seed_xa;
104    }
105}