1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolAlloc, 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::{
 13    ScratchTakeCore,
 14    layouts::{
 15        AutomorphismKey, AutomorphismKeyToMut, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, LWEInfos,
 16    },
 17};
 18
 19impl AutomorphismKey<Vec<u8>> {
 20    pub fn encrypt_sk_tmp_bytes<BE: Backend, A>(module: &Module<BE>, infos: &A) -> usize
 21    where
 22        A: GGLWEInfos,
 23        Module<BE>: ModuleN + SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolAlloc<BE>,
 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_tmp_bytes(module, infos) + GLWESecret::bytes_of_from_infos(module, &infos.glwe_layout())
 31    }
 32
 33    pub fn encrypt_pk_tmp_bytes<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_tmp_bytes(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>: ModuleN 
 83        + VecZnxAddScalarInplace
 84        + VecZnxDftBytesOf
 85        + VecZnxBigNormalize<BE>
 86        + VecZnxDftApply<BE>
 87        + SvpApplyDftToDftInplace<BE>
 88        + VecZnxIdftApplyConsume<BE>
 89        + VecZnxNormalizeTmpBytes
 90        + VecZnxFillUniform
 91        + VecZnxSubInplace
 92        + VecZnxAddInplace
 93        + VecZnxNormalizeInplace<BE>
 94        + VecZnxAddNormal
 95        + VecZnxNormalize<BE>
 96        + VecZnxSub
 97        + SvpPrepare<BE>
 98        + VecZnxSwitchRing
 99        + SvpPPolBytesOf
100        + VecZnxAutomorphism
101        + SvpPPolAlloc<BE>
102        + SvpPPolBytesOf,
103    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
104{
105    fn gglwe_automorphism_key_encrypt_sk<A, B>(
106        &self,
107        res: &mut A,
108        p: i64,
109        sk: &B,
110        source_xa: &mut Source,
111        source_xe: &mut Source,
112        scratch: &mut Scratch<BE>,
113    ) where
114        A: AutomorphismKeyToMut,
115        B: GLWESecretToRef,
116    {
117        let res: &mut AutomorphismKey<&mut [u8]> = &mut res.to_mut();
118        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
119
120        #[cfg(debug_assertions)]
121        {
122            use crate::layouts::{GLWEInfos, LWEInfos};
123
124            assert_eq!(res.n(), sk.n());
125            assert_eq!(res.rank_out(), res.rank_in());
126            assert_eq!(sk.rank(), res.rank_out());
127            assert!(
128                scratch.available() >= AutomorphismKey::encrypt_sk_tmp_bytes(self, res),
129                "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {:?}",
130                scratch.available(),
131                AutomorphismKey::encrypt_sk_tmp_bytes(self, res)
132            )
133        }
134
135        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(self, sk.rank());
136
137        {
138            (0..res.rank_out().into()).for_each(|i| {
139                self.vec_znx_automorphism(
140                    self.galois_element_inv(p),
141                    &mut sk_out.data.as_vec_znx_mut(),
142                    i,
143                    &sk.data.as_vec_znx(),
144                    i,
145                );
146            });
147        }
148
149        res.key
150            .encrypt_sk(self, sk, &sk_out, source_xa, source_xe, scratch_1);
151
152        res.p = p;
153    }
154}