1use poulpy_hal::{
  2    layouts::{Backend, DataMut, Module, Scratch},
  3    source::Source,
  4};
  5
  6use crate::{
  7    encryption::{GLWEEncryptSk, GLWEEncryptSkInternal, SIGMA},
  8    layouts::{
  9        GLWECompressedSeedMut, GLWEInfos, GLWEPlaintextToRef, LWEInfos,
 10        compressed::{GLWECompressed, GLWECompressedToMut},
 11        prepared::GLWESecretPreparedToRef,
 12    },
 13};
 14
 15impl GLWECompressed<Vec<u8>> {
 16    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 17    where
 18        A: GLWEInfos,
 19        M: GLWECompressedEncryptSk<BE>,
 20    {
 21        module.glwe_compressed_encrypt_sk_tmp_bytes(infos)
 22    }
 23}
 24
 25impl<D: DataMut> GLWECompressed<D> {
 26    #[allow(clippy::too_many_arguments)]
 27    pub fn encrypt_sk<M, P, S, BE: Backend>(
 28        &mut self,
 29        module: &M,
 30        pt: &P,
 31        sk: &S,
 32        seed_xa: [u8; 32],
 33        source_xe: &mut Source,
 34        scratch: &mut Scratch<BE>,
 35    ) where
 36        M: GLWECompressedEncryptSk<BE>,
 37        P: GLWEPlaintextToRef,
 38        S: GLWESecretPreparedToRef<BE>,
 39    {
 40        module.glwe_compressed_encrypt_sk(self, pt, sk, seed_xa, source_xe, scratch);
 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 + GLWECompressedSeedMut,
 59        P: GLWEPlaintextToRef,
 60        S: GLWESecretPreparedToRef<BE>;
 61}
 62
 63impl<BE: Backend> GLWECompressedEncryptSk<BE> for Module<BE>
 64where
 65    Self: GLWEEncryptSkInternal<BE> + GLWEEncryptSk<BE>,
 66{
 67    fn glwe_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 68    where
 69        A: GLWEInfos,
 70    {
 71        self.glwe_encrypt_sk_tmp_bytes(infos)
 72    }
 73
 74    fn glwe_compressed_encrypt_sk<R, P, S>(
 75        &self,
 76        res: &mut R,
 77        pt: &P,
 78        sk: &S,
 79        seed_xa: [u8; 32],
 80        source_xe: &mut Source,
 81        scratch: &mut Scratch<BE>,
 82    ) where
 83        R: GLWECompressedToMut + GLWECompressedSeedMut,
 84        P: GLWEPlaintextToRef,
 85        S: GLWESecretPreparedToRef<BE>,
 86    {
 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
107        res.seed_mut().copy_from_slice(&seed_xa);
108    }
109}