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