1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxDftBytesOf,
  4        VecZnxNormalizeTmpBytes, VecZnxSwitchRing,
  5    },
  6    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
  7    source::Source,
  8};
  9
 10use crate::{
 11    ScratchTakeCore,
 12    encryption::compressed::gglwe_ct::GGLWECompressedEncryptSk,
 13    layouts::{
 14        GGLWE, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, RingDegree,
 15        compressed::{GLWESwitchingKeyCompressed, GLWESwitchingKeyCompressedToMut},
 16        prepared::GLWESecretPrepared,
 17    },
 18};
 19
 20impl GLWESwitchingKeyCompressed<Vec<u8>> {
 21    pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 22    where
 23        A: GGLWEInfos,
 24        Module<B>:
 25            ModuleN + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
 26    {
 27        (GGLWE::encrypt_sk_tmp_bytes(module, infos) | ScalarZnx::bytes_of(module.n(), 1))
 28            + ScalarZnx::bytes_of(module.n(), infos.rank_in().into())
 29            + GLWESecretPrepared::bytes_of(module, infos.rank_out())
 30    }
 31}
 32
 33impl<DataSelf: DataMut> GLWESwitchingKeyCompressed<DataSelf> {
 34    #[allow(clippy::too_many_arguments)]
 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        scratch: &mut Scratch<B>,
 43    ) where
 44        Module<B>: GGLWEKeyCompressedEncryptSk<B>,
 45    {
 46        module.gglwe_key_compressed_encrypt_sk(self, sk_in, sk_out, seed_xa, source_xe, scratch);
 47    }
 48}
 49
 50pub trait GGLWEKeyCompressedEncryptSk<B: Backend> {
 51    fn gglwe_key_compressed_encrypt_sk<R, SI, SO>(
 52        &self,
 53        res: &mut R,
 54        sk_in: &SI,
 55        sk_out: &SO,
 56        seed_xa: [u8; 32],
 57        source_xe: &mut Source,
 58        scratch: &mut Scratch<B>,
 59    ) where
 60        R: GLWESwitchingKeyCompressedToMut,
 61        SI: GLWESecretToRef,
 62        SO: GLWESecretToRef;
 63}
 64
 65impl<B: Backend> GGLWEKeyCompressedEncryptSk<B> for Module<B>
 66where
 67    Module<B>: ModuleN
 68        + GGLWECompressedEncryptSk<B>
 69        + SvpPPolBytesOf
 70        + VecZnxNormalizeTmpBytes
 71        + VecZnxDftBytesOf
 72        + VecZnxSwitchRing
 73        + SvpPrepare<B>
 74        + SvpPPolAlloc<B>,
 75    Scratch<B>: ScratchAvailable + ScratchTakeBasic + ScratchTakeCore<B>,
 76{
 77    fn gglwe_key_compressed_encrypt_sk<R, SI, SO>(
 78        &self,
 79        res: &mut R,
 80        sk_in: &SI,
 81        sk_out: &SO,
 82        seed_xa: [u8; 32],
 83        source_xe: &mut Source,
 84        scratch: &mut Scratch<B>,
 85    ) where
 86        R: GLWESwitchingKeyCompressedToMut,
 87        SI: GLWESecretToRef,
 88        SO: GLWESecretToRef,
 89    {
 90        let res: &mut GLWESwitchingKeyCompressed<&mut [u8]> = &mut res.to_mut();
 91        let sk_in: &GLWESecret<&[u8]> = &sk_in.to_ref();
 92        let sk_out: &GLWESecret<&[u8]> = &sk_out.to_ref();
 93
 94        #[cfg(debug_assertions)]
 95        {
 96            use crate::layouts::GLWESwitchingKey;
 97
 98            assert!(sk_in.n().0 <= self.n() as u32);
 99            assert!(sk_out.n().0 <= self.n() as u32);
100            assert!(
101                scratch.available() >= GLWESwitchingKey::encrypt_sk_tmp_bytes(self, res),
102                "scratch.available()={} < GLWESwitchingKey::encrypt_sk_tmp_bytes={}",
103                scratch.available(),
104                GLWESwitchingKey::encrypt_sk_tmp_bytes(self, res)
105            )
106        }
107
108        let n: usize = sk_in.n().max(sk_out.n()).into();
109
110        let (mut sk_in_tmp, scratch_1) = scratch.take_scalar_znx(self, sk_in.rank().into());
111        (0..sk_in.rank().into()).for_each(|i| {
112            self.vec_znx_switch_ring(
113                &mut sk_in_tmp.as_vec_znx_mut(),
114                i,
115                &sk_in.data.as_vec_znx(),
116                i,
117            );
118        });
119
120        let (mut sk_out_tmp, scratch_2) = scratch_1.take_glwe_secret_prepared(self, sk_out.rank());
121        {
122            let (mut tmp, _) = scratch_2.take_scalar_znx(self, 1);
123            (0..sk_out.rank().into()).for_each(|i| {
124                self.vec_znx_switch_ring(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
125                self.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
126            });
127        }
128
129        self.gglwe_compressed_encrypt_sk(
130            &mut res.key,
131            &sk_in_tmp,
132            &sk_out_tmp,
133            seed_xa,
134            source_xe,
135            scratch_2,
136        );
137        res.sk_in_n = sk_in.n().into();
138        res.sk_out_n = sk_out.n().into();
139    }
140}