From f9dcddcce1b9babffc6f9fe5b21d25305f3c38ff Mon Sep 17 00:00:00 2001 From: Pro7ech Date: Sun, 16 Nov 2025 13:55:40 +0100 Subject: [PATCH] Add offset to blind retrieval --- .../tfhe/bdd_arithmetic/blind_retrieval.rs | 44 +++++++++++++++---- .../src/tfhe/bdd_arithmetic/blind_rotation.rs | 14 +++--- .../src/tfhe/bdd_arithmetic/eval.rs | 6 +-- .../tests/test_suite/glwe_blind_rotation.rs | 2 +- .../bdd_arithmetic/tests/test_suite/swap.rs | 16 ++++--- 5 files changed, 56 insertions(+), 26 deletions(-) diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_retrieval.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_retrieval.rs index 87b8852..f993975 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_retrieval.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_retrieval.rs @@ -17,21 +17,31 @@ impl GLWEBlindRetriever { where A: GLWEInfos, { - let log2_max_address: usize = (u32::BITS - (size as u32 - 1).leading_zeros()) as usize; + let bit_size: usize = (u32::BITS - (size as u32 - 1).leading_zeros()) as usize; Self { - accumulators: (0..log2_max_address) + accumulators: (0..bit_size) .map(|_| Accumulator::alloc(infos)) .collect_vec(), counter: 0, } } + pub fn retrieve_tmp_bytes(module: &M, res: &R, selector: &S) -> usize + where + M: Cmux, + R: GLWEInfos, + S: GGSWInfos, + { + module.cmux_tmp_bytes(res, res, selector) + } + pub fn retrieve( &mut self, module: &M, res: &mut R, data: &[A], selector: &S, + offset: usize, scratch: &mut Scratch, ) where M: GLWECopy + Cmux, @@ -41,14 +51,13 @@ impl GLWEBlindRetriever { Scratch: ScratchTakeCore, { self.reset(); - for ct in data { - self.add(module, ct, selector, scratch); + self.add(module, ct, selector, offset, scratch); } - self.flush(module, res, selector, scratch); + self.flush(module, res, selector, offset, scratch); } - pub fn add(&mut self, module: &M, a: &A, selector: &S, scratch: &mut Scratch) + pub fn add(&mut self, module: &M, a: &A, selector: &S, offset: usize, scratch: &mut Scratch) where A: GLWEToRef, S: GetGGSWBit, @@ -61,11 +70,19 @@ impl GLWEBlindRetriever { 1 << self.accumulators.len() ); - add_core(module, a, &mut self.accumulators, 0, selector, scratch); + add_core( + module, + a, + &mut self.accumulators, + 0, + selector, + offset, + scratch, + ); self.counter += 1; } - pub fn flush(&mut self, module: &M, res: &mut R, selector: &S, scratch: &mut Scratch) + pub fn flush(&mut self, module: &M, res: &mut R, selector: &S, offset: usize, scratch: &mut Scratch) where R: GLWEToMut, S: GetGGSWBit, @@ -81,6 +98,7 @@ impl GLWEBlindRetriever { acc_next, i + 1, selector, + offset, scratch, ); acc_prev[0].num = 0 @@ -94,6 +112,7 @@ impl GLWEBlindRetriever { for acc in self.accumulators.iter_mut() { acc.num = 0; } + self.counter = 0; } } @@ -120,6 +139,7 @@ fn add_core( accumulators: &mut [Accumulator], i: usize, selector: &S, + offset: usize, scratch: &mut Scratch, ) where A: GLWEToRef, @@ -136,7 +156,12 @@ fn add_core( acc_prev[0].num = 1; } 1 => { - module.cmux_inplace_neg(&mut acc_prev[0].data, a, &selector.get_bit(i), scratch); + module.cmux_inplace_neg( + &mut acc_prev[0].data, + a, + &selector.get_bit(i + offset), + scratch, + ); if !acc_next.is_empty() { add_core( @@ -145,6 +170,7 @@ fn add_core( acc_next, i + 1, selector, + offset, scratch, ); } diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_rotation.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_rotation.rs index 593abf7..fe1ae39 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_rotation.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/blind_rotation.rs @@ -11,14 +11,14 @@ use crate::tfhe::bdd_arithmetic::{Cmux, GetGGSWBit, UnsignedInteger}; impl GGSWBlindRotation for Module where - Self: GLWEBlindRotation + VecZnxAddScalarInplace + VecZnxNormalizeInplace, + Self: GLWEBlindRotation + VecZnxAddScalarInplace + VecZnxNormalizeInplace, Scratch: ScratchTakeCore, { } pub trait GGSWBlindRotation where - Self: GLWEBlindRotation + VecZnxAddScalarInplace + VecZnxNormalizeInplace, + Self: GLWEBlindRotation + VecZnxAddScalarInplace + VecZnxNormalizeInplace, { fn ggsw_to_ggsw_blind_rotation_tmp_bytes(&self, res_infos: &R, k_infos: &K) -> usize where @@ -161,14 +161,14 @@ where } } -impl GLWEBlindRotation for Module +impl GLWEBlindRotation for Module where Self: GLWECopy + GLWERotate + Cmux, Scratch: ScratchTakeCore, { } -pub trait GLWEBlindRotation +pub trait GLWEBlindRotation where Self: GLWECopy + GLWERotate + Cmux, { @@ -184,7 +184,7 @@ where fn glwe_blind_rotation_inplace( &self, res: &mut R, - fhe_uint: &K, + value: &K, sign: bool, bit_rsh: usize, bit_mask: usize, @@ -195,8 +195,6 @@ where K: GetGGSWBit, Scratch: ScratchTakeCore, { - assert!(bit_rsh + bit_mask <= T::BITS as usize); - let mut res: GLWE<&mut [u8]> = res.to_mut(); let (mut tmp_res, scratch_1) = scratch.take_glwe(&res); @@ -219,7 +217,7 @@ where } // b <- (b - a) * GGSW(b[i]) + a - self.cmux_inplace(b, a, &fhe_uint.get_bit(i + bit_rsh), scratch_1); + self.cmux_inplace(b, a, &value.get_bit(i + bit_rsh), scratch_1); // ping-pong roles for next iter a_is_res = !a_is_res; diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs index 75ba7cf..b92b783 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/eval.rs @@ -471,16 +471,16 @@ where + VecZnxBigNormalize + VecZnxBigNormalizeTmpBytes, { - fn cmux_tmp_bytes(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize + fn cmux_tmp_bytes(&self, res_infos: &R, a_infos: &A, selector_infos: &B) -> usize where R: GLWEInfos, A: GLWEInfos, B: GGSWInfos, { - let res_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank() + 1).into(), b_infos.size()); + let res_dft: usize = self.bytes_of_vec_znx_dft((selector_infos.rank() + 1).into(), selector_infos.size()); res_dft + self - .glwe_external_product_internal_tmp_bytes(res_infos, a_infos, b_infos) + .glwe_external_product_internal_tmp_bytes(res_infos, a_infos, selector_infos) .max(self.vec_znx_big_normalize_tmp_bytes()) } diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/glwe_blind_rotation.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/glwe_blind_rotation.rs index de7b37f..54baf97 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/glwe_blind_rotation.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/glwe_blind_rotation.rs @@ -26,7 +26,7 @@ where + GLWESecretPreparedFactory + GGSWPreparedFactory + GGSWEncryptSk - + GLWEBlindRotation + + GLWEBlindRotation + GLWEDecrypt + GLWEEncryptSk, ScratchOwned: ScratchOwnedAlloc + ScratchOwnedBorrow, diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/swap.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/swap.rs index a9e5d0c..3535254 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/swap.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/swap.rs @@ -178,22 +178,28 @@ where }) .collect_vec(); + let mut retriever: GLWEBlindRetriever = GLWEBlindRetriever::alloc(&glwe_infos, data.len()); for idx in 0..data.len() as u32 { + let offset = 2; let mut idx_enc: FheUintPrepared, u32, BE> = FheUintPrepared::alloc_from_infos(module, &ggsw_infos); idx_enc.encrypt_sk( module, - idx, + idx << offset, sk, &mut source_xa, &mut source_xe, scratch.borrow(), ); - let mut retriever: GLWEBlindRetriever = GLWEBlindRetriever::alloc(&glwe_infos, 25); let mut res: FheUint, u32> = FheUint::alloc_from_infos(&glwe_infos); - retriever.retrieve(module, &mut res, &data_enc, &idx_enc, scratch.borrow()); - - println!("{}", res.decrypt(module, sk, scratch.borrow())); + retriever.retrieve( + module, + &mut res, + &data_enc, + &idx_enc, + offset, + scratch.borrow(), + ); assert_eq!( data[idx as usize],