1use backend::hal::{
  2    api::{ScratchAvailable, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxSwithcDegree},
  3    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    TakeGLWESecretExec,
  9    layouts::{
 10        GGLWECiphertext, GLWESecret, GGLWESwitchingKey, Infos, compressed::GGLWESwitchingKeyCompressed, prepared::GLWESecretExec,
 11    },
 12};
 13
 14use crate::trait_families::{GGLWESwitchingKeyEncryptSkFamily, GLWESecretExecModuleFamily};
 15
 16impl GGLWESwitchingKeyCompressed<Vec<u8>> {
 17    pub fn encrypt_sk_scratch_space<B: Backend>(
 18        module: &Module<B>,
 19        n: usize,
 20        basek: usize,
 21        k: usize,
 22        rank_in: usize,
 23        rank_out: usize,
 24    ) -> usize
 25    where
 26        Module<B>: GGLWESwitchingKeyEncryptSkFamily<B> + GLWESecretExecModuleFamily<B>,
 27    {
 28        (GGLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k) | ScalarZnx::alloc_bytes(n, 1))
 29            + ScalarZnx::alloc_bytes(n, rank_in)
 30            + GLWESecretExec::bytes_of(module, n, rank_out)
 31    }
 32}
 33
 34impl<DataSelf: DataMut> GGLWESwitchingKeyCompressed<DataSelf> {
 35    pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
 36        &mut self,
 37        module: &Module<B>,
 38        sk_in: &GLWESecret<DataSkIn>,
 39        sk_out: &GLWESecret<DataSkOut>,
 40        seed_xa: [u8; 32],
 41        source_xe: &mut Source,
 42        sigma: f64,
 43        scratch: &mut Scratch<B>,
 44    ) where
 45        Module<B>:
 46            GGLWESwitchingKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAddScalarInplace + GLWESecretExecModuleFamily<B>,
 47        Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + ScratchAvailable + TakeVecZnx,
 48    {
 49        #[cfg(debug_assertions)]
 50        {
 51            assert!(sk_in.n() <= module.n());
 52            assert!(sk_out.n() <= module.n());
 53            assert!(
 54                scratch.available()
 55                    >= GGLWESwitchingKey::encrypt_sk_scratch_space(
 56                        module,
 57                        sk_out.n(),
 58                        self.basek(),
 59                        self.k(),
 60                        self.rank_in(),
 61                        self.rank_out()
 62                    ),
 63                "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
 64                scratch.available(),
 65                GGLWESwitchingKey::encrypt_sk_scratch_space(
 66                    module,
 67                    sk_out.n(),
 68                    self.basek(),
 69                    self.k(),
 70                    self.rank_in(),
 71                    self.rank_out()
 72                )
 73            )
 74        }
 75
 76        let n: usize = sk_in.n().max(sk_out.n());
 77
 78        let (mut sk_in_tmp, scratch1) = scratch.take_scalar_znx(n, sk_in.rank());
 79        (0..sk_in.rank()).for_each(|i| {
 80            module.vec_znx_switch_degree(
 81                &mut sk_in_tmp.as_vec_znx_mut(),
 82                i,
 83                &sk_in.data.as_vec_znx(),
 84                i,
 85            );
 86        });
 87
 88        let (mut sk_out_tmp, scratch2) = scratch1.take_glwe_secret_exec(n, sk_out.rank());
 89        {
 90            let (mut tmp, _) = scratch2.take_scalar_znx(n, 1);
 91            (0..sk_out.rank()).for_each(|i| {
 92                module.vec_znx_switch_degree(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
 93                module.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
 94            });
 95        }
 96
 97        self.key.encrypt_sk(
 98            module,
 99            &sk_in_tmp,
100            &sk_out_tmp,
101            seed_xa,
102            source_xe,
103            sigma,
104            scratch2,
105        );
106        self.sk_in_n = sk_in.n();
107        self.sk_out_n = sk_out.n();
108    }
109}