1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftApply,
  5        VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub,
  6        VecZnxSubInplace, VecZnxSwitchRing,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
  9    source::Source,
 10};
 11
 12use crate::{
 13    TakeGLWESecretPrepared,
 14    layouts::{Degree, GGLWE, GGLWEInfos, GLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, prepared::GLWESecretPrepared},
 15};
 16
 17impl GLWESwitchingKey<Vec<u8>> {
 18    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 19    where
 20        A: GGLWEInfos,
 21        Module<B>: SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
 22    {
 23        (GGLWE::encrypt_sk_scratch_space(module, infos) | ScalarZnx::bytes_of(module.n(), 1))
 24            + ScalarZnx::bytes_of(module.n(), infos.rank_in().into())
 25            + GLWESecretPrepared::bytes_of_from_infos(module, &infos.glwe_layout())
 26    }
 27
 28    pub fn encrypt_pk_scratch_space<B: Backend, A>(module: &Module<B>, _infos: &A) -> usize
 29    where
 30        A: GGLWEInfos,
 31    {
 32        GGLWE::encrypt_pk_scratch_space(module, _infos)
 33    }
 34}
 35
 36impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
 37    #[allow(clippy::too_many_arguments)]
 38    pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
 39        &mut self,
 40        module: &Module<B>,
 41        sk_in: &GLWESecret<DataSkIn>,
 42        sk_out: &GLWESecret<DataSkOut>,
 43        source_xa: &mut Source,
 44        source_xe: &mut Source,
 45        scratch: &mut Scratch<B>,
 46    ) where
 47        Module<B>: VecZnxAddScalarInplace
 48            + VecZnxDftAllocBytes
 49            + VecZnxBigNormalize<B>
 50            + VecZnxDftApply<B>
 51            + SvpApplyDftToDftInplace<B>
 52            + VecZnxIdftApplyConsume<B>
 53            + VecZnxNormalizeTmpBytes
 54            + VecZnxFillUniform
 55            + VecZnxSubInplace
 56            + VecZnxAddInplace
 57            + VecZnxNormalizeInplace<B>
 58            + VecZnxAddNormal
 59            + VecZnxNormalize<B>
 60            + VecZnxSub
 61            + SvpPrepare<B>
 62            + VecZnxSwitchRing
 63            + SvpPPolAllocBytes,
 64        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
 65    {
 66        #[cfg(debug_assertions)]
 67        {
 68            assert!(sk_in.n().0 <= module.n() as u32);
 69            assert!(sk_out.n().0 <= module.n() as u32);
 70            assert!(
 71                scratch.available() >= GLWESwitchingKey::encrypt_sk_scratch_space(module, self),
 72                "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
 73                scratch.available(),
 74                GLWESwitchingKey::encrypt_sk_scratch_space(module, self)
 75            )
 76        }
 77
 78        let n: usize = sk_in.n().max(sk_out.n()).into();
 79
 80        let (mut sk_in_tmp, scratch_1) = scratch.take_scalar_znx(n, sk_in.rank().into());
 81        (0..sk_in.rank().into()).for_each(|i| {
 82            module.vec_znx_switch_ring(
 83                &mut sk_in_tmp.as_vec_znx_mut(),
 84                i,
 85                &sk_in.data.as_vec_znx(),
 86                i,
 87            );
 88        });
 89
 90        let (mut sk_out_tmp, scratch_2) = scratch_1.take_glwe_secret_prepared(Degree(n as u32), sk_out.rank());
 91        {
 92            let (mut tmp, _) = scratch_2.take_scalar_znx(n, 1);
 93            (0..sk_out.rank().into()).for_each(|i| {
 94                module.vec_znx_switch_ring(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
 95                module.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
 96            });
 97        }
 98
 99        self.key.encrypt_sk(
100            module,
101            &sk_in_tmp,
102            &sk_out_tmp,
103            source_xa,
104            source_xe,
105            scratch_2,
106        );
107        self.sk_in_n = sk_in.n().into();
108        self.sk_out_n = sk_out.n().into();
109    }
110}