1use poulpy_hal::{
  2    api::{ModuleN, ScratchAvailable, VecZnxAddScalarInplace, VecZnxDftBytesOf, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, DataMut, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxInfos, ZnxZero},
  4    source::Source,
  5};
  6
  7use crate::{
  8    ScratchTakeCore,
  9    encryption::{
 10        SIGMA,
 11        glwe_ct::{GLWEEncryptSk, GLWEEncryptSkInternal},
 12    },
 13    layouts::{
 14        GGLWECompressedSeedMut, GGLWEInfos, GLWEPlaintextAlloc, GLWESecretPrepared, LWEInfos,
 15        compressed::{GGLWECompressed, GGLWECompressedToMut},
 16        prepared::GLWESecretPreparedToRef,
 17    },
 18};
 19
 20impl<D: DataMut> GGLWECompressed<D> {
 21    #[allow(clippy::too_many_arguments)]
 22    pub fn encrypt_sk<M, P, S, BE: Backend>(
 23        &mut self,
 24        module: &M,
 25        pt: &P,
 26        sk: &S,
 27        seed: [u8; 32],
 28        source_xe: &mut Source,
 29        scratch: &mut Scratch<BE>,
 30    ) where
 31        P: ScalarZnxToRef,
 32        S: GLWESecretPreparedToRef<BE>,
 33        M: GGLWECompressedEncryptSk<BE>,
 34    {
 35        module.gglwe_compressed_encrypt_sk(self, pt, sk, seed, source_xe, scratch);
 36    }
 37}
 38
 39impl GGLWECompressed<Vec<u8>> {
 40    pub fn encrypt_sk_tmp_bytes<M, BE: Backend, A>(module: &M, infos: &A) -> usize
 41    where
 42        A: GGLWEInfos,
 43        M: GGLWECompressedEncryptSk<BE>,
 44    {
 45        module.gglwe_compressed_encrypt_sk_tmp_bytes(infos)
 46    }
 47}
 48
 49pub trait GGLWECompressedEncryptSk<BE: Backend> {
 50    fn gglwe_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 51    where
 52        A: GGLWEInfos;
 53
 54    fn gglwe_compressed_encrypt_sk<R, P, S>(
 55        &self,
 56        res: &mut R,
 57        pt: &P,
 58        sk: &S,
 59        seed: [u8; 32],
 60        source_xe: &mut Source,
 61        scratch: &mut Scratch<BE>,
 62    ) where
 63        R: GGLWECompressedToMut + GGLWECompressedSeedMut,
 64        P: ScalarZnxToRef,
 65        S: GLWESecretPreparedToRef<BE>;
 66}
 67
 68impl<BE: Backend> GGLWECompressedEncryptSk<BE> for Module<BE>
 69where
 70    Self: ModuleN
 71        + GLWEEncryptSkInternal<BE>
 72        + GLWEEncryptSk<BE>
 73        + VecZnxDftBytesOf
 74        + VecZnxNormalizeInplace<BE>
 75        + VecZnxAddScalarInplace
 76        + VecZnxNormalizeTmpBytes,
 77    Scratch<BE>: ScratchTakeCore<BE>,
 78{
 79    fn gglwe_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 80    where
 81        A: GGLWEInfos,
 82    {
 83        self.glwe_encrypt_sk_tmp_bytes(infos)
 84            .max(self.vec_znx_normalize_tmp_bytes())
 85            + self.bytes_of_glwe_plaintext_from_infos(infos)
 86    }
 87
 88    fn gglwe_compressed_encrypt_sk<R, P, S>(
 89        &self,
 90        res: &mut R,
 91        pt: &P,
 92        sk: &S,
 93        seed: [u8; 32],
 94        source_xe: &mut Source,
 95        scratch: &mut Scratch<BE>,
 96    ) where
 97        R: GGLWECompressedToMut + GGLWECompressedSeedMut,
 98        P: ScalarZnxToRef,
 99        S: GLWESecretPreparedToRef<BE>,
100    {
101        let mut seeds: Vec<[u8; 32]> = vec![[0u8; 32]; res.seed_mut().len()];
102
103        {
104            let res: &mut GGLWECompressed<&mut [u8]> = &mut res.to_mut();
105            let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
106            let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();
107
108            assert_eq!(
109                res.rank_in(),
110                pt.cols() as u32,
111                "res.rank_in(): {} != pt.cols(): {}",
112                res.rank_in(),
113                pt.cols()
114            );
115            assert_eq!(
116                res.rank_out(),
117                sk.rank(),
118                "res.rank_out(): {} != sk.rank(): {}",
119                res.rank_out(),
120                sk.rank()
121            );
122            assert_eq!(res.n(), sk.n());
123            assert_eq!(pt.n() as u32, sk.n());
124            assert!(
125                scratch.available() >= GGLWECompressed::encrypt_sk_tmp_bytes(self, res),
126                "scratch.available: {} < GGLWECiphertext::encrypt_sk_tmp_bytes: {}",
127                scratch.available(),
128                GGLWECompressed::encrypt_sk_tmp_bytes(self, res)
129            );
130            assert!(
131                res.dnum().0 * res.dsize().0 * res.base2k().0 <= res.k().0,
132                "res.dnum() : {} * res.dsize() : {} * res.base2k() : {} = {} >= res.k() = {}",
133                res.dnum(),
134                res.dsize(),
135                res.base2k(),
136                res.dnum().0 * res.dsize().0 * res.base2k().0,
137                res.k()
138            );
139
140            let dnum: usize = res.dnum().into();
141            let dsize: usize = res.dsize().into();
142            let base2k: usize = res.base2k().into();
143            let rank_in: usize = res.rank_in().into();
144            let cols: usize = (res.rank_out() + 1).into();
145
146            let mut source_xa = Source::new(seed);
147
148            let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, res);
149            for col_i in 0..rank_in {
150                for d_i in 0..dnum {
151                    // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
152                    tmp_pt.data.zero(); // zeroes for next iteration
153                    self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + d_i * dsize, pt, col_i);
154                    self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
155
156                    let (seed, mut source_xa_tmp) = source_xa.branch();
157                    seeds[col_i * dnum + d_i] = seed;
158
159                    self.glwe_encrypt_sk_internal(
160                        res.base2k().into(),
161                        res.k().into(),
162                        &mut res.at_mut(d_i, col_i).data,
163                        cols,
164                        true,
165                        Some((&tmp_pt, 0)),
166                        sk,
167                        &mut source_xa_tmp,
168                        source_xe,
169                        SIGMA,
170                        scrach_1,
171                    );
172                }
173            }
174        }
175
176        res.seed_mut().copy_from_slice(&seeds);
177    }
178}