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