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, GLWEToLWEKey, LWEInfos, LWESecret, LWESecretToRef, Rank,
 11        prepared::{GLWESecretPrepared, GLWESecretPreparedFactory},
 12    },
 13};
 14
 15impl GLWEToLWEKey<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: GLWEToLWESwitchingKeyEncryptSk<BE>,
 20    {
 21        module.glwe_to_lwe_key_encrypt_sk_tmp_bytes(infos)
 22    }
 23}
 24
 25impl<D: DataMut> GLWEToLWEKey<D> {
 26    pub fn encrypt_sk<M, S1, S2, 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        M: GLWEToLWESwitchingKeyEncryptSk<BE>,
 36        S1: LWESecretToRef,
 37        S2: GLWESecretToRef,
 38        Scratch<BE>: ScratchTakeCore<BE>,
 39    {
 40        module.glwe_to_lwe_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
 41    }
 42}
 43
 44pub trait GLWEToLWESwitchingKeyEncryptSk<BE: Backend> {
 45    fn glwe_to_lwe_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 46    where
 47        A: GGLWEInfos;
 48
 49    fn glwe_to_lwe_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: GLWESecretToRef,
 60        R: GGLWEToMut;
 61}
 62
 63impl<BE: Backend> GLWEToLWESwitchingKeyEncryptSk<BE> for Module<BE>
 64where
 65    Self: ModuleN
 66        + GGLWEEncryptSk<BE>
 67        + GLWESecretPreparedFactory<BE>
 68        + VecZnxAutomorphismInplace<BE>
 69        + VecZnxAutomorphismInplaceTmpBytes,
 70    Scratch<BE>: ScratchTakeCore<BE>,
 71{
 72    fn glwe_to_lwe_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 73    where
 74        A: GGLWEInfos,
 75    {
 76        GLWESecretPrepared::bytes_of(self, infos.rank_in())
 77            + GGLWE::encrypt_sk_tmp_bytes(self, infos)
 78                .max(GLWESecret::bytes_of(self.n().into(), infos.rank_in()) + self.vec_znx_automorphism_inplace_tmp_bytes())
 79    }
 80
 81    fn glwe_to_lwe_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: GGLWEToMut,
 93    {
 94        let sk_lwe: &LWESecret<&[u8]> = &sk_lwe.to_ref();
 95        let sk_glwe: &GLWESecret<&[u8]> = &sk_glwe.to_ref();
 96
 97        assert!(sk_lwe.n().0 <= self.n() as u32);
 98
 99        let (mut sk_lwe_as_glwe_prep, scratch_1) = scratch.take_glwe_secret_prepared(self, Rank(1));
100
101        {
102            let (mut sk_lwe_as_glwe, scratch_2) = scratch_1.take_glwe_secret(self.n().into(), sk_lwe_as_glwe_prep.rank());
103            sk_lwe_as_glwe.dist = sk_lwe.dist;
104            sk_lwe_as_glwe.data.zero();
105            sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n().into()].copy_from_slice(sk_lwe.data.at(0, 0));
106            self.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0, scratch_2);
107            sk_lwe_as_glwe_prep.prepare(self, &sk_lwe_as_glwe);
108        }
109
110        self.gglwe_encrypt_sk(
111            res,
112            &sk_glwe.data,
113            &sk_lwe_as_glwe_prep,
114            source_xa,
115            source_xe,
116            scratch_1,
117        );
118    }
119}