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