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