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