1use poulpy_hal::{
  2    api::{ModuleN, VecZnxAutomorphismInplace, VecZnxAutomorphismInplaceTmpBytes},
  3    layouts::{Backend, DataMut, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GGLWEEncryptSk, ScratchTakeCore,
  9    layouts::{
 10        GGLWE, GGLWEInfos, GGLWEToMut, GLWESecret, GLWESecretToRef, GLWEToLWESwitchingKey, LWEInfos, LWESecret, LWESecretToRef,
 11        Rank,
 12        prepared::{GLWESecretPrepared, GLWESecretPreparedFactory},
 13    },
 14};
 15
 16impl GLWEToLWESwitchingKey<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: GLWEToLWESwitchingKeyEncryptSk<BE>,
 21    {
 22        module.glwe_to_lwe_switching_key_encrypt_sk_tmp_bytes(infos)
 23    }
 24}
 25
 26impl<D: DataMut> GLWEToLWESwitchingKey<D> {
 27    pub fn encrypt_sk<M, S1, S2, 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        M: GLWEToLWESwitchingKeyEncryptSk<BE>,
 37        S1: LWESecretToRef,
 38        S2: GLWESecretToRef,
 39        Scratch<BE>: ScratchTakeCore<BE>,
 40    {
 41        module.glwe_to_lwe_switching_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
 42    }
 43}
 44
 45pub trait GLWEToLWESwitchingKeyEncryptSk<BE: Backend> {
 46    fn glwe_to_lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 47    where
 48        A: GGLWEInfos;
 49
 50    fn glwe_to_lwe_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: GGLWEToMut;
 62}
 63
 64impl<BE: Backend> GLWEToLWESwitchingKeyEncryptSk<BE> for Module<BE>
 65where
 66    Self: ModuleN
 67        + GGLWEEncryptSk<BE>
 68        + GLWESecretPreparedFactory<BE>
 69        + VecZnxAutomorphismInplace<BE>
 70        + VecZnxAutomorphismInplaceTmpBytes,
 71    Scratch<BE>: ScratchTakeCore<BE>,
 72{
 73    fn glwe_to_lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 74    where
 75        A: GGLWEInfos,
 76    {
 77        GLWESecretPrepared::bytes_of(self, infos.rank_in())
 78            + GGLWE::encrypt_sk_tmp_bytes(self, infos)
 79                .max(GLWESecret::bytes_of(self.n().into(), infos.rank_in()) + self.vec_znx_automorphism_inplace_tmp_bytes())
 80    }
 81
 82    fn glwe_to_lwe_switching_key_encrypt_sk<R, S1, S2>(
 83        &self,
 84        res: &mut R,
 85        sk_lwe: &S1,
 86        sk_glwe: &S2,
 87        source_xa: &mut Source,
 88        source_xe: &mut Source,
 89        scratch: &mut Scratch<BE>,
 90    ) where
 91        S1: LWESecretToRef,
 92        S2: GLWESecretToRef,
 93        R: GGLWEToMut,
 94    {
 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_prep, scratch_1) = scratch.take_glwe_secret_prepared(self, Rank(1));
101
102        {
103            let (mut sk_lwe_as_glwe, scratch_2) = scratch_1.take_glwe_secret(self.n().into(), sk_lwe_as_glwe_prep.rank());
104            sk_lwe_as_glwe.dist = sk_lwe.dist;
105            sk_lwe_as_glwe.data.zero();
106            sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n().into()].copy_from_slice(sk_lwe.data.at(0, 0));
107            self.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0, scratch_2);
108            sk_lwe_as_glwe_prep.prepare(self, &sk_lwe_as_glwe);
109        }
110
111        self.gglwe_encrypt_sk(
112            res,
113            &sk_glwe.data,
114            &sk_lwe_as_glwe_prep,
115            source_xa,
116            source_xe,
117            scratch_1,
118        );
119    }
120}