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    TakeGLWESecretPrepared,
  9    layouts::{GGLWECiphertext, GGLWESwitchingKey, GLWESecret, Infos, prepared::GLWESecretPrepared},
 10};
 11
 12use crate::trait_families::{GGLWESwitchingKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
 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> + GLWESecretPreparedModuleFamily<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            + GLWESecretPrepared::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 + GLWESecretPreparedModuleFamily<B>,
 56        Scratch<B>:
 57            ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretPrepared<B> + ScratchAvailable + TakeVecZnx,
 58    {
 59        #[cfg(debug_assertions)]
 60        {
 61            assert!(sk_in.n() <= module.n());
 62            assert!(sk_out.n() <= module.n());
 63            assert!(
 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                "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
 74                scratch.available(),
 75                GGLWESwitchingKey::encrypt_sk_scratch_space(
 76                    module,
 77                    sk_out.n(),
 78                    self.basek(),
 79                    self.k(),
 80                    self.rank_in(),
 81                    self.rank_out()
 82                )
 83            )
 84        }
 85
 86        let n: usize = sk_in.n().max(sk_out.n());
 87
 88        let (mut sk_in_tmp, scratch1) = scratch.take_scalar_znx(n, sk_in.rank());
 89        (0..sk_in.rank()).for_each(|i| {
 90            module.vec_znx_switch_degree(
 91                &mut sk_in_tmp.as_vec_znx_mut(),
 92                i,
 93                &sk_in.data.as_vec_znx(),
 94                i,
 95            );
 96        });
 97
 98        let (mut sk_out_tmp, scratch2) = scratch1.take_glwe_secret_prepared(n, sk_out.rank());
 99        {
100            let (mut tmp, _) = scratch2.take_scalar_znx(n, 1);
101            (0..sk_out.rank()).for_each(|i| {
102                module.vec_znx_switch_degree(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
103                module.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
104            });
105        }
106
107        self.key.encrypt_sk(
108            module,
109            &sk_in_tmp,
110            &sk_out_tmp,
111            source_xa,
112            source_xe,
113            sigma,
114            scratch2,
115        );
116        self.sk_in_n = sk_in.n();
117        self.sk_out_n = sk_out.n();
118    }
119}