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