Files
poulpy/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/block_prepared.rs
Pro7ech 2b2b994f7d wip
2025-10-16 09:31:58 +02:00

281 lines
8.7 KiB
Rust

use std::marker::PhantomData;
use poulpy_core::layouts::{Base2K, Dnum, Dsize, GGSWInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision, prepared::GGSWPrepared};
#[cfg(test)]
use poulpy_core::{
TakeGGSW,
layouts::{GGSW, prepared::GLWESecretPrepared},
};
use poulpy_hal::{
api::VmpPMatAlloc,
layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
};
#[cfg(test)]
use poulpy_hal::{
api::{
ScratchAvailable, SvpApplyDftToDftInplace, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal,
VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigBytesOf,
VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub,
VecZnxSubInplace, VmpPrepare,
},
oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl},
source::Source,
};
use crate::tfhe::bdd_arithmetic::{FheUintBlocks, FheUintPrepare, ToBits, UnsignedInteger};
#[cfg(test)]
pub(crate) struct FheUintBlocksPrepDebug<D: Data, T: UnsignedInteger> {
pub(crate) blocks: Vec<GGSW<D>>,
pub(crate) _base: u8,
pub(crate) _phantom: PhantomData<T>,
}
#[cfg(test)]
impl<T: UnsignedInteger> FheUintBlocksPrepDebug<Vec<u8>, T> {
#[allow(dead_code)]
pub(crate) fn alloc<A, BE: Backend>(module: &Module<BE>, infos: &A) -> Self
where
A: GGSWInfos,
{
Self::alloc_with(
module,
infos.base2k(),
infos.k(),
infos.dnum(),
infos.dsize(),
infos.rank(),
)
}
#[allow(dead_code)]
pub(crate) fn alloc_with<BE: Backend>(
module: &Module<BE>,
base2k: Base2K,
k: TorusPrecision,
dnum: Dnum,
dsize: Dsize,
rank: Rank,
) -> Self {
Self {
blocks: (0..T::WORD_SIZE)
.map(|_| GGSW::alloc(module.n().into(), base2k, k, rank, dnum, dsize))
.collect(),
_base: 1,
_phantom: PhantomData,
}
}
}
/// A prepared FHE ciphertext encrypting the bits of an [UnsignedInteger].
pub struct FheUintBlocksPrep<D: Data, B: Backend, T: UnsignedInteger> {
pub(crate) blocks: Vec<GGSWPrepared<D, B>>,
pub(crate) _base: u8,
pub(crate) _phantom: PhantomData<T>,
}
impl<T: UnsignedInteger, BE: Backend> FheUintBlocksPrep<Vec<u8>, BE, T>
where
Module<BE>: VmpPMatAlloc<BE>,
{
#[allow(dead_code)]
pub(crate) fn alloc<A>(module: &Module<BE>, infos: &A) -> Self
where
A: GGSWInfos,
{
Self::alloc_with(
module,
infos.base2k(),
infos.k(),
infos.dnum(),
infos.dsize(),
infos.rank(),
)
}
#[allow(dead_code)]
pub(crate) fn alloc_with(module: &Module<BE>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
where
Module<BE>: VmpPMatAlloc<BE>,
{
Self {
blocks: (0..T::WORD_SIZE)
.map(|_| GGSWPrepared::alloc(module, base2k, k, dnum, dsize, rank))
.collect(),
_base: 1,
_phantom: PhantomData,
}
}
}
impl<D: DataMut, T: UnsignedInteger + ToBits, BE: Backend> FheUintBlocksPrep<D, BE, T> {
#[allow(dead_code)]
#[cfg(test)]
pub(crate) fn encrypt_sk<S>(
&mut self,
module: &Module<BE>,
value: T,
sk: &GLWESecretPrepared<S, BE>,
source_xa: &mut Source,
source_xe: &mut Source,
scratch: &mut Scratch<BE>,
) where
S: DataRef,
Module<BE>: VecZnxAddScalarInplace
+ VecZnxDftBytesOf
+ VecZnxBigNormalize<BE>
+ VecZnxDftApply<BE>
+ SvpApplyDftToDftInplace<BE>
+ VecZnxIdftApplyConsume<BE>
+ VecZnxNormalizeTmpBytes
+ VecZnxFillUniform
+ VecZnxSubInplace
+ VecZnxAddInplace
+ VecZnxNormalizeInplace<BE>
+ VecZnxAddNormal
+ VecZnxNormalize<BE>
+ VecZnxSub
+ VmpPrepare<BE>,
Scratch<BE>: TakeVecZnxDft<BE> + ScratchAvailable + TakeVecZnx + TakeGGSW + TakeScalarZnx,
{
#[cfg(debug_assertions)]
{
assert!(module.n().is_multiple_of(T::WORD_SIZE));
assert_eq!(self.n(), module.n() as u32);
assert_eq!(sk.n(), module.n() as u32);
}
let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(self);
let (mut pt, scratch_2) = scratch_1.take_scalar_znx(module.n(), 1);
for i in 0..T::WORD_SIZE {
use poulpy_core::layouts::prepared::Prepare;
use poulpy_hal::layouts::ZnxViewMut;
pt.at_mut(0, 0)[0] = value.bit(i) as i64;
tmp_ggsw.encrypt_sk(&module, &pt, sk, source_xa, source_xe, scratch_2);
self.blocks[i].prepare(module, &tmp_ggsw, scratch_2);
}
}
/// Prepares [FheUintBits] to [FheUintBitsPrep].
pub fn prepare<BIT, KEY>(&mut self, module: &Module<BE>, bits: &FheUintBlocks<BIT, T>, key: &KEY, scratch: &mut Scratch<BE>)
where
BIT: DataRef,
KEY: FheUintPrepare<BE, FheUintBlocksPrep<D, BE, T>, FheUintBlocks<BIT, T>>,
{
key.prepare(module, self, bits, scratch);
}
}
#[cfg(test)]
impl<D: DataMut, T: UnsignedInteger + ToBits> FheUintBlocksPrepDebug<D, T> {
pub(crate) fn prepare<BIT, KEY, BE: Backend>(
&mut self,
module: &Module<BE>,
bits: &FheUintBlocks<BIT, T>,
key: &KEY,
scratch: &mut Scratch<BE>,
) where
BIT: DataRef,
KEY: FheUintPrepare<BE, FheUintBlocksPrepDebug<D, T>, FheUintBlocks<BIT, T>>,
{
key.prepare(module, self, bits, scratch);
}
}
#[cfg(test)]
impl<D: DataRef, T: UnsignedInteger + ToBits> FheUintBlocksPrepDebug<D, T> {
#[allow(dead_code)]
pub(crate) fn noise<S: DataRef, BE: Backend>(&self, module: &Module<BE>, sk: &GLWESecretPrepared<S, BE>, want: T)
where
Module<BE>: VecZnxDftBytesOf
+ VecZnxBigBytesOf
+ VecZnxDftApply<BE>
+ SvpApplyDftToDftInplace<BE>
+ VecZnxIdftApplyConsume<BE>
+ VecZnxBigAddInplace<BE>
+ VecZnxBigAddSmallInplace<BE>
+ VecZnxBigNormalize<BE>
+ VecZnxNormalizeTmpBytes
+ VecZnxBigAlloc<BE>
+ VecZnxDftAlloc<BE>
+ VecZnxBigNormalizeTmpBytes
+ VecZnxIdftApplyTmpA<BE>
+ VecZnxAddScalarInplace
+ VecZnxSubInplace,
BE: Backend + TakeVecZnxDftImpl<BE> + TakeVecZnxBigImpl<BE> + ScratchOwnedAllocImpl<BE> + ScratchOwnedBorrowImpl<BE>,
{
for (i, ggsw) in self.blocks.iter().enumerate() {
use poulpy_hal::layouts::{ScalarZnx, ZnxViewMut};
let mut pt_want = ScalarZnx::alloc(self.n().into(), 1);
pt_want.at_mut(0, 0)[0] = want.bit(i) as i64;
ggsw.print_noise(module, sk, &pt_want);
}
}
}
impl<D: DataRef, T: UnsignedInteger, B: Backend> LWEInfos for FheUintBlocksPrep<D, B, T> {
fn base2k(&self) -> poulpy_core::layouts::Base2K {
self.blocks[0].base2k()
}
fn k(&self) -> poulpy_core::layouts::TorusPrecision {
self.blocks[0].k()
}
fn n(&self) -> poulpy_core::layouts::RingDegree {
self.blocks[0].n()
}
}
impl<D: DataRef, T: UnsignedInteger, B: Backend> GLWEInfos for FheUintBlocksPrep<D, B, T> {
fn rank(&self) -> poulpy_core::layouts::Rank {
self.blocks[0].rank()
}
}
impl<D: DataRef, T: UnsignedInteger, B: Backend> GGSWInfos for FheUintBlocksPrep<D, B, T> {
fn dsize(&self) -> poulpy_core::layouts::Dsize {
self.blocks[0].dsize()
}
fn dnum(&self) -> poulpy_core::layouts::Dnum {
self.blocks[0].dnum()
}
}
#[cfg(test)]
impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintBlocksPrepDebug<D, T> {
fn base2k(&self) -> poulpy_core::layouts::Base2K {
self.blocks[0].base2k()
}
fn k(&self) -> poulpy_core::layouts::TorusPrecision {
self.blocks[0].k()
}
fn n(&self) -> poulpy_core::layouts::RingDegree {
self.blocks[0].n()
}
}
#[cfg(test)]
impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintBlocksPrepDebug<D, T> {
fn rank(&self) -> poulpy_core::layouts::Rank {
self.blocks[0].rank()
}
}
#[cfg(test)]
impl<D: DataRef, T: UnsignedInteger> GGSWInfos for FheUintBlocksPrepDebug<D, T> {
fn dsize(&self) -> poulpy_core::layouts::Dsize {
self.blocks[0].dsize()
}
fn dnum(&self) -> poulpy_core::layouts::Dnum {
self.blocks[0].dnum()
}
}