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