1use std::marker::PhantomData;
  2
  3use poulpy_core::layouts::{Base2K, Dnum, Dsize, GGSWInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision, prepared::GGSWPrepared};
  4#[cfg(test)]
  5use poulpy_core::{
  6    layouts::{prepared::GLWESecretPrepared, GGSW}, ScratchTakeCore,
  7};
  8use poulpy_hal::{
  9    api::VmpPMatAlloc,
 10    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 11};
 12#[cfg(test)]
 13use poulpy_hal::{
 14    api::{
 15        SvpApplyDftToDftInplace, VecZnxAddInplace, VecZnxAddNormal,
 16        VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigBytesOf,
 17        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
 18        VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub,
 19        VecZnxSubInplace, VmpPrepare,
 20    },
 21    source::Source,
 22};
 23
 24use crate::tfhe::bdd_arithmetic::{FheUintBlocks, FheUintPrepare, ToBits, UnsignedInteger};
 25
 26#[cfg(test)]
 27pub(crate) struct FheUintBlocksPrepDebug<D: Data, T: UnsignedInteger> {
 28    pub(crate) blocks: Vec<GGSW<D>>,
 29    pub(crate) _base: u8,
 30    pub(crate) _phantom: PhantomData<T>,
 31}
 32
 33#[cfg(test)]
 34impl<T: UnsignedInteger> FheUintBlocksPrepDebug<Vec<u8>, T> {
 35    #[allow(dead_code)]
 36    pub(crate) fn alloc<A, BE: Backend>(module: &Module<BE>, infos: &A) -> Self
 37    where
 38        A: GGSWInfos,
 39    {
 40        Self::alloc_with(
 41            module,
 42            infos.base2k(),
 43            infos.k(),
 44            infos.dnum(),
 45            infos.dsize(),
 46            infos.rank(),
 47        )
 48    }
 49
 50    #[allow(dead_code)]
 51    pub(crate) fn alloc_with<BE: Backend>(
 52        module: &Module<BE>,
 53        base2k: Base2K,
 54        k: TorusPrecision,
 55        dnum: Dnum,
 56        dsize: Dsize,
 57        rank: Rank,
 58    ) -> Self {
 59        Self {
 60            blocks: (0..T::WORD_SIZE)
 61                .map(|_| GGSW::alloc(module.n().into(), base2k, k, rank, dnum, dsize))
 62                .collect(),
 63            _base: 1,
 64            _phantom: PhantomData,
 65        }
 66    }
 67}
 68
 69/// A prepared FHE ciphertext encrypting the bits of an [UnsignedInteger].
 70pub struct FheUintBlocksPrep<D: Data, B: Backend, T: UnsignedInteger> {
 71    pub(crate) blocks: Vec<GGSWPrepared<D, B>>,
 72    pub(crate) _base: u8,
 73    pub(crate) _phantom: PhantomData<T>,
 74}
 75
 76impl<T: UnsignedInteger, BE: Backend> FheUintBlocksPrep<Vec<u8>, BE, T>
 77where
 78    Module<BE>: VmpPMatAlloc<BE>,
 79{
 80    #[allow(dead_code)]
 81    pub(crate) fn alloc<A>(module: &Module<BE>, infos: &A) -> Self
 82    where
 83        A: GGSWInfos,
 84    {
 85        Self::alloc_with(
 86            module,
 87            infos.base2k(),
 88            infos.k(),
 89            infos.dnum(),
 90            infos.dsize(),
 91            infos.rank(),
 92        )
 93    }
 94
 95    #[allow(dead_code)]
 96    pub(crate) fn alloc_with(module: &Module<BE>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
 97    where
 98        Module<BE>: VmpPMatAlloc<BE>,
 99    {
100        Self {
101            blocks: (0..T::WORD_SIZE)
102                .map(|_| GGSWPrepared::alloc(module, base2k, k, dnum, dsize, rank))
103                .collect(),
104            _base: 1,
105            _phantom: PhantomData,
106        }
107    }
108}
109
110impl<D: DataMut, T: UnsignedInteger + ToBits, BE: Backend> FheUintBlocksPrep<D, BE, T> {
111    #[allow(dead_code)]
112    #[cfg(test)]
113    pub(crate) fn encrypt_sk<S>(
114        &mut self,
115        module: &Module<BE>,
116        value: T,
117        sk: &GLWESecretPrepared<S, BE>,
118        source_xa: &mut Source,
119        source_xe: &mut Source,
120        scratch: &mut Scratch<BE>,
121    ) where
122        S: DataRef,
123        Module<BE>: VecZnxAddScalarInplace
124            + VecZnxDftBytesOf
125            + VecZnxBigNormalize<BE>
126            + VecZnxDftApply<BE>
127            + SvpApplyDftToDftInplace<BE>
128            + VecZnxIdftApplyConsume<BE>
129            + VecZnxNormalizeTmpBytes
130            + VecZnxFillUniform
131            + VecZnxSubInplace
132            + VecZnxAddInplace
133            + VecZnxNormalizeInplace<BE>
134            + VecZnxAddNormal
135            + VecZnxNormalize<BE>
136            + VecZnxSub
137            + VmpPrepare<BE>,
138        Scratch<BE>: ScratchTakeCore<BE>,
139    {
140        #[cfg(debug_assertions)]
141        {
142            assert!(module.n().is_multiple_of(T::WORD_SIZE));
143            assert_eq!(self.n(), module.n() as u32);
144            assert_eq!(sk.n(), module.n() as u32);
145        }
146
147        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(module, self);
148        let (mut pt, scratch_2) = scratch_1.take_scalar_znx(module.n(), 1);
149
150        for i in 0..T::WORD_SIZE {
151
152            use poulpy_hal::layouts::ZnxViewMut;
153
154            pt.at_mut(0, 0)[0] = value.bit(i) as i64;
155            tmp_ggsw.encrypt_sk(&module, &pt, sk, source_xa, source_xe, scratch_2);
156            self.blocks[i].prepare(module, &tmp_ggsw, scratch_2);
157        }
158    }
159
160    /// Prepares [FheUintBits] to [FheUintBitsPrep].
161    pub fn prepare<BIT, KEY>(&mut self, module: &Module<BE>, bits: &FheUintBlocks<BIT, T>, key: &KEY, scratch: &mut Scratch<BE>)
162    where
163        BIT: DataRef,
164        KEY: FheUintPrepare<BE, FheUintBlocksPrep<D, BE, T>, FheUintBlocks<BIT, T>>,
165    {
166        key.prepare(module, self, bits, scratch);
167    }
168}
169
170#[cfg(test)]
171impl<D: DataMut, T: UnsignedInteger + ToBits> FheUintBlocksPrepDebug<D, T> {
172    pub(crate) fn prepare<BIT, KEY, BE: Backend>(
173        &mut self,
174        module: &Module<BE>,
175        bits: &FheUintBlocks<BIT, T>,
176        key: &KEY,
177        scratch: &mut Scratch<BE>,
178    ) where
179        BIT: DataRef,
180        KEY: FheUintPrepare<BE, FheUintBlocksPrepDebug<D, T>, FheUintBlocks<BIT, T>>,
181    {
182        key.prepare(module, self, bits, scratch);
183    }
184}
185
186#[cfg(test)]
187impl<D: DataRef, T: UnsignedInteger + ToBits> FheUintBlocksPrepDebug<D, T> {
188    #[allow(dead_code)]
189    pub(crate) fn noise<S: DataRef, BE: Backend>(&self, module: &Module<BE>, sk: &GLWESecretPrepared<S, BE>, want: T)
190    where
191        Module<BE>: VecZnxDftBytesOf
192            + VecZnxBigBytesOf
193            + VecZnxDftApply<BE>
194            + SvpApplyDftToDftInplace<BE>
195            + VecZnxIdftApplyConsume<BE>
196            + VecZnxBigAddInplace<BE>
197            + VecZnxBigAddSmallInplace<BE>
198            + VecZnxBigNormalize<BE>
199            + VecZnxNormalizeTmpBytes
200            + VecZnxBigAlloc<BE>
201            + VecZnxDftAlloc<BE>
202            + VecZnxBigNormalizeTmpBytes
203            + VecZnxIdftApplyTmpA<BE>
204            + VecZnxAddScalarInplace
205            + VecZnxSubInplace,
206    {
207        for (i, ggsw) in self.blocks.iter().enumerate() {
208            use poulpy_hal::layouts::{ScalarZnx, ZnxViewMut};
209            let mut pt_want = ScalarZnx::alloc(self.n().into(), 1);
210            pt_want.at_mut(0, 0)[0] = want.bit(i) as i64;
211            ggsw.print_noise(module, sk, &pt_want);
212        }
213    }
214}
215
216impl<D: DataRef, T: UnsignedInteger, B: Backend> LWEInfos for FheUintBlocksPrep<D, B, T> {
217    fn base2k(&self) -> poulpy_core::layouts::Base2K {
218        self.blocks[0].base2k()
219    }
220
221    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
222        self.blocks[0].k()
223    }
224
225    fn n(&self) -> poulpy_core::layouts::Degree {
226        self.blocks[0].n()
227    }
228}
229
230impl<D: DataRef, T: UnsignedInteger, B: Backend> GLWEInfos for FheUintBlocksPrep<D, B, T> {
231    fn rank(&self) -> poulpy_core::layouts::Rank {
232        self.blocks[0].rank()
233    }
234}
235
236impl<D: DataRef, T: UnsignedInteger, B: Backend> GGSWInfos for FheUintBlocksPrep<D, B, T> {
237    fn dsize(&self) -> poulpy_core::layouts::Dsize {
238        self.blocks[0].dsize()
239    }
240
241    fn dnum(&self) -> poulpy_core::layouts::Dnum {
242        self.blocks[0].dnum()
243    }
244}
245
246#[cfg(test)]
247impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintBlocksPrepDebug<D, T> {
248    fn base2k(&self) -> poulpy_core::layouts::Base2K {
249        self.blocks[0].base2k()
250    }
251
252    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
253        self.blocks[0].k()
254    }
255
256    fn n(&self) -> poulpy_core::layouts::Degree {
257        self.blocks[0].n()
258    }
259}
260
261#[cfg(test)]
262impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintBlocksPrepDebug<D, T> {
263    fn rank(&self) -> poulpy_core::layouts::Rank {
264        self.blocks[0].rank()
265    }
266}
267
268#[cfg(test)]
269impl<D: DataRef, T: UnsignedInteger> GGSWInfos for FheUintBlocksPrepDebug<D, T> {
270    fn dsize(&self) -> poulpy_core::layouts::Dsize {
271        self.blocks[0].dsize()
272    }
273
274    fn dnum(&self) -> poulpy_core::layouts::Dnum {
275        self.blocks[0].dnum()
276    }
277}