1use poulpy_backend::hal::{
  2    api::{
  3        ScratchAvailable, SvpApplyInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftAllocBytes,
  5        VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace,
  6        VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  9    source::Source,
 10};
 11
 12use crate::{
 13    TakeGLWESecret, TakeGLWESecretPrepared,
 14    layouts::{
 15        GLWESecret,
 16        compressed::{GGLWEAutomorphismKeyCompressed, GGLWESwitchingKeyCompressed},
 17    },
 18};
 19
 20impl GGLWEAutomorphismKeyCompressed<Vec<u8>> {
 21    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
 22    where
 23        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes + SvpPPolAllocBytes,
 24    {
 25        GGLWESwitchingKeyCompressed::encrypt_sk_scratch_space(module, n, basek, k, rank, rank) + GLWESecret::bytes_of(n, rank)
 26    }
 27}
 28
 29impl<DataSelf: DataMut> GGLWEAutomorphismKeyCompressed<DataSelf> {
 30    #[allow(clippy::too_many_arguments)]
 31    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
 32        &mut self,
 33        module: &Module<B>,
 34        p: i64,
 35        sk: &GLWESecret<DataSk>,
 36        seed_xa: [u8; 32],
 37        source_xe: &mut Source,
 38        sigma: f64,
 39        scratch: &mut Scratch<B>,
 40    ) where
 41        Module<B>: VecZnxAutomorphism
 42            + SvpPrepare<B>
 43            + SvpPPolAllocBytes
 44            + VecZnxSwithcDegree
 45            + VecZnxDftAllocBytes
 46            + VecZnxBigNormalize<B>
 47            + VecZnxDftFromVecZnx<B>
 48            + SvpApplyInplace<B>
 49            + VecZnxDftToVecZnxBigConsume<B>
 50            + VecZnxNormalizeTmpBytes
 51            + VecZnxFillUniform
 52            + VecZnxSubABInplace
 53            + VecZnxAddInplace
 54            + VecZnxNormalizeInplace<B>
 55            + VecZnxAddNormal
 56            + VecZnxNormalize<B>
 57            + VecZnxSub
 58            + VecZnxAddScalarInplace,
 59        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
 60    {
 61        #[cfg(debug_assertions)]
 62        {
 63            use crate::layouts::Infos;
 64
 65            assert_eq!(self.n(), sk.n());
 66            assert_eq!(self.rank_out(), self.rank_in());
 67            assert_eq!(sk.rank(), self.rank());
 68            assert!(
 69                scratch.available()
 70                    >= GGLWEAutomorphismKeyCompressed::encrypt_sk_scratch_space(
 71                        module,
 72                        sk.n(),
 73                        self.basek(),
 74                        self.k(),
 75                        self.rank()
 76                    ),
 77                "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
 78                scratch.available(),
 79                self.rank(),
 80                self.size(),
 81                GGLWEAutomorphismKeyCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank())
 82            )
 83        }
 84
 85        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
 86
 87        {
 88            (0..self.rank()).for_each(|i| {
 89                module.vec_znx_automorphism(
 90                    module.galois_element_inv(p),
 91                    &mut sk_out.data.as_vec_znx_mut(),
 92                    i,
 93                    &sk.data.as_vec_znx(),
 94                    i,
 95                );
 96            });
 97        }
 98
 99        self.key
100            .encrypt_sk(module, sk, &sk_out, seed_xa, source_xe, sigma, scratch_1);
101
102        self.p = p;
103    }
104}