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