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