1use backend::hal::{
  2    api::{
  3        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes,
  4        ZnxZero,
  5    },
  6    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
  7};
  8use sampling::source::Source;
  9
 10use crate::{
 11    TakeGLWEPt,
 12    layouts::{GGLWECiphertext, GLWECiphertext, GLWEPlaintext, Infos, prepared::GLWESecretExec},
 13};
 14
 15use crate::trait_families::GGLWEEncryptSkFamily;
 16
 17impl GGLWECiphertext<Vec<u8>> {
 18    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
 19    where
 20        Module<B>: GGLWEEncryptSkFamily<B>,
 21    {
 22        GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
 23            + (GLWEPlaintext::byte_of(n, basek, k) | module.vec_znx_normalize_tmp_bytes(n))
 24    }
 25
 26    pub fn encrypt_pk_scratch_space<B: Backend>(_module: &Module<B>, _n: usize, _basek: usize, _k: usize, _rank: usize) -> usize {
 27        unimplemented!()
 28    }
 29}
 30
 31impl<DataSelf: DataMut> GGLWECiphertext<DataSelf> {
 32    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
 33        &mut self,
 34        module: &Module<B>,
 35        pt: &ScalarZnx<DataPt>,
 36        sk: &GLWESecretExec<DataSk, B>,
 37        source_xa: &mut Source,
 38        source_xe: &mut Source,
 39        sigma: f64,
 40        scratch: &mut Scratch<B>,
 41    ) where
 42        Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAddScalarInplace,
 43        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 44    {
 45        #[cfg(debug_assertions)]
 46        {
 47            use backend::hal::api::ZnxInfos;
 48
 49            assert_eq!(
 50                self.rank_in(),
 51                pt.cols(),
 52                "self.rank_in(): {} != pt.cols(): {}",
 53                self.rank_in(),
 54                pt.cols()
 55            );
 56            assert_eq!(
 57                self.rank_out(),
 58                sk.rank(),
 59                "self.rank_out(): {} != sk.rank(): {}",
 60                self.rank_out(),
 61                sk.rank()
 62            );
 63            assert_eq!(self.n(), sk.n());
 64            assert_eq!(pt.n(), sk.n());
 65            assert!(
 66                scratch.available() >= GGLWECiphertext::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k()),
 67                "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
 68                scratch.available(),
 69                self.rank(),
 70                self.size(),
 71                GGLWECiphertext::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k())
 72            );
 73            assert!(
 74                self.rows() * self.digits() * self.basek() <= self.k(),
 75                "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
 76                self.rows(),
 77                self.digits(),
 78                self.basek(),
 79                self.rows() * self.digits() * self.basek(),
 80                self.k()
 81            );
 82        }
 83
 84        let rows: usize = self.rows();
 85        let digits: usize = self.digits();
 86        let basek: usize = self.basek();
 87        let k: usize = self.k();
 88        let rank_in: usize = self.rank_in();
 89
 90        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(sk.n(), basek, k);
 91        // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
 92        //
 93        // Example for ksk rank 2 to rank 3:
 94        //
 95        // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
 96        // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
 97        //
 98        // Example ksk rank 2 to rank 1
 99        //
100        // (-(a*s) + s0, a)
101        // (-(b*s) + s1, b)
102        (0..rank_in).for_each(|col_i| {
103            (0..rows).for_each(|row_i| {
104                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
105                tmp_pt.data.zero(); // zeroes for next iteration
106                module.vec_znx_add_scalar_inplace(
107                    &mut tmp_pt.data,
108                    0,
109                    (digits - 1) + row_i * digits,
110                    pt,
111                    col_i,
112                );
113                module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scrach_1);
114
115                // rlwe encrypt of vec_znx_pt into vec_znx_ct
116                self.at_mut(row_i, col_i)
117                    .encrypt_sk(module, &tmp_pt, sk, source_xa, source_xe, sigma, scrach_1);
118            });
119        });
120    }
121}