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