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, GLWESecretPrepare, GLWESecretToRef, GLWEToLWESwitchingKey, LWEInfos,
 11        LWESecret, LWESecretToRef, Rank,
 12        prepared::{GLWESecretPrepared, GLWESecretPreparedAlloc},
 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: GLWEToLWESwitchingKeyEncrypt<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: GLWEToLWESwitchingKeyEncrypt<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 GLWEToLWESwitchingKeyEncrypt<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> GLWEToLWESwitchingKeyEncrypt<BE> for Module<BE>
 65where
 66    Self: ModuleN
 67        + GGLWEEncryptSk<BE>
 68        + GLWESecretPreparedAlloc<BE>
 69        + VecZnxAutomorphismInplace<BE>
 70        + VecZnxAutomorphismInplaceTmpBytes
 71        + GLWESecretPrepare<BE>,
 72    Scratch<BE>: ScratchTakeCore<BE>,
 73{
 74    fn glwe_to_lwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 75    where
 76        A: GGLWEInfos,
 77    {
 78        GLWESecretPrepared::bytes_of(self, infos.rank_in())
 79            + GGLWE::encrypt_sk_tmp_bytes(self, infos)
 80                .max(GLWESecret::bytes_of(self.n().into(), infos.rank_in()) + self.vec_znx_automorphism_inplace_tmp_bytes())
 81    }
 82
 83    fn glwe_to_lwe_switching_key_encrypt_sk<R, S1, S2>(
 84        &self,
 85        res: &mut R,
 86        sk_lwe: &S1,
 87        sk_glwe: &S2,
 88        source_xa: &mut Source,
 89        source_xe: &mut Source,
 90        scratch: &mut Scratch<BE>,
 91    ) where
 92        S1: LWESecretToRef,
 93        S2: GLWESecretToRef,
 94        R: GGLWEToMut,
 95    {
 96        let sk_lwe: &LWESecret<&[u8]> = &sk_lwe.to_ref();
 97        let sk_glwe: &GLWESecret<&[u8]> = &sk_glwe.to_ref();
 98
 99        assert!(sk_lwe.n().0 <= self.n() as u32);
100
101        let (mut sk_lwe_as_glwe_prep, scratch_1) = scratch.take_glwe_secret_prepared(self, Rank(1));
102
103        {
104            let (mut sk_lwe_as_glwe, scratch_2) = scratch_1.take_glwe_secret(self, Rank(1));
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}