1use poulpy_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        scratch: &mut Scratch<B>,
 39    ) where
 40        Module<B>: VecZnxAutomorphism
 41            + SvpPrepare<B>
 42            + SvpPPolAllocBytes
 43            + VecZnxSwithcDegree
 44            + VecZnxDftAllocBytes
 45            + VecZnxBigNormalize<B>
 46            + VecZnxDftFromVecZnx<B>
 47            + SvpApplyInplace<B>
 48            + VecZnxDftToVecZnxBigConsume<B>
 49            + VecZnxNormalizeTmpBytes
 50            + VecZnxFillUniform
 51            + VecZnxSubABInplace
 52            + VecZnxAddInplace
 53            + VecZnxNormalizeInplace<B>
 54            + VecZnxAddNormal
 55            + VecZnxNormalize<B>
 56            + VecZnxSub
 57            + VecZnxAddScalarInplace,
 58        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
 59    {
 60        #[cfg(debug_assertions)]
 61        {
 62            use crate::layouts::Infos;
 63
 64            assert_eq!(self.n(), sk.n());
 65            assert_eq!(self.rank_out(), self.rank_in());
 66            assert_eq!(sk.rank(), self.rank());
 67            assert!(
 68                scratch.available()
 69                    >= GGLWEAutomorphismKeyCompressed::encrypt_sk_scratch_space(
 70                        module,
 71                        sk.n(),
 72                        self.basek(),
 73                        self.k(),
 74                        self.rank()
 75                    ),
 76                "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
 77                scratch.available(),
 78                self.rank(),
 79                self.size(),
 80                GGLWEAutomorphismKeyCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank())
 81            )
 82        }
 83
 84        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
 85
 86        {
 87            (0..self.rank()).for_each(|i| {
 88                module.vec_znx_automorphism(
 89                    module.galois_element_inv(p),
 90                    &mut sk_out.data.as_vec_znx_mut(),
 91                    i,
 92                    &sk.data.as_vec_znx(),
 93                    i,
 94                );
 95            });
 96        }
 97
 98        self.key
 99            .encrypt_sk(module, sk, &sk_out, seed_xa, source_xe, scratch_1);
100
101        self.p = p;
102    }
103}