1use poulpy_hal::{
  2    api::{ScratchAvailable, VecZnxAutomorphism},
  3    layouts::{Backend, DataMut, GaloisElement, Module, Scratch},
  4    source::Source,
  5};
  6
  7use crate::{
  8    ScratchTakeCore,
  9    encryption::gglwe_ksk::GLWESwitchingKeyEncryptSk,
 10    layouts::{AutomorphismKey, AutomorphismKeyToMut, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos},
 11};
 12
 13impl AutomorphismKey<Vec<u8>> {
 14    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 15    where
 16        A: GGLWEInfos,
 17        M: AutomorphismKeyEncryptSk<BE>,
 18    {
 19        module.automorphism_key_encrypt_sk_tmp_bytes(infos)
 20    }
 21
 22    pub fn encrypt_pk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 23    where
 24        A: GGLWEInfos,
 25        M: GGLWEAutomorphismKeyEncryptPk<BE>,
 26    {
 27        module.automorphism_key_encrypt_pk_tmp_bytes(infos)
 28    }
 29}
 30
 31impl<DM: DataMut> AutomorphismKey<DM>
 32where
 33    Self: AutomorphismKeyToMut,
 34{
 35    pub fn encrypt_sk<S, M, BE: Backend>(
 36        &mut self,
 37        module: &M,
 38        p: i64,
 39        sk: &S,
 40        source_xa: &mut Source,
 41        source_xe: &mut Source,
 42        scratch: &mut Scratch<BE>,
 43    ) where
 44        S: GLWESecretToRef,
 45        M: AutomorphismKeyEncryptSk<BE>,
 46    {
 47        module.automorphism_key_encrypt_sk(self, p, sk, source_xa, source_xe, scratch);
 48    }
 49}
 50
 51pub trait AutomorphismKeyEncryptSk<BE: Backend> {
 52    fn automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 53    where
 54        A: GGLWEInfos;
 55
 56    fn automorphism_key_encrypt_sk<R, S>(
 57        &self,
 58        res: &mut R,
 59        p: i64,
 60        sk: &S,
 61        source_xa: &mut Source,
 62        source_xe: &mut Source,
 63        scratch: &mut Scratch<BE>,
 64    ) where
 65        R: AutomorphismKeyToMut,
 66        S: GLWESecretToRef;
 67}
 68
 69impl<BE: Backend> AutomorphismKeyEncryptSk<BE> for Module<BE>
 70where
 71    Self: GLWESwitchingKeyEncryptSk<BE> + VecZnxAutomorphism + GaloisElement,
 72    Scratch<BE>: ScratchTakeCore<BE>,
 73{
 74    fn automorphism_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 75    where
 76        A: GGLWEInfos,
 77    {
 78        assert_eq!(
 79            infos.rank_in(),
 80            infos.rank_out(),
 81            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
 82        );
 83        self.glwe_switching_key_encrypt_sk_tmp_bytes(infos) + GLWESecret::bytes_of_from_infos(self, infos)
 84    }
 85
 86    fn automorphism_key_encrypt_sk<R, S>(
 87        &self,
 88        res: &mut R,
 89        p: i64,
 90        sk: &S,
 91        source_xa: &mut Source,
 92        source_xe: &mut Source,
 93        scratch: &mut Scratch<BE>,
 94    ) where
 95        R: AutomorphismKeyToMut,
 96        S: GLWESecretToRef,
 97    {
 98        let res: &mut AutomorphismKey<&mut [u8]> = &mut res.to_mut();
 99        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
100
101        assert_eq!(res.n(), sk.n());
102        assert_eq!(res.rank_out(), res.rank_in());
103        assert_eq!(sk.rank(), res.rank_out());
104        assert!(
105            scratch.available() >= self.automorphism_key_encrypt_sk_tmp_bytes(res),
106            "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {:?}",
107            scratch.available(),
108            self.automorphism_key_encrypt_sk_tmp_bytes(res)
109        );
110
111        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(self, sk.rank());
112
113        {
114            (0..res.rank_out().into()).for_each(|i| {
115                self.vec_znx_automorphism(
116                    self.galois_element_inv(p),
117                    &mut sk_out.data.as_vec_znx_mut(),
118                    i,
119                    &sk.data.as_vec_znx(),
120                    i,
121                );
122            });
123        }
124
125        res.key
126            .encrypt_sk(self, sk, &sk_out, source_xa, source_xe, scratch_1);
127
128        res.p = p;
129    }
130}
131
132pub trait GGLWEAutomorphismKeyEncryptPk<BE: Backend> {
133    fn automorphism_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
134    where
135        A: GGLWEInfos;
136}
137
138impl<BE: Backend> GGLWEAutomorphismKeyEncryptPk<BE> for Module<BE>
139where
140    Self:,
141    Scratch<BE>: ScratchTakeCore<BE>,
142{
143    fn automorphism_key_encrypt_pk_tmp_bytes<A>(&self, _infos: &A) -> usize
144    where
145        A: GGLWEInfos,
146    {
147        unimplemented!()
148    }
149}