From 187756a495db09b3fdc6724bcd8aae8e7db9be4c Mon Sep 17 00:00:00 2001 From: Jean-Philippe Bossuat Date: Wed, 28 May 2025 15:11:04 +0200 Subject: [PATCH] Added extension trait for struct allocation --- core/src/glwe_plaintext.rs | 7 +- core/src/keys.rs | 14 +-- core/src/lib.rs | 236 +++++++++++++++++++++++++++++++++++++ 3 files changed, 249 insertions(+), 8 deletions(-) diff --git a/core/src/glwe_plaintext.rs b/core/src/glwe_plaintext.rs index 82bcac7..3bf0060 100644 --- a/core/src/glwe_plaintext.rs +++ b/core/src/glwe_plaintext.rs @@ -1,6 +1,11 @@ use backend::{Backend, Module, VecZnx, VecZnxAlloc, VecZnxToMut, VecZnxToRef}; -use crate::{elem::{Infos, SetMetaData}, glwe_ciphertext::{GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef}, glwe_ops::GLWEOps, utils::derive_size}; +use crate::{ + elem::{Infos, SetMetaData}, + glwe_ciphertext::{GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef}, + glwe_ops::GLWEOps, + utils::derive_size, +}; pub struct GLWEPlaintext { pub data: VecZnx, diff --git a/core/src/keys.rs b/core/src/keys.rs index a6eb46c..cd69d28 100644 --- a/core/src/keys.rs +++ b/core/src/keys.rs @@ -7,7 +7,7 @@ use sampling::source::Source; use crate::{elem::Infos, glwe_ciphertext_fourier::GLWECiphertextFourier}; #[derive(Clone, Copy, Debug)] -pub enum SecretDistribution { +pub(crate) enum SecretDistribution { TernaryFixed(usize), // Ternary with fixed Hamming weight TernaryProb(f64), // Ternary with probabilistic Hamming weight ZERO, // Debug mod @@ -15,8 +15,8 @@ pub enum SecretDistribution { } pub struct SecretKey { - pub data: ScalarZnx, - pub dist: SecretDistribution, + pub(crate) data: ScalarZnx, + pub(crate) dist: SecretDistribution, } impl SecretKey> { @@ -64,8 +64,8 @@ impl + AsRef<[u8]>> SecretKey { } pub struct SecretKeyFourier { - pub data: ScalarZnxDft, - pub dist: SecretDistribution, + pub(crate) data: ScalarZnxDft, + pub(crate) dist: SecretDistribution, } impl SecretKeyFourier { @@ -113,8 +113,8 @@ impl + AsMut<[u8]>> SecretKeyFourier { } pub struct GLWEPublicKey { - pub data: GLWECiphertextFourier, - pub dist: SecretDistribution, + pub(crate) data: GLWECiphertextFourier, + pub(crate) dist: SecretDistribution, } impl GLWEPublicKey, B> { diff --git a/core/src/lib.rs b/core/src/lib.rs index 82b3c4b..b13ab8b 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -5,6 +5,7 @@ pub mod ggsw_ciphertext; pub mod glwe_ciphertext; pub mod glwe_ciphertext_fourier; pub mod glwe_ops; +pub mod glwe_packing; pub mod glwe_plaintext; pub mod keys; pub mod keyswitch_key; @@ -14,4 +15,239 @@ mod test_fft64; pub mod trace; mod utils; +pub use automorphism::*; +use backend::Backend; +use backend::FFT64; +use backend::Module; +pub use elem::*; +pub use gglwe_ciphertext::*; +pub use ggsw_ciphertext::*; +pub use glwe_ciphertext::*; +pub use glwe_ciphertext_fourier::*; +pub use glwe_ops::*; +pub use glwe_packing::*; +pub use glwe_plaintext::*; +pub use keys::*; +pub use keyswitch_key::*; +pub use tensor_key::*; + +pub use backend::Scratch; +pub use backend::ScratchOwned; +use utils::derive_size; + pub(crate) const SIX_SIGMA: f64 = 6.0; + +pub trait ScratchCore { + fn tmp_glwe(&mut self, module: &Module, basek: usize, k: usize, rank: usize) -> (GLWECiphertext<&mut [u8]>, &mut Self); + fn tmp_gglwe( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank_in: usize, + rank_out: usize, + ) -> (GGLWECiphertext<&mut [u8], B>, &mut Self); + fn tmp_ggsw( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank: usize, + ) -> (GGSWCiphertext<&mut [u8], B>, &mut Self); + fn tmp_glwe_fourier( + &mut self, + module: &Module, + basek: usize, + k: usize, + rank: usize, + ) -> (GLWECiphertextFourier<&mut [u8], B>, &mut Self); + fn tmp_sk(&mut self, module: &Module, rank: usize) -> (SecretKey<&mut [u8]>, &mut Self); + fn tmp_sk_fourier(&mut self, module: &Module, rank: usize) -> (SecretKeyFourier<&mut [u8], B>, &mut Self); + fn tmp_glwe_pk( + &mut self, + module: &Module, + basek: usize, + k: usize, + rank: usize, + ) -> (GLWEPublicKey<&mut [u8], B>, &mut Self); + fn tmp_glwe_ksk( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank_in: usize, + rank_out: usize, + ) -> (GLWESwitchingKey<&mut [u8], B>, &mut Self); + fn tmp_tsk( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank: usize, + ) -> (TensorKey<&mut [u8], B>, &mut Self); + fn tmp_autokey( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank: usize, + ) -> (AutomorphismKey<&mut [u8], B>, &mut Self); +} + +impl ScratchCore for Scratch { + fn tmp_glwe( + &mut self, + module: &Module, + basek: usize, + k: usize, + rank: usize, + ) -> (GLWECiphertext<&mut [u8]>, &mut Self) { + let (data, scratch) = self.tmp_vec_znx(module, rank + 1, derive_size(basek, k)); + (GLWECiphertext { data, basek, k }, scratch) + } + + fn tmp_gglwe( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank_in: usize, + rank_out: usize, + ) -> (GGLWECiphertext<&mut [u8], FFT64>, &mut Self) { + let (data, scratch) = self.tmp_mat_znx_dft(module, rows, rank_in, rank_out + 1, derive_size(basek, k)); + ( + GGLWECiphertext { + data: data, + basek: basek, + k, + }, + scratch, + ) + } + + fn tmp_ggsw( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank: usize, + ) -> (GGSWCiphertext<&mut [u8], FFT64>, &mut Self) { + let (data, scratch) = self.tmp_mat_znx_dft(module, rows, rank + 1, rank + 1, derive_size(basek, k)); + ( + GGSWCiphertext { + data: data, + basek: basek, + k, + }, + scratch, + ) + } + + fn tmp_glwe_fourier( + &mut self, + module: &Module, + basek: usize, + k: usize, + rank: usize, + ) -> (GLWECiphertextFourier<&mut [u8], FFT64>, &mut Self) { + let (data, scratch) = self.tmp_vec_znx_dft(module, rank + 1, derive_size(basek, k)); + (GLWECiphertextFourier { data, basek, k }, scratch) + } + + fn tmp_glwe_pk( + &mut self, + module: &Module, + basek: usize, + k: usize, + rank: usize, + ) -> (GLWEPublicKey<&mut [u8], FFT64>, &mut Self) { + let (data, scratch) = self.tmp_glwe_fourier(module, basek, k, rank); + ( + GLWEPublicKey { + data, + dist: SecretDistribution::NONE, + }, + scratch, + ) + } + + fn tmp_sk(&mut self, module: &Module, rank: usize) -> (SecretKey<&mut [u8]>, &mut Self) { + let (data, scratch) = self.tmp_scalar_znx(module, rank + 1); + ( + SecretKey { + data, + dist: SecretDistribution::NONE, + }, + scratch, + ) + } + + fn tmp_sk_fourier(&mut self, module: &Module, rank: usize) -> (SecretKeyFourier<&mut [u8], FFT64>, &mut Self) { + let (data, scratch) = self.tmp_scalar_znx_dft(module, rank + 1); + ( + SecretKeyFourier { + data, + dist: SecretDistribution::NONE, + }, + scratch, + ) + } + + fn tmp_glwe_ksk( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank_in: usize, + rank_out: usize, + ) -> (GLWESwitchingKey<&mut [u8], FFT64>, &mut Self) { + let (data, scratch) = self.tmp_gglwe(module, basek, k, rows, rank_in, rank_out); + (GLWESwitchingKey(data), scratch) + } + + fn tmp_autokey( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank: usize, + ) -> (AutomorphismKey<&mut [u8], FFT64>, &mut Self) { + let (data, scratch) = self.tmp_glwe_ksk(module, basek, k, rows, rank, rank); + (AutomorphismKey { key: data, p: 0 }, scratch) + } + + fn tmp_tsk( + &mut self, + module: &Module, + basek: usize, + k: usize, + rows: usize, + rank: usize, + ) -> (TensorKey<&mut [u8], FFT64>, &mut Self) { + let mut keys: Vec> = Vec::new(); + let pairs: usize = (((rank + 1) * rank) >> 1).max(1); + + let mut scratch: &mut Scratch = self; + + if pairs != 0 { + let (gglwe, s) = scratch.tmp_glwe_ksk(module, basek, k, rows, 1, rank); + scratch = s; + keys.push(gglwe); + } + for _ in 1..pairs { + let (gglwe, s) = scratch.tmp_glwe_ksk(module, basek, k, rows, 1, rank); + scratch = s; + keys.push(gglwe); + } + (TensorKey { keys }, scratch) + } +}