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::glwe_switching_key::GLWESwitchingKeyEncryptSk,
 10    layouts::{
 11        GGLWEInfos, GGLWEToMut, GLWESecret, GLWESwitchingKey, GLWESwitchingKeyDegreesMut, LWEInfos, LWESecret, LWESecretToRef,
 12        LWESwitchingKey, Rank,
 13        prepared::{GLWESecretPrepared, GLWESecretPreparedFactory},
 14    },
 15};
 16
 17impl LWESwitchingKey<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: LWESwitchingKeyEncrypt<BE>,
 22    {
 23        module.lwe_switching_key_encrypt_sk_tmp_bytes(infos)
 24    }
 25}
 26
 27impl<D: DataMut> LWESwitchingKey<D> {
 28    pub fn encrypt_sk<S1, S2, M, BE: Backend>(
 29        &mut self,
 30        module: &M,
 31        sk_lwe_in: &S1,
 32        sk_lwe_out: &S2,
 33        source_xa: &mut Source,
 34        source_xe: &mut Source,
 35        scratch: &mut Scratch<BE>,
 36    ) where
 37        S1: LWESecretToRef,
 38        S2: LWESecretToRef,
 39        M: LWESwitchingKeyEncrypt<BE>,
 40    {
 41        module.lwe_switching_key_encrypt_sk(self, sk_lwe_in, sk_lwe_out, source_xa, source_xe, scratch);
 42    }
 43}
 44
 45pub trait LWESwitchingKeyEncrypt<BE: Backend> {
 46    fn lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 47    where
 48        A: GGLWEInfos;
 49
 50    fn lwe_switching_key_encrypt_sk<R, S1, S2>(
 51        &self,
 52        res: &mut R,
 53        sk_lwe_in: &S1,
 54        sk_lwe_out: &S2,
 55        source_xa: &mut Source,
 56        source_xe: &mut Source,
 57        scratch: &mut Scratch<BE>,
 58    ) where
 59        R: GGLWEToMut + GLWESwitchingKeyDegreesMut + GGLWEInfos,
 60        S1: LWESecretToRef,
 61        S2: LWESecretToRef;
 62}
 63
 64impl<BE: Backend> LWESwitchingKeyEncrypt<BE> for Module<BE>
 65where
 66    Self: ModuleN + GLWESwitchingKeyEncryptSk<BE> + GLWESecretPreparedFactory<BE> + VecZnxAutomorphismInplace<BE>,
 67    Scratch<BE>: ScratchTakeCore<BE>,
 68{
 69    fn lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 70    where
 71        A: GGLWEInfos,
 72    {
 73        assert_eq!(infos.dsize().0, 1, "dsize > 1 is not supported for LWESwitchingKey");
 74        assert_eq!(infos.rank_in().0, 1, "rank_in > 1 is not supported for LWESwitchingKey");
 75        assert_eq!(infos.rank_out().0, 1, "rank_out > 1 is not supported for LWESwitchingKey");
 76        GLWESecret::bytes_of(self.n().into(), Rank(1))
 77            + GLWESecretPrepared::bytes_of(self, Rank(1))
 78            + GLWESwitchingKey::encrypt_sk_tmp_bytes(self, infos)
 79    }
 80
 81    #[allow(clippy::too_many_arguments)]
 82    fn lwe_switching_key_encrypt_sk<R, S1, S2>(
 83        &self,
 84        res: &mut R,
 85        sk_lwe_in: &S1,
 86        sk_lwe_out: &S2,
 87        source_xa: &mut Source,
 88        source_xe: &mut Source,
 89        scratch: &mut Scratch<BE>,
 90    ) where
 91        R: GGLWEToMut + GLWESwitchingKeyDegreesMut + GGLWEInfos,
 92        S1: LWESecretToRef,
 93        S2: LWESecretToRef,
 94    {
 95        let sk_lwe_in: &LWESecret<&[u8]> = &sk_lwe_in.to_ref();
 96        let sk_lwe_out: &LWESecret<&[u8]> = &sk_lwe_out.to_ref();
 97
 98        assert!(sk_lwe_in.n().0 <= res.n().0);
 99        assert!(sk_lwe_out.n().0 <= res.n().0);
100        assert!(res.n() <= self.n() as u32);
101
102        let (mut sk_glwe_in, scratch_1) = scratch.take_glwe_secret(self.n().into(), Rank(1));
103        let (mut sk_glwe_out, scratch_2) = scratch_1.take_glwe_secret(self.n().into(), Rank(1));
104
105        sk_glwe_in.dist = sk_lwe_in.dist;
106        sk_glwe_out.dist = sk_lwe_out.dist;
107
108        sk_glwe_out.data.at_mut(0, 0)[..sk_lwe_out.n().into()].copy_from_slice(sk_lwe_out.data.at(0, 0));
109        sk_glwe_out.data.at_mut(0, 0)[sk_lwe_out.n().into()..].fill(0);
110        self.vec_znx_automorphism_inplace(-1, &mut sk_glwe_out.data.as_vec_znx_mut(), 0, scratch_2);
111
112        sk_glwe_in.data.at_mut(0, 0)[..sk_lwe_in.n().into()].copy_from_slice(sk_lwe_in.data.at(0, 0));
113        sk_glwe_in.data.at_mut(0, 0)[sk_lwe_in.n().into()..].fill(0);
114        self.vec_znx_automorphism_inplace(-1, &mut sk_glwe_in.data.as_vec_znx_mut(), 0, scratch_2);
115
116        self.glwe_switching_key_encrypt_sk(res, &sk_glwe_in, &sk_glwe_out, source_xa, source_xe, scratch_2);
117    }
118}