Update to custom fheuint prepare

This commit is contained in:
Pro7ech
2025-11-07 08:49:32 +01:00
parent 5cf184d950
commit 1d23dfc078
10 changed files with 112 additions and 34 deletions

View File

@@ -217,18 +217,19 @@ pub trait FheUintPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend
DB: DataRef,
DK: DataRef,
K: BDDKeyHelper<DK, BRA, BE>;
fn fhe_uint_prepare_partial<DM, DB, DK, K>(
fn fhe_uint_prepare_custom<DM, DB, DK, K>(
&self,
res: &mut FheUintPrepared<DM, T, BE>,
bits: &FheUint<DB, T>,
count: usize,
bit_start: usize,
bit_end: usize,
key: &K,
scratch: &mut Scratch<BE>,
) where
DM: DataMut,
DB: DataRef,
DK: DataRef,
K: BDDKeyHelper<DK, BRA, BE>;
K: BDDKeyHelper<DK, BRA, BE>;
}
impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintPrepare<BRA, T, BE> for Module<BE>
@@ -261,22 +262,15 @@ where
DK: DataRef,
K: BDDKeyHelper<DK, BRA, BE>,
{
let (cbt, ks) = key.get_cbt_key();
let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(bits); //TODO: add TakeLWE
let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
for (bit, dst) in res.bits.iter_mut().enumerate() {
bits.get_bit_lwe(self, bit, &mut lwe, ks, scratch_1);
cbt.execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
dst.prepare(self, &tmp_ggsw, scratch_1);
}
self.fhe_uint_prepare_custom(res, bits, 0, T::BITS as usize, key, scratch);
}
fn fhe_uint_prepare_partial<DM, DB, DK, K>(
fn fhe_uint_prepare_custom<DM, DB, DK, K>(
&self,
res: &mut FheUintPrepared<DM, T, BE>,
bits: &FheUint<DB, T>,
count: usize,
bit_start: usize,
bit_end: usize,
key: &K,
scratch: &mut Scratch<BE>,
) where
@@ -289,12 +283,21 @@ where
let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(bits); //TODO: add TakeLWE
let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
for (bit, dst) in res.bits[0..count].iter_mut().enumerate() { // TODO: set the rest of the bits to a prepared zero GGSW
for (bit, dst) in res.bits[bit_start..bit_end].iter_mut().enumerate() {
// TODO: set the rest of the bits to a prepared zero GGSW
bits.get_bit_lwe(self, bit, &mut lwe, ks, scratch_1);
cbt.execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
dst.prepare(self, &tmp_ggsw, scratch_1);
}
}
for i in 0..bit_start {
res.bits[i].zero(self);
}
for i in bit_end..T::BITS as usize {
res.bits[i].zero(self);
}
}
}
impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
@@ -309,8 +312,15 @@ impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
{
module.fhe_uint_prepare(self, other, key, scratch);
}
pub fn prepare_partial<BRA, M, O, K, DK>(&mut self, module: &M, other: &FheUint<O, T>, count: usize, key: &K, scratch: &mut Scratch<BE>)
where
pub fn prepare_partial<BRA, M, O, K, DK>(
&mut self,
module: &M,
other: &FheUint<O, T>,
bit_start: usize,
bit_end: usize,
key: &K,
scratch: &mut Scratch<BE>,
) where
BRA: BlindRotationAlgo,
O: DataRef,
DK: DataRef,
@@ -318,6 +328,6 @@ impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
M: FheUintPrepare<BRA, T, BE>,
Scratch<BE>: ScratchTakeCore<BE>,
{
module.fhe_uint_prepare_partial(self, other, count, key, scratch);
}
module.fhe_uint_prepare_custom(self, other, bit_start, bit_end, key, scratch);
}
}