1use base2k::{
  2    Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft,
  3    ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnxAlloc, VecZnxDftAlloc, VecZnxDftOps, ZnxView, ZnxViewMut,
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    elem::Infos,
  9    elem_grlwe::GRLWECt,
 10    elem_rlwe::{RLWECt, RLWECtDft, RLWEPt},
 11    keys::SecretKeyDft,
 12    utils::derive_size,
 13};
 14
 15pub struct RGSWCt<C, B: Backend> {
 16    pub data: MatZnxDft<C, B>,
 17    pub log_base2k: usize,
 18    pub log_k: usize,
 19}
 20
 21impl<B: Backend> RGSWCt<Vec<u8>, B> {
 22    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
 23        Self {
 24            data: module.new_mat_znx_dft(rows, 2, 2, derive_size(log_base2k, log_k)),
 25            log_base2k: log_base2k,
 26            log_k: log_k,
 27        }
 28    }
 29}
 30
 31impl<T, B: Backend> Infos for RGSWCt<T, B> {
 32    type Inner = MatZnxDft<T, B>;
 33
 34    fn inner(&self) -> &Self::Inner {
 35        &self.data
 36    }
 37
 38    fn log_base2k(&self) -> usize {
 39        self.log_base2k
 40    }
 41
 42    fn log_k(&self) -> usize {
 43        self.log_k
 44    }
 45}
 46
 47impl<C, B: Backend> MatZnxDftToMut<B> for RGSWCt<C, B>
 48where
 49    MatZnxDft<C, B>: MatZnxDftToMut<B>,
 50{
 51    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
 52        self.data.to_mut()
 53    }
 54}
 55
 56impl<C, B: Backend> MatZnxDftToRef<B> for RGSWCt<C, B>
 57where
 58    MatZnxDft<C, B>: MatZnxDftToRef<B>,
 59{
 60    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
 61        self.data.to_ref()
 62    }
 63}
 64
 65impl GRLWECt<Vec<u8>, FFT64> {
 66    pub fn encrypt_sk_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
 67        RLWECt::encrypt_sk_scratch_bytes(module, size)
 68            + module.bytes_of_vec_znx(2, size)
 69            + module.bytes_of_vec_znx(1, size)
 70            + module.bytes_of_vec_znx_dft(2, size)
 71    }
 72
 73    pub fn encrypt_pk_scratch_bytes(module: &Module<FFT64>, pk_size: usize) -> usize {
 74        RLWECt::encrypt_pk_scratch_bytes(module, pk_size)
 75    }
 76
 77    pub fn decrypt_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
 78        RLWECtDft::decrypt_scratch_bytes(module, size)
 79    }
 80}
 81
 82pub fn encrypt_grlwe_sk<C, P, S>(
 83    module: &Module<FFT64>,
 84    ct: &mut GRLWECt<C, FFT64>,
 85    pt: &ScalarZnx<P>,
 86    sk: &SecretKeyDft<S, FFT64>,
 87    source_xa: &mut Source,
 88    source_xe: &mut Source,
 89    sigma: f64,
 90    bound: f64,
 91    scratch: &mut Scratch,
 92) where
 93    MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
 94    ScalarZnx<P>: ScalarZnxToRef,
 95    ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
 96{
 97    let rows: usize = ct.rows();
 98    let size: usize = ct.size();
 99
100    let (tmp_znx_pt, scrach_1) = scratch.tmp_vec_znx(module, 1, size);
101    let (tmp_znx_ct, scrach_2) = scrach_1.tmp_vec_znx(module, 2, size);
102    let (mut tmp_dft, scratch_3) = scrach_2.tmp_vec_znx_dft(module, 2, size);
103
104    let mut tmp_pt: RLWEPt<&mut [u8]> = RLWEPt {
105        data: tmp_znx_pt,
106        log_base2k: ct.log_base2k(),
107        log_k: ct.log_k(),
108    };
109
110    let mut tmp_ct: RLWECt<&mut [u8]> = RLWECt {
111        data: tmp_znx_ct,
112        log_base2k: ct.log_base2k(),
113        log_k: ct.log_k(),
114    };
115
116    (0..rows).for_each(|row_i| {
117        tmp_pt
118            .data
119            .at_mut(0, row_i)
120            .copy_from_slice(&pt.to_ref().raw());
121
122        tmp_ct.encrypt_sk(
123            module,
124            Some(&tmp_pt),
125            sk,
126            source_xa,
127            source_xe,
128            sigma,
129            bound,
130            scratch_3,
131        );
132
133        tmp_pt.data.at_mut(0, row_i).fill(0);
134
135        module.vec_znx_dft(&mut tmp_dft, 0, &tmp_ct, 0);
136        module.vec_znx_dft(&mut tmp_dft, 1, &tmp_ct, 1);
137
138        module.vmp_prepare_row(ct, row_i, 0, &tmp_dft);
139    });
140}