Start of full rewrite of rlwe crate

This commit is contained in:
Jean-Philippe Bossuat
2025-05-06 14:38:22 +02:00
parent 08e81f50c9
commit 4efe22e723
6 changed files with 223 additions and 9 deletions

View File

@@ -10,7 +10,3 @@ base2k = {path="../base2k"}
sampling = {path="../sampling"}
rand_distr = {workspace = true}
itertools = {workspace = true}
[[bench]]
name = "gadget_product"
harness = false

View File

@@ -1,4 +0,0 @@
pub struct Ciphertext{
x
}

85
rlwe/src/elem.rs Normal file
View File

@@ -0,0 +1,85 @@
use base2k::{Backend, FFT64, MatZnxDft, MatZnxDftAlloc, Module, VecZnx, VecZnxAlloc};
pub struct Ciphertext<T> {
data: T,
log_base2k: usize,
log_q: usize,
}
impl<T> Ciphertext<T> {
pub fn log_base2k(&self) -> usize {
self.log_base2k
}
pub fn log_q(&self) -> usize {
self.log_q
}
pub fn data(&self) -> &T {
&self.data
}
pub fn data_mut(&mut self) -> &mut T {
&mut self.data
}
}
pub struct Plaintext<T> {
data: T,
log_base2k: usize,
log_q: usize,
}
impl<T> Plaintext<T> {
pub fn log_base2k(&self) -> usize {
self.log_base2k
}
pub fn log_q(&self) -> usize {
self.log_q
}
pub fn data(&self) -> &T {
&self.data
}
pub fn data_mut(&mut self) -> &mut T {
&mut self.data
}
}
pub(crate) type CipherVecZnx<C> = Ciphertext<VecZnx<C>>;
impl Ciphertext<VecZnx<Vec<u8>>> {
pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_q: usize, cols: usize) -> Self {
Self {
data: module.new_vec_znx(cols, derive_size(log_base2k, log_q)),
log_base2k: log_base2k,
log_q: log_q,
}
}
}
impl Plaintext<VecZnx<Vec<u8>>> {
pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_q: usize) -> Self {
Self {
data: module.new_vec_znx(1, derive_size(log_base2k, log_q)),
log_base2k: log_base2k,
log_q: log_q,
}
}
}
impl<B: Backend> Ciphertext<MatZnxDft<Vec<u8>, B>> {
pub fn new(module: &Module<B>, log_base2k: usize, rows: usize, cols_in: usize, cols_out: usize, log_q: usize) -> Self {
Self {
data: module.new_mat_znx_dft(rows, cols_in, cols_out, derive_size(log_base2k, log_q)),
log_base2k: log_base2k,
log_q: log_q,
}
}
}
pub(crate) fn derive_size(log_base2k: usize, log_q: usize) -> usize {
(log_q + log_base2k - 1) / log_base2k
}

71
rlwe/src/encryption.rs Normal file
View File

@@ -0,0 +1,71 @@
use base2k::{
AddNormal, Backend, FillUniform, Module, VecZnxDftOps, ScalarZnxDftOps, ScalarZnxDftToRef, VecZnxBigOps, Scratch, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, FFT64
};
use sampling::source::Source;
use crate::{
elem::{CipherVecZnx, Plaintext},
keys::SecretKey,
};
pub trait EncryptSk<B: Backend, D> {
fn encrypt<P, S>(
module: &Module<B>,
res: &mut D,
pt: Option<&Plaintext<P>>,
sk: &SecretKey<S>,
source_xa: &mut Source,
source_xe: &mut Source,
scratch: &mut Scratch,
sigma: f64,
bound: f64,
) where
P: VecZnxToRef,
S: ScalarZnxDftToRef<B>;
}
impl<C> EncryptSk<FFT64, CipherVecZnx<C>> for CipherVecZnx<C>
where
VecZnx<C>: VecZnxToMut + VecZnxToRef,
{
fn encrypt<P, S>(
module: &Module<FFT64>,
ct: &mut CipherVecZnx<C>,
pt: Option<&Plaintext<P>>,
sk: &SecretKey<S>,
source_xa: &mut Source,
source_xe: &mut Source,
scratch: &mut Scratch,
sigma: f64,
bound: f64,
) where
P: VecZnxToRef,
S: ScalarZnxDftToRef<FFT64>,
{
let log_base2k: usize = ct.log_base2k();
let log_q: usize = ct.log_q();
let mut ct_mut: VecZnx<&mut [u8]> = ct.data_mut().to_mut();
let size: usize = ct_mut.size();
ct_mut.fill_uniform(log_base2k, 1, size, source_xa);
// c1_dft = DFT(a) * DFT(s)
let (mut c1_dft, scratch_1) = scratch.tmp_vec_znx_dft(module, 1, size);
module.svp_apply(&mut c1_dft, 0, &sk.data().to_ref(), 0, &ct_mut, 1);
// c1_big = IDFT(c1_dft)
let (mut c1_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, size);
module.vec_znx_idft_tmp_a(&mut c1_big, 0, &mut c1_dft, 0);
// c1_big = m - c1_big
if let Some(pt) = pt {
module.vec_znx_big_sub_small_b_inplace(&mut c1_big, 0, &pt.data().to_ref(), 0);
}
// c1_big += e
c1_big.add_normal(log_base2k, 0, log_q, source_xe, sigma, bound);
// c0 = norm(c1_big)
module.vec_znx_big_normalize(log_base2k, &mut ct_mut, 0, &c1_big, 0, scratch_2);
}
}

64
rlwe/src/keys.rs Normal file
View File

@@ -0,0 +1,64 @@
use base2k::{
Backend, FFT64, Module, Scalar, ScalarAlloc, ScalarZnxDft, ScalarZnxDftOps, ScalarZnxDftToMut, Scratch, VecZnx, VecZnxDft,
VecZnxDftAlloc, VecZnxDftToMut,
};
use sampling::source::Source;
use crate::elem::derive_size;
pub struct SecretKey<T> {
data: T,
}
impl<T> SecretKey<T> {
pub fn data(&self) -> &T {
&self.data
}
pub fn data_mut(&self) -> &mut T {
&mut self.data
}
}
impl SecretKey<Scalar<Vec<u8>>> {
pub fn new<B: Backend>(module: &Module<B>) -> Self {
Self {
data: module.new_scalar(1),
}
}
pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
self.data.fill_ternary_prob(0, prob, source);
}
pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
self.data.fill_ternary_hw(0, hw, source);
}
pub fn svp_prepare<D>(&self, module: &Module<FFT64>, sk_prep: &mut SecretKey<ScalarZnxDft<D, FFT64>>)
where
ScalarZnxDft<D, base2k::FFT64>: ScalarZnxDftToMut<base2k::FFT64>,
{
module.svp_prepare(&mut sk_prep.data, 0, &self.data, 0)
}
}
pub struct PublicKey<D, B: Backend> {
data: VecZnxDft<D, B>,
}
impl<B: Backend> PublicKey<Vec<u8>, B> {
pub fn new(module: &Module<B>, log_base2k: usize, log_q: usize) -> Self {
Self {
data: module.new_vec_znx_dft(2, derive_size(log_base2k, log_q)),
}
}
}
impl<B: Backend, D: VecZnxDftToMut<B>> PublicKey<D, B> {
pub fn generate<S>(&mut self, module: &Module<B>, sk: &SecretKey<ScalarZnxDft<S, B>>)
where
ScalarZnxDft<S, B>: ScalarZnxDftToMut<B>,
{
}
}

View File

@@ -1 +1,3 @@
pub mod ciphertext;
pub mod elem;
pub mod encryption;
pub mod keys;