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