mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 21:26:41 +01:00
refactor
This commit is contained in:
@@ -1,7 +1,6 @@
|
||||
use backend::{
|
||||
Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDftAlloc,
|
||||
ScalarZnxDftOps, ScalarZnxOps, ScalarZnxToRef, Scratch, VecZnx, VecZnxBigAlloc, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut,
|
||||
VecZnxDftToRef, VecZnxOps, ZnxZero,
|
||||
Backend, FFT64, MatZnxDft, MatZnxDftOps, Module, ScalarZnxDftAlloc, ScalarZnxDftOps, ScalarZnxOps, Scratch, VecZnx,
|
||||
VecZnxBigAlloc, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, ZnxZero,
|
||||
};
|
||||
use sampling::source::Source;
|
||||
|
||||
@@ -63,45 +62,27 @@ impl<T, B: Backend> AutomorphismKey<T, B> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf, B: Backend> MatZnxDftToMut<B> for AutomorphismKey<DataSelf, B>
|
||||
where
|
||||
MatZnxDft<DataSelf, B>: MatZnxDftToMut<B>,
|
||||
{
|
||||
fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
|
||||
self.key.to_mut()
|
||||
impl<C: AsRef<[u8]>> GetRow<FFT64> for AutomorphismKey<C, FFT64> {
|
||||
fn get_row<R: AsMut<[u8]> + AsRef<[u8]>>(
|
||||
&self,
|
||||
module: &Module<FFT64>,
|
||||
row_i: usize,
|
||||
col_j: usize,
|
||||
res: &mut GLWECiphertextFourier<R, FFT64>,
|
||||
) {
|
||||
module.vmp_extract_row(&mut res.data, &self.key.0.data, row_i, col_j);
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf, B: Backend> MatZnxDftToRef<B> for AutomorphismKey<DataSelf, B>
|
||||
where
|
||||
MatZnxDft<DataSelf, B>: MatZnxDftToRef<B>,
|
||||
{
|
||||
fn to_ref(&self) -> MatZnxDft<&[u8], B> {
|
||||
self.key.to_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> GetRow<FFT64> for AutomorphismKey<C, FFT64>
|
||||
where
|
||||
MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
{
|
||||
fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut R)
|
||||
where
|
||||
R: VecZnxDftToMut<FFT64>,
|
||||
{
|
||||
module.vmp_extract_row(res, self, row_i, col_j);
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> SetRow<FFT64> for AutomorphismKey<C, FFT64>
|
||||
where
|
||||
MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
|
||||
{
|
||||
fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &R)
|
||||
where
|
||||
R: VecZnxDftToRef<FFT64>,
|
||||
{
|
||||
module.vmp_prepare_row(self, row_i, col_j, a);
|
||||
impl<C: AsMut<[u8]> + AsRef<[u8]>> SetRow<FFT64> for AutomorphismKey<C, FFT64> {
|
||||
fn set_row<R: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
row_i: usize,
|
||||
col_j: usize,
|
||||
a: &GLWECiphertextFourier<R, FFT64>,
|
||||
) {
|
||||
module.vmp_prepare_row(&mut self.key.0.data, row_i, col_j, &a.data);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -166,11 +147,8 @@ impl AutomorphismKey<Vec<u8>, FFT64> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
|
||||
where
|
||||
MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
|
||||
{
|
||||
pub fn generate_from_sk<DataSk>(
|
||||
impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> AutomorphismKey<DataSelf, FFT64> {
|
||||
pub fn generate_from_sk<DataSk: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
p: i64,
|
||||
@@ -179,9 +157,7 @@ where
|
||||
source_xe: &mut Source,
|
||||
sigma: f64,
|
||||
scratch: &mut Scratch,
|
||||
) where
|
||||
ScalarZnx<DataSk>: ScalarZnxToRef,
|
||||
{
|
||||
) {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(self.n(), module.n());
|
||||
@@ -209,12 +185,18 @@ where
|
||||
{
|
||||
(0..self.rank()).for_each(|i| {
|
||||
let (mut sk_inv_auto, _) = scratch_1.tmp_scalar_znx(module, 1);
|
||||
module.scalar_znx_automorphism(module.galois_element_inv(p), &mut sk_inv_auto, 0, sk, i);
|
||||
module.svp_prepare(&mut sk_out_dft, i, &sk_inv_auto, 0);
|
||||
module.scalar_znx_automorphism(
|
||||
module.galois_element_inv(p),
|
||||
&mut sk_inv_auto,
|
||||
0,
|
||||
&sk.data,
|
||||
i,
|
||||
);
|
||||
module.svp_prepare(&mut sk_out_dft.data, i, &sk_inv_auto, 0);
|
||||
});
|
||||
}
|
||||
|
||||
self.key.encrypt_sk(
|
||||
self.key.generate_from_sk(
|
||||
module,
|
||||
&sk,
|
||||
&sk_out_dft,
|
||||
@@ -228,20 +210,14 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
|
||||
where
|
||||
MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
|
||||
{
|
||||
pub fn automorphism<DataLhs, DataRhs>(
|
||||
impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> AutomorphismKey<DataSelf, FFT64> {
|
||||
pub fn automorphism<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
lhs: &AutomorphismKey<DataLhs, FFT64>,
|
||||
rhs: &AutomorphismKey<DataRhs, FFT64>,
|
||||
scratch: &mut Scratch,
|
||||
) where
|
||||
MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
{
|
||||
) {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(
|
||||
@@ -311,8 +287,8 @@ where
|
||||
// Applies back the automorphism X^{k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) -> (-pi^{-1}_{k'+k}(s)a + s, a)
|
||||
// and switches back to DFT domain
|
||||
(0..self.rank_out() + 1).for_each(|i| {
|
||||
module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft, i);
|
||||
module.vec_znx_dft(&mut tmp_dft, i, &tmp_idft, i);
|
||||
module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft.data, i);
|
||||
module.vec_znx_dft(&mut tmp_dft.data, i, &tmp_idft.data, i);
|
||||
});
|
||||
|
||||
// Sets back the relevant row
|
||||
@@ -331,65 +307,53 @@ where
|
||||
self.p = (lhs.p * rhs.p) % (module.cyclotomic_order() as i64);
|
||||
}
|
||||
|
||||
pub fn automorphism_inplace<DataRhs>(
|
||||
pub fn automorphism_inplace<DataRhs: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
rhs: &AutomorphismKey<DataRhs, FFT64>,
|
||||
scratch: &mut Scratch,
|
||||
) where
|
||||
MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
{
|
||||
) {
|
||||
unsafe {
|
||||
let self_ptr: *mut AutomorphismKey<DataSelf, FFT64> = self as *mut AutomorphismKey<DataSelf, FFT64>;
|
||||
self.automorphism(&module, &*self_ptr, rhs, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn keyswitch<DataLhs, DataRhs>(
|
||||
pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
lhs: &AutomorphismKey<DataLhs, FFT64>,
|
||||
rhs: &GLWESwitchingKey<DataRhs, FFT64>,
|
||||
scratch: &mut Scratch,
|
||||
) where
|
||||
MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
{
|
||||
) {
|
||||
self.key.keyswitch(module, &lhs.key, rhs, scratch);
|
||||
}
|
||||
|
||||
pub fn keyswitch_inplace<DataRhs>(
|
||||
pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
rhs: &AutomorphismKey<DataRhs, FFT64>,
|
||||
scratch: &mut Scratch,
|
||||
) where
|
||||
MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
{
|
||||
) {
|
||||
self.key.keyswitch_inplace(module, &rhs.key, scratch);
|
||||
}
|
||||
|
||||
pub fn external_product<DataLhs, DataRhs>(
|
||||
pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
lhs: &AutomorphismKey<DataLhs, FFT64>,
|
||||
rhs: &GGSWCiphertext<DataRhs, FFT64>,
|
||||
scratch: &mut Scratch,
|
||||
) where
|
||||
MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
{
|
||||
) {
|
||||
self.key.external_product(module, &lhs.key, rhs, scratch);
|
||||
}
|
||||
|
||||
pub fn external_product_inplace<DataRhs>(
|
||||
pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
module: &Module<FFT64>,
|
||||
rhs: &GGSWCiphertext<DataRhs, FFT64>,
|
||||
scratch: &mut Scratch,
|
||||
) where
|
||||
MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
|
||||
{
|
||||
) {
|
||||
self.key.external_product_inplace(module, rhs, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user