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        GGLWEInfos, GLWEPlaintextAlloc, 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,
 64        P: ScalarZnxToRef,
 65        S: GLWESecretPreparedToRef<BE>;
 66}
 67
 68impl<BE: Backend> GGLWECompressedEncryptSk<BE> for Module<BE>
 69where
 70    Module<BE>: 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,
 98        P: ScalarZnxToRef,
 99        S: GLWESecretPreparedToRef<BE>,
100    {
101        let res: &mut GGLWECompressed<&mut [u8]> = &mut res.to_mut();
102        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
103
104        let sk = &sk.to_ref();
105
106        assert_eq!(
107            res.rank_in(),
108            pt.cols() as u32,
109            "res.rank_in(): {} != pt.cols(): {}",
110            res.rank_in(),
111            pt.cols()
112        );
113        assert_eq!(
114            res.rank_out(),
115            sk.rank(),
116            "res.rank_out(): {} != sk.rank(): {}",
117            res.rank_out(),
118            sk.rank()
119        );
120        assert_eq!(res.n(), sk.n());
121        assert_eq!(pt.n() as u32, sk.n());
122        assert!(
123            scratch.available() >= GGLWECompressed::encrypt_sk_tmp_bytes(self, res),
124            "scratch.available: {} < GGLWECiphertext::encrypt_sk_tmp_bytes: {}",
125            scratch.available(),
126            GGLWECompressed::encrypt_sk_tmp_bytes(self, res)
127        );
128        assert!(
129            res.dnum().0 * res.dsize().0 * res.base2k().0 <= res.k().0,
130            "res.dnum() : {} * res.dsize() : {} * res.base2k() : {} = {} >= res.k() = {}",
131            res.dnum(),
132            res.dsize(),
133            res.base2k(),
134            res.dnum().0 * res.dsize().0 * res.base2k().0,
135            res.k()
136        );
137
138        let dnum: usize = res.dnum().into();
139        let dsize: usize = res.dsize().into();
140        let base2k: usize = res.base2k().into();
141        let rank_in: usize = res.rank_in().into();
142        let cols: usize = (res.rank_out() + 1).into();
143
144        let mut source_xa = Source::new(seed);
145
146        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, res);
147        (0..rank_in).for_each(|col_i| {
148            (0..dnum).for_each(|d_i| {
149                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
150                tmp_pt.data.zero(); // zeroes for next iteration
151                self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + d_i * dsize, pt, col_i);
152                self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
153
154                let (seed, mut source_xa_tmp) = source_xa.branch();
155                res.seed[col_i * dnum + d_i] = seed;
156
157                self.glwe_encrypt_sk_internal(
158                    res.base2k().into(),
159                    res.k().into(),
160                    &mut res.at_mut(d_i, col_i).data,
161                    cols,
162                    true,
163                    Some((&tmp_pt, 0)),
164                    sk,
165                    &mut source_xa_tmp,
166                    source_xe,
167                    SIGMA,
168                    scrach_1,
169                );
170            });
171        });
172    }
173}