1use backend::hal::{
  2    api::{
  3        ScratchAvailable, SvpApplyInplace, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
  4        VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform,
  5        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, ZnxZero,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
  8};
  9use sampling::source::Source;
 10
 11use crate::{
 12    TakeGLWEPt,
 13    encryption::glwe_encrypt_sk_internal,
 14    layouts::{GGLWECiphertext, Infos, compressed::GGLWECiphertextCompressed, prepared::GLWESecretPrepared},
 15};
 16
 17impl GGLWECiphertextCompressed<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>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
 21    {
 22        GGLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
 23    }
 24}
 25
 26impl<D: DataMut> GGLWECiphertextCompressed<D> {
 27    #[allow(clippy::too_many_arguments)]
 28    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
 29        &mut self,
 30        module: &Module<B>,
 31        pt: &ScalarZnx<DataPt>,
 32        sk: &GLWESecretPrepared<DataSk, B>,
 33        seed: [u8; 32],
 34        source_xe: &mut Source,
 35        sigma: f64,
 36        scratch: &mut Scratch<B>,
 37    ) where
 38        Module<B>: VecZnxAddScalarInplace
 39            + VecZnxDftAllocBytes
 40            + VecZnxBigNormalize<B>
 41            + VecZnxDftFromVecZnx<B>
 42            + SvpApplyInplace<B>
 43            + VecZnxDftToVecZnxBigConsume<B>
 44            + VecZnxNormalizeTmpBytes
 45            + VecZnxFillUniform
 46            + VecZnxSubABInplace
 47            + VecZnxAddInplace
 48            + VecZnxNormalizeInplace<B>
 49            + VecZnxAddNormal
 50            + VecZnxNormalize<B>
 51            + VecZnxSub,
 52        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 53    {
 54        #[cfg(debug_assertions)]
 55        {
 56            use backend::hal::api::ZnxInfos;
 57
 58            assert_eq!(
 59                self.rank_in(),
 60                pt.cols(),
 61                "self.rank_in(): {} != pt.cols(): {}",
 62                self.rank_in(),
 63                pt.cols()
 64            );
 65            assert_eq!(
 66                self.rank_out(),
 67                sk.rank(),
 68                "self.rank_out(): {} != sk.rank(): {}",
 69                self.rank_out(),
 70                sk.rank()
 71            );
 72            assert_eq!(self.n(), sk.n());
 73            assert_eq!(pt.n(), sk.n());
 74            assert!(
 75                scratch.available()
 76                    >= GGLWECiphertextCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k()),
 77                "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
 78                scratch.available(),
 79                self.rank(),
 80                self.size(),
 81                GGLWECiphertextCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k())
 82            );
 83            assert!(
 84                self.rows() * self.digits() * self.basek() <= self.k(),
 85                "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
 86                self.rows(),
 87                self.digits(),
 88                self.basek(),
 89                self.rows() * self.digits() * self.basek(),
 90                self.k()
 91            );
 92        }
 93
 94        let rows: usize = self.rows();
 95        let digits: usize = self.digits();
 96        let basek: usize = self.basek();
 97        let k: usize = self.k();
 98        let rank_in: usize = self.rank_in();
 99        let cols: usize = self.rank_out() + 1;
100
101        let mut source_xa = Source::new(seed);
102
103        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(sk.n(), basek, k);
104        (0..rank_in).for_each(|col_i| {
105            (0..rows).for_each(|row_i| {
106                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
107                tmp_pt.data.zero(); // zeroes for next iteration
108                module.vec_znx_add_scalar_inplace(
109                    &mut tmp_pt.data,
110                    0,
111                    (digits - 1) + row_i * digits,
112                    pt,
113                    col_i,
114                );
115                module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scrach_1);
116
117                let (seed, mut source_xa_tmp) = source_xa.branch();
118                self.seed[col_i * rows + row_i] = seed;
119
120                glwe_encrypt_sk_internal(
121                    module,
122                    self.basek(),
123                    self.k(),
124                    &mut self.at_mut(row_i, col_i).data,
125                    cols,
126                    true,
127                    Some((&tmp_pt, 0)),
128                    sk,
129                    &mut source_xa_tmp,
130                    source_xe,
131                    sigma,
132                    scrach_1,
133                );
134            });
135        });
136    }
137}