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