Add support for blind retrieval

This commit is contained in:
Pro7ech
2025-11-15 22:41:11 +01:00
parent 28102b684f
commit b062c722a0
12 changed files with 882 additions and 25 deletions

View File

@@ -1,16 +1,16 @@
use poulpy_hal::{
api::{
ModuleN, ScratchTakeBasic, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply, VecZnxDftBytesOf,
VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
VmpApplyDftToDftTmpBytes,
ModuleN, ScratchTakeBasic, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply,
VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
},
layouts::{Backend, DataMut, DataViewMut, Module, Scratch, VecZnx, VecZnxBig, VecZnxDft},
};
use crate::{
ScratchTakeCore,
GLWENormalize, ScratchTakeCore,
layouts::{
GGSWInfos, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos,
GGSWInfos, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, LWEInfos,
prepared::{GGSWPrepared, GGSWPreparedToRef},
},
};
@@ -67,11 +67,22 @@ pub trait GLWEExternalProduct<BE: Backend> {
A: GLWEToRef,
D: GGSWPreparedToRef<BE> + GGSWInfos,
Scratch<BE>: ScratchTakeCore<BE>;
fn glwe_external_product_add<R, A, D>(&self, res: &mut R, lhs: &A, rhs: &D, scratch: &mut Scratch<BE>)
where
R: GLWEToMut,
A: GLWEToRef,
D: GGSWPreparedToRef<BE> + GGSWInfos,
Scratch<BE>: ScratchTakeCore<BE>;
}
impl<BE: Backend> GLWEExternalProduct<BE> for Module<BE>
where
Self: GLWEExternalProductInternal<BE> + VecZnxDftBytesOf + VecZnxBigNormalize<BE> + VecZnxBigNormalizeTmpBytes,
Self: GLWEExternalProductInternal<BE>
+ VecZnxDftBytesOf
+ VecZnxBigNormalize<BE>
+ VecZnxBigNormalizeTmpBytes
+ VecZnxBigAddSmallInplace<BE>
+ GLWENormalize<BE>,
{
fn glwe_external_product_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
where
@@ -163,6 +174,80 @@ where
);
}
}
fn glwe_external_product_add<R, A, D>(&self, res: &mut R, a: &A, key: &D, scratch: &mut Scratch<BE>)
where
R: GLWEToMut,
A: GLWEToRef,
D: GGSWPreparedToRef<BE>,
Scratch<BE>: ScratchTakeCore<BE>,
{
let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
let a: &GLWE<&[u8]> = &a.to_ref();
let key: &GGSWPrepared<&[u8], BE> = &key.to_ref();
assert_eq!(a.base2k(), res.base2k());
let res_base2k: usize = res.base2k().into();
let key_base2k: usize = key.base2k().into();
#[cfg(debug_assertions)]
{
use poulpy_hal::api::ScratchAvailable;
assert_eq!(key.rank(), a.rank());
assert_eq!(key.rank(), res.rank());
assert_eq!(key.n(), res.n());
assert_eq!(a.n(), res.n());
assert!(scratch.available() >= self.glwe_external_product_tmp_bytes(res, a, key));
}
if res_base2k == key_base2k {
let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // Todo optimise
let mut res_big = self.glwe_external_product_internal(res_dft, a, key, scratch_1);
for j in 0..(res.rank() + 1).into() {
self.vec_znx_big_add_small_inplace(&mut res_big, j, res.data(), j);
self.vec_znx_big_normalize(
res_base2k,
&mut res.data,
j,
key_base2k,
&res_big,
j,
scratch_1,
);
}
} else {
let (mut a_conv, scratch_1) = scratch.take_glwe(&GLWELayout {
n: a.n(),
base2k: key.base2k(),
k: a.k(),
rank: a.rank(),
});
let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
n: res.n(),
base2k: key.base2k(),
k: res.k(),
rank: res.rank(),
});
self.glwe_normalize(&mut a_conv, a, scratch_2);
self.glwe_normalize(&mut res_conv, res, scratch_2);
let (res_dft, scratch_2) = scratch_2.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // Todo optimise
let mut res_big = self.glwe_external_product_internal(res_dft, &a_conv, key, scratch_2);
for j in 0..(res.rank() + 1).into() {
self.vec_znx_big_add_small_inplace(&mut res_big, j, res_conv.data(), j);
self.vec_znx_big_normalize(
res_base2k,
&mut res.data,
j,
key_base2k,
&res_big,
j,
scratch_2,
);
}
}
}
}
pub trait GLWEExternalProductInternal<BE: Backend> {

View File

@@ -7,7 +7,7 @@ use poulpy_hal::{
use crate::{
GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, GLWETrace, ScratchTakeCore,
layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement},
layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEToMut, GetGaloisElement},
};
pub trait GLWEPacking<BE: Backend> {
/// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
@@ -21,7 +21,7 @@ pub trait GLWEPacking<BE: Backend> {
scratch: &mut Scratch<BE>,
) where
R: GLWEToMut + GLWEInfos,
A: GLWEToMut + GLWEToRef + GLWEInfos,
A: GLWEToMut + GLWEInfos,
K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
H: GLWEAutomorphismKeyHelper<K, BE>;
}
@@ -51,7 +51,7 @@ where
scratch: &mut Scratch<BE>,
) where
R: GLWEToMut + GLWEInfos,
A: GLWEToMut + GLWEToRef + GLWEInfos,
A: GLWEToMut + GLWEInfos,
K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
H: GLWEAutomorphismKeyHelper<K, BE>,
{
@@ -97,8 +97,8 @@ fn pack_internal<M, A, B, K, BE: Backend>(
scratch: &mut Scratch<BE>,
) where
M: GLWEAutomorphism<BE> + GLWERotate<BE> + GLWESub + GLWEShift<BE> + GLWEAdd + GLWENormalize<BE>,
A: GLWEToMut + GLWEToRef + GLWEInfos,
B: GLWEToMut + GLWEToRef + GLWEInfos,
A: GLWEToMut + GLWEInfos,
B: GLWEToMut + GLWEInfos,
K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{

View File

@@ -189,7 +189,7 @@ impl<D: DataRef> WriterTo for GLWE<D> {
}
}
pub trait GLWEToRef {
pub trait GLWEToRef: Sized {
fn to_ref(&self) -> GLWE<&[u8]>;
}
@@ -203,14 +203,11 @@ impl<D: DataRef> GLWEToRef for GLWE<D> {
}
}
pub trait GLWEToMut {
pub trait GLWEToMut: GLWEToRef {
fn to_mut(&mut self) -> GLWE<&mut [u8]>;
}
impl<D: DataMut> GLWEToMut for GLWE<D>
where
Self: GLWEToRef,
{
impl<D: DataMut> GLWEToMut for GLWE<D> {
fn to_mut(&mut self) -> GLWE<&mut [u8]> {
GLWE {
k: self.k,

View File

@@ -402,6 +402,84 @@ where
self.vec_znx_normalize_tmp_bytes()
}
/// Usage:
/// let mut tmp_b: Option<GLWE<&mut [u8]>> = None;
/// let (b_conv, scratch_1) = glwe_maybe_convert_in_place(self, b, res.base2k().as_u32(), &mut tmp_b, scratch);
fn glwe_maybe_cross_normalize_to_ref<'a, A>(
&self,
glwe: &'a A,
target_base2k: usize,
tmp_slot: &'a mut Option<GLWE<&'a mut [u8]>>, // caller-owned scratch-backed temp
scratch: &'a mut Scratch<BE>,
) -> (GLWE<&'a [u8]>, &'a mut Scratch<BE>)
where
A: GLWEToRef + GLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
// No conversion: just use the original GLWE
if glwe.base2k().as_usize() == target_base2k {
// Drop any previous temp; it's stale for this base
tmp_slot.take();
return (glwe.to_ref(), scratch);
}
// Conversion: allocate a temporary GLWE in scratch
let mut layout = glwe.glwe_layout();
layout.base2k = target_base2k.into();
let (tmp, scratch2) = scratch.take_glwe(&layout);
*tmp_slot = Some(tmp);
// Get a mutable handle to the temp and normalize into it
let tmp_ref: &mut GLWE<&mut [u8]> = tmp_slot
.as_mut()
.expect("tmp_slot just set to Some, but found None");
self.glwe_normalize(tmp_ref, glwe, scratch2);
// Return a trait-object view of the temp
(tmp_ref.to_ref(), scratch2)
}
/// Usage:
/// let mut tmp_b: Option<GLWE<&mut [u8]>> = None;
/// let (b_conv, scratch_1) = glwe_maybe_convert_in_place(self, b, res.base2k().as_u32(), &mut tmp_b, scratch);
fn glwe_maybe_cross_normalize_to_mut<'a, A>(
&self,
glwe: &'a mut A,
target_base2k: usize,
tmp_slot: &'a mut Option<GLWE<&'a mut [u8]>>, // caller-owned scratch-backed temp
scratch: &'a mut Scratch<BE>,
) -> (GLWE<&'a mut [u8]>, &'a mut Scratch<BE>)
where
A: GLWEToMut + GLWEInfos,
Scratch<BE>: ScratchTakeCore<BE>,
{
// No conversion: just use the original GLWE
if glwe.base2k().as_usize() == target_base2k {
// Drop any previous temp; it's stale for this base
tmp_slot.take();
return (glwe.to_mut(), scratch);
}
// Conversion: allocate a temporary GLWE in scratch
let mut layout = glwe.glwe_layout();
layout.base2k = target_base2k.into();
let (tmp, scratch2) = scratch.take_glwe(&layout);
*tmp_slot = Some(tmp);
// Get a mutable handle to the temp and normalize into it
let tmp_ref: &mut GLWE<&mut [u8]> = tmp_slot
.as_mut()
.expect("tmp_slot just set to Some, but found None");
self.glwe_normalize(tmp_ref, glwe, scratch2);
// Return a trait-object view of the temp
(tmp_ref.to_mut(), scratch2)
}
fn glwe_normalize<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
where
R: GLWEToMut,