1use poulpy_hal::{
  2    api::{ScratchAvailable, SvpPPolAllocBytes, VecZnxAutomorphism, VecZnxDftAllocBytes, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    TakeGLWESecret,
  9    encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
 10    layouts::{
 11        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos,
 12        compressed::{AutomorphismKeyCompressed, AutomorphismKeyCompressedToMut, GLWESwitchingKeyCompressed},
 13    },
 14};
 15
 16impl AutomorphismKeyCompressed<Vec<u8>> {
 17    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 18    where
 19        A: GGLWEInfos,
 20        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes + SvpPPolAllocBytes,
 21    {
 22        assert_eq!(module.n() as u32, infos.n());
 23        GLWESwitchingKeyCompressed::encrypt_sk_scratch_space(module, infos)
 24            + GLWESecret::alloc_bytes_with(infos.n(), infos.rank_out())
 25    }
 26}
 27
 28pub trait GGLWEAutomorphismKeyCompressedEncryptSk<B: Backend> {
 29    fn gglwe_automorphism_key_compressed_encrypt_sk<R, S>(
 30        &self,
 31        res: &mut R,
 32        p: i64,
 33        sk: &S,
 34        seed_xa: [u8; 32],
 35        source_xe: &mut Source,
 36        scratch: &mut Scratch<B>,
 37    ) where
 38        R: AutomorphismKeyCompressedToMut,
 39        S: GLWESecretToRef;
 40}
 41
 42impl<B: Backend> GGLWEAutomorphismKeyCompressedEncryptSk<B> for Module<B>
 43where
 44    Module<B>:
 45        GGLWEKeyCompressedEncryptSk<B> + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + SvpPPolAllocBytes + VecZnxAutomorphism,
 46    Scratch<B>: TakeGLWESecret + ScratchAvailable,
 47{
 48    fn gglwe_automorphism_key_compressed_encrypt_sk<R, S>(
 49        &self,
 50        res: &mut R,
 51        p: i64,
 52        sk: &S,
 53        seed_xa: [u8; 32],
 54        source_xe: &mut Source,
 55        scratch: &mut Scratch<B>,
 56    ) where
 57        R: AutomorphismKeyCompressedToMut,
 58        S: GLWESecretToRef,
 59    {
 60        let res: &mut AutomorphismKeyCompressed<&mut [u8]> = &mut res.to_mut();
 61        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
 62
 63        #[cfg(debug_assertions)]
 64        {
 65            assert_eq!(res.n(), sk.n());
 66            assert_eq!(res.rank_out(), res.rank_in());
 67            assert_eq!(sk.rank(), res.rank_out());
 68            assert!(
 69                scratch.available() >= AutomorphismKeyCompressed::encrypt_sk_scratch_space(self, res),
 70                "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space: {}",
 71                scratch.available(),
 72                AutomorphismKeyCompressed::encrypt_sk_scratch_space(self, res)
 73            )
 74        }
 75
 76        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
 77
 78        {
 79            (0..res.rank_out().into()).for_each(|i| {
 80                self.vec_znx_automorphism(
 81                    self.galois_element_inv(p),
 82                    &mut sk_out.data.as_vec_znx_mut(),
 83                    i,
 84                    &sk.data.as_vec_znx(),
 85                    i,
 86                );
 87            });
 88        }
 89
 90        self.gglwe_key_compressed_encrypt_sk(&mut res.key, sk, &sk_out, seed_xa, source_xe, scratch_1);
 91
 92        res.p = p;
 93    }
 94}
 95
 96impl<DataSelf: DataMut> AutomorphismKeyCompressed<DataSelf> {
 97    #[allow(clippy::too_many_arguments)]
 98    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
 99        &mut self,
100        module: &Module<B>,
101        p: i64,
102        sk: &GLWESecret<DataSk>,
103        seed_xa: [u8; 32],
104        source_xe: &mut Source,
105        scratch: &mut Scratch<B>,
106    ) where
107        Module<B>: GGLWEAutomorphismKeyCompressedEncryptSk<B>,
108    {
109        module.gglwe_automorphism_key_compressed_encrypt_sk(self, p, sk, seed_xa, source_xe, scratch);
110    }
111}