1use poulpy_hal::{
  2    api::{ModuleN, VecZnxAutomorphismInplace},
  3    layouts::{Backend, DataMut, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
  4    source::Source,
  5};
  6
  7use crate::{
  8    ScratchTakeCore,
  9    encryption::gglwe_ksk::GLWESwitchingKeyEncryptSk,
 10    layouts::{
 11        GGLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, GLWEToLWESwitchingKey, GLWEToLWESwitchingKeyToMut, LWEInfos,
 12        LWESecret, LWESecretToRef, Rank,
 13        prepared::{GLWESecretPrepared, GLWESecretPreparedAlloc},
 14    },
 15};
 16
 17impl GLWEToLWESwitchingKey<Vec<u8>> {
 18    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 19    where
 20        A: GGLWEInfos,
 21        M: GLWEToLWESwitchingKeyEncrypt<BE>,
 22    {
 23        module.glwe_to_lwe_switching_key_encrypt_sk_tmp_bytes(infos)
 24    }
 25}
 26
 27impl<D: DataMut> GLWEToLWESwitchingKey<D> {
 28    pub fn encrypt_sk<M, S1, S2, BE: Backend>(
 29        &mut self,
 30        module: &M,
 31        sk_lwe: &S1,
 32        sk_glwe: &S2,
 33        source_xa: &mut Source,
 34        source_xe: &mut Source,
 35        scratch: &mut Scratch<BE>,
 36    ) where
 37        M: GLWEToLWESwitchingKeyEncrypt<BE>,
 38        S1: LWESecretToRef,
 39        S2: GLWESecretToRef,
 40        Scratch<BE>: ScratchTakeCore<BE>,
 41    {
 42        module.glwe_to_lwe_switching_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
 43    }
 44}
 45
 46pub trait GLWEToLWESwitchingKeyEncrypt<BE: Backend> {
 47    fn glwe_to_lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 48    where
 49        A: GGLWEInfos;
 50
 51    fn glwe_to_lwe_switching_key_encrypt_sk<R, S1, S2>(
 52        &self,
 53        res: &mut R,
 54        sk_lwe: &S1,
 55        sk_glwe: &S2,
 56        source_xa: &mut Source,
 57        source_xe: &mut Source,
 58        scratch: &mut Scratch<BE>,
 59    ) where
 60        S1: LWESecretToRef,
 61        S2: GLWESecretToRef,
 62        R: GLWEToLWESwitchingKeyToMut;
 63}
 64
 65impl<BE: Backend> GLWEToLWESwitchingKeyEncrypt<BE> for Module<BE>
 66where
 67    Self: ModuleN + GLWESwitchingKeyEncryptSk<BE> + GLWESecretPreparedAlloc<BE> + VecZnxAutomorphismInplace<BE>,
 68    Scratch<BE>: ScratchTakeCore<BE>,
 69{
 70    fn glwe_to_lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 71    where
 72        A: GGLWEInfos,
 73    {
 74        GLWESecretPrepared::bytes_of(self, infos.rank_in())
 75            + (GLWESwitchingKey::encrypt_sk_tmp_bytes(self, infos) | GLWESecret::bytes_of(self.n().into(), infos.rank_in()))
 76    }
 77
 78    fn glwe_to_lwe_switching_key_encrypt_sk<R, S1, S2>(
 79        &self,
 80        res: &mut R,
 81        sk_lwe: &S1,
 82        sk_glwe: &S2,
 83        source_xa: &mut Source,
 84        source_xe: &mut Source,
 85        scratch: &mut Scratch<BE>,
 86    ) where
 87        S1: LWESecretToRef,
 88        S2: GLWESecretToRef,
 89        R: GLWEToLWESwitchingKeyToMut,
 90    {
 91        let res: &mut GLWEToLWESwitchingKey<&mut [u8]> = &mut res.to_mut();
 92        let sk_lwe: &LWESecret<&[u8]> = &sk_lwe.to_ref();
 93        let sk_glwe: &GLWESecret<&[u8]> = &sk_glwe.to_ref();
 94
 95        assert!(sk_lwe.n().0 <= self.n() as u32);
 96
 97        let (mut sk_lwe_as_glwe, scratch_1) = scratch.take_glwe_secret(self, Rank(1));
 98        sk_lwe_as_glwe.data.zero();
 99        sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n().into()].copy_from_slice(sk_lwe.data.at(0, 0));
100        self.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0, scratch_1);
101
102        res.0.encrypt_sk(
103            self,
104            sk_glwe,
105            &sk_lwe_as_glwe,
106            source_xa,
107            source_xe,
108            scratch_1,
109        );
110    }
111}