1use backend::hal::{
 2    api::{
 3        ScratchAvailable, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxAutomorphism,
 4        VecZnxSwithcDegree,
 5    },
 6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 7};
 8use sampling::source::Source;
 9
10use crate::{
11    TakeGLWESecret, TakeGLWESecretExec,
12    layouts::{
13        GLWESecret, Infos,
14        compressed::{GGLWEAutomorphismKeyCompressed, GGLWESwitchingKeyCompressed},
15    },
16};
17
18use crate::trait_families::{GGLWEAutomorphismKeyEncryptSkFamily, GLWESecretExecModuleFamily};
19
20impl GGLWEAutomorphismKeyCompressed<Vec<u8>> {
21    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
22    where
23        Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B> + GLWESecretExecModuleFamily<B>,
24    {
25        GGLWESwitchingKeyCompressed::encrypt_sk_scratch_space(module, n, basek, k, rank, rank) + GLWESecret::bytes_of(n, rank)
26    }
27}
28
29impl<DataSelf: DataMut> GGLWEAutomorphismKeyCompressed<DataSelf> {
30    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
31        &mut self,
32        module: &Module<B>,
33        p: i64,
34        sk: &GLWESecret<DataSk>,
35        seed_xa: [u8; 32],
36        source_xe: &mut Source,
37        sigma: f64,
38        scratch: &mut Scratch<B>,
39    ) where
40        Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
41            + VecZnxSwithcDegree
42            + VecZnxAutomorphism
43            + VecZnxAddScalarInplace
44            + GLWESecretExecModuleFamily<B>,
45        Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
46    {
47        #[cfg(debug_assertions)]
48        {
49            assert_eq!(self.n(), sk.n());
50            assert_eq!(self.rank_out(), self.rank_in());
51            assert_eq!(sk.rank(), self.rank());
52            assert!(
53                scratch.available()
54                    >= GGLWEAutomorphismKeyCompressed::encrypt_sk_scratch_space(
55                        module,
56                        sk.n(),
57                        self.basek(),
58                        self.k(),
59                        self.rank()
60                    ),
61                "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
62                scratch.available(),
63                self.rank(),
64                self.size(),
65                GGLWEAutomorphismKeyCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank())
66            )
67        }
68
69        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
70
71        {
72            (0..self.rank()).for_each(|i| {
73                module.vec_znx_automorphism(
74                    module.galois_element_inv(p),
75                    &mut sk_out.data.as_vec_znx_mut(),
76                    i,
77                    &sk.data.as_vec_znx(),
78                    i,
79                );
80            });
81        }
82
83        self.key
84            .encrypt_sk(module, &sk, &sk_out, seed_xa, source_xe, sigma, scratch_1);
85
86        self.p = p;
87    }
88}