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