1use backend::hal::{
 2    api::{ScratchAvailable, TakeVecZnx, TakeVecZnxDft},
 3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 4};
 5use sampling::source::Source;
 6
 7use crate::{
 8    encryption::glwe_ct::glwe_encrypt_sk_internal,
 9    layouts::{GLWECiphertext, GLWEPlaintext, Infos, compressed::GLWECiphertextCompressed, prepared::GLWESecretPrepared},
10};
11
12use crate::trait_families::GLWEEncryptSkFamily;
13
14impl GLWECiphertextCompressed<Vec<u8>> {
15    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
16    where
17        Module<B>: GLWEEncryptSkFamily<B>,
18    {
19        GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
20    }
21}
22
23impl<D: DataMut> GLWECiphertextCompressed<D> {
24    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
25        &mut self,
26        module: &Module<B>,
27        pt: &GLWEPlaintext<DataPt>,
28        sk: &GLWESecretPrepared<DataSk, B>,
29        seed_xa: [u8; 32],
30        source_xe: &mut Source,
31        sigma: f64,
32        scratch: &mut Scratch<B>,
33    ) where
34        Module<B>: GLWEEncryptSkFamily<B>,
35        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
36    {
37        self.encrypt_sk_internal(
38            module,
39            Some((pt, 0)),
40            sk,
41            seed_xa,
42            source_xe,
43            sigma,
44            scratch,
45        );
46    }
47
48    pub(crate) fn encrypt_sk_internal<DataPt: DataRef, DataSk: DataRef, B: Backend>(
49        &mut self,
50        module: &Module<B>,
51        pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
52        sk: &GLWESecretPrepared<DataSk, B>,
53        seed_xa: [u8; 32],
54        source_xe: &mut Source,
55        sigma: f64,
56        scratch: &mut Scratch<B>,
57    ) where
58        Module<B>: GLWEEncryptSkFamily<B>,
59        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
60    {
61        let mut source_xa = Source::new(seed_xa);
62        let cols: usize = self.rank() + 1;
63        glwe_encrypt_sk_internal(
64            module,
65            self.basek(),
66            self.k(),
67            &mut self.data,
68            cols,
69            true,
70            pt,
71            sk,
72            &mut source_xa,
73            source_xe,
74            sigma,
75            scratch,
76        );
77        self.seed = seed_xa;
78    }
79}