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