1use poulpy_hal::{
  2    api::{VecZnxDftBytesOf, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    encryption::{SIGMA, glwe_ct::{GLWEEncryptSk, GLWEEncryptSkInternal}},
  9    layouts::{
 10        GLWE, GLWEInfos, GLWEPlaintext, GLWEPlaintextToRef, LWEInfos,
 11        compressed::{GLWECompressed, GLWECompressedToMut},
 12        prepared::{GLWESecretPrepared, GLWESecretPreparedToRef},
 13    },
 14};
 15
 16impl GLWECompressed<Vec<u8>> {
 17    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 18    where
 19        A: GLWEInfos,
 20        M: GLWECompressedEncryptSk<BE>,
 21    {
 22        module.glwe_compressed_encrypt_sk_tmp_bytes(infos)
 23    }
 24}
 25
 26impl<D: DataMut> GLWECompressed<D> {
 27    #[allow(clippy::too_many_arguments)]
 28    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, BE: Backend>(
 29        &mut self,
 30        module: &Module<BE>,
 31        pt: &GLWEPlaintext<DataPt>,
 32        sk: &GLWESecretPrepared<DataSk, BE>,
 33        seed_xa: [u8; 32],
 34        source_xe: &mut Source,
 35        scratch: &mut Scratch<BE>,
 36    ) where
 37        Module<BE>: GLWECompressedEncryptSk<BE>,
 38    {
 39        module.glwe_compressed_encrypt_sk(self, pt, sk, seed_xa, source_xe, scratch);
 40    }
 41}
 42
 43
 44pub trait GLWECompressedEncryptSk<BE: Backend> {
 45    fn glwe_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 46    where
 47        A: GLWEInfos;
 48
 49    fn glwe_compressed_encrypt_sk<R, P, S>(
 50        &self,
 51        res: &mut R,
 52        pt: &P,
 53        sk: &S,
 54        seed_xa: [u8; 32],
 55        source_xe: &mut Source,
 56        scratch: &mut Scratch<BE>,
 57    ) where
 58        R: GLWECompressedToMut,
 59        P: GLWEPlaintextToRef,
 60        S: GLWESecretPreparedToRef<BE>;
 61}
 62
 63impl<BE: Backend> GLWECompressedEncryptSk<BE> for Module<BE>
 64where
 65    Module<BE>: GLWEEncryptSkInternal<BE> + GLWEEncryptSk<BE>,
 66{
 67
 68    fn glwe_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 69    where
 70        A: GLWEInfos,
 71    {
 72        self.glwe_encrypt_sk_tmp_bytes(infos)
 73    }
 74
 75    fn glwe_compressed_encrypt_sk<R, P, S>(
 76        &self,
 77        res: &mut R,
 78        pt: &P,
 79        sk: &S,
 80        seed_xa: [u8; 32],
 81        source_xe: &mut Source,
 82        scratch: &mut Scratch<BE>,
 83    ) where
 84        R: GLWECompressedToMut,
 85        P: GLWEPlaintextToRef,
 86        S: GLWESecretPreparedToRef<BE>,
 87    {
 88        let res: &mut GLWECompressed<&mut [u8]> = &mut res.to_mut();
 89        let mut source_xa: Source = Source::new(seed_xa);
 90        let cols: usize = (res.rank() + 1).into();
 91
 92        self.glwe_encrypt_sk_internal(
 93            res.base2k().into(),
 94            res.k().into(),
 95            &mut res.data,
 96            cols,
 97            true,
 98            Some((pt, 0)),
 99            sk,
100            &mut source_xa,
101            source_xe,
102            SIGMA,
103            scratch,
104        );
105
106        res.seed = seed_xa;
107    }
108}