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