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