1use std::marker::PhantomData;
  2
  3use poulpy_core::layouts::{Base2K, GLWE, GLWEInfos, GLWEPlaintextLayout, LWEInfos, Rank, TorusPrecision};
  4
  5#[cfg(test)]
  6use poulpy_core::ScratchTakeCore;
  7use poulpy_core::{layouts::prepared::GLWESecretPrepared};
  8use poulpy_hal::api::VecZnxBigBytesOf;
  9#[cfg(test)]
 10use poulpy_hal::api::{
 11    VecZnxAddInplace, VecZnxAddNormal, VecZnxFillUniform, VecZnxNormalize, VecZnxSub,
 12};
 13#[cfg(test)]
 14use poulpy_hal::source::Source;
 15use poulpy_hal::{
 16    api::{
 17        VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxDftApply,
 18        VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalizeTmpBytes,
 19    },
 20    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 21};
 22
 23use poulpy_hal::api::{SvpApplyDftToDftInplace, VecZnxNormalizeInplace, VecZnxSubInplace};
 24
 25use crate::tfhe::bdd_arithmetic::{FromBits, ToBits, UnsignedInteger};
 26
 27/// An FHE ciphertext encrypting the bits of an [UnsignedInteger].
 28pub struct FheUintBlocks<D: Data, T: UnsignedInteger> {
 29    pub(crate) blocks: Vec<GLWE<D>>,
 30    pub(crate) _base: u8,
 31    pub(crate) _phantom: PhantomData<T>,
 32}
 33
 34impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintBlocks<D, T> {
 35    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 36        self.blocks[0].base2k()
 37    }
 38
 39    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 40        self.blocks[0].k()
 41    }
 42
 43    fn n(&self) -> poulpy_core::layouts::Degree {
 44        self.blocks[0].n()
 45    }
 46}
 47
 48impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintBlocks<D, T> {
 49    fn rank(&self) -> poulpy_core::layouts::Rank {
 50        self.blocks[0].rank()
 51    }
 52}
 53
 54impl<T: UnsignedInteger> FheUintBlocks<Vec<u8>, T> {
 55    #[allow(dead_code)]
 56    pub(crate) fn alloc<A, BE: Backend>(module: &Module<BE>, infos: &A) -> Self
 57    where
 58        A: GLWEInfos,
 59    {
 60        Self::alloc_with(module, infos.base2k(), infos.k(), infos.rank())
 61    }
 62
 63    #[allow(dead_code)]
 64    pub(crate) fn alloc_with<BE: Backend>(module: &Module<BE>, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 65        Self {
 66            blocks: (0..T::WORD_SIZE)
 67                .map(|_| GLWE::alloc(module.n().into(), base2k, k, rank))
 68                .collect(),
 69            _base: 1,
 70            _phantom: PhantomData,
 71        }
 72    }
 73}
 74
 75impl<D: DataMut, T: UnsignedInteger + ToBits> FheUintBlocks<D, T> {
 76    #[allow(dead_code)]
 77    #[cfg(test)]
 78    pub(crate) fn encrypt_sk<S, BE: Backend>(
 79        &mut self,
 80        module: &Module<BE>,
 81        value: T,
 82        sk: &GLWESecretPrepared<S, BE>,
 83        source_xa: &mut Source,
 84        source_xe: &mut Source,
 85        scratch: &mut Scratch<BE>,
 86    ) where
 87        S: DataRef,
 88        Module<BE>: VecZnxDftBytesOf
 89            + VecZnxBigNormalize<BE>
 90            + VecZnxDftApply<BE>
 91            + SvpApplyDftToDftInplace<BE>
 92            + VecZnxIdftApplyConsume<BE>
 93            + VecZnxNormalizeTmpBytes
 94            + VecZnxFillUniform
 95            + VecZnxSubInplace
 96            + VecZnxAddInplace
 97            + VecZnxNormalizeInplace<BE>
 98            + VecZnxAddNormal
 99            + VecZnxNormalize<BE>
100            + VecZnxSub,
101        Scratch<BE>: ScratchTakeCore<BE>,
102    {
103        use poulpy_core::layouts::GLWEPlaintextLayout;
104
105        #[cfg(debug_assertions)]
106        {
107            assert!(module.n().is_multiple_of(T::WORD_SIZE));
108            assert_eq!(self.n(), module.n() as u32);
109            assert_eq!(sk.n(), module.n() as u32);
110        }
111
112        let pt_infos = GLWEPlaintextLayout {
113            n: self.n(),
114            base2k: self.base2k(),
115            k: 1_usize.into(),
116        };
117
118        let (mut pt, scratch_1) = scratch.take_glwe_pt(&pt_infos);
119
120        for i in 0..T::WORD_SIZE {
121            pt.encode_coeff_i64(value.bit(i) as i64, TorusPrecision(1), 0);
122            self.blocks[i].encrypt_sk(&module, &pt, sk, source_xa, source_xe, scratch_1);
123        }
124    }
125}
126
127impl<D: DataRef, T: UnsignedInteger + FromBits + ToBits> FheUintBlocks<D, T> {
128    pub fn decrypt<S: DataRef, BE: Backend>(
129        &self,
130        module: &Module<BE>,
131        sk: &GLWESecretPrepared<S, BE>,
132        scratch: &mut Scratch<BE>,
133    ) -> T
134    where
135        Module<BE>: VecZnxDftApply<BE>
136            + SvpApplyDftToDftInplace<BE>
137            + VecZnxIdftApplyConsume<BE>
138            + VecZnxBigAddInplace<BE>
139            + VecZnxBigAddSmallInplace<BE>
140            + VecZnxBigNormalize<BE>,
141        Scratch<BE>: ScratchTakeCore<BE>,
142    {
143        #[cfg(debug_assertions)]
144        {
145            assert!(module.n().is_multiple_of(T::WORD_SIZE));
146            assert_eq!(self.n(), module.n() as u32);
147            assert_eq!(sk.n(), module.n() as u32);
148        }
149
150        let pt_infos = GLWEPlaintextLayout {
151            n: self.n(),
152            base2k: self.base2k(),
153            k: self.k(),
154        };
155
156        let (mut pt, scratch_1) = scratch.take_glwe_pt(&pt_infos);
157
158        let mut bits: Vec<u8> = vec![0u8; T::WORD_SIZE];
159
160        let base2k: usize = self.base2k().into();
161        let scale: f64 = 4.0 / ((1 << base2k) as f64);
162
163        for (i, bit) in bits.iter_mut().enumerate().take(T::WORD_SIZE) {
164            self.blocks[i].decrypt(module, &mut pt, sk, scratch_1);
165            let value: i64 = pt.decode_coeff_i64(base2k.into(), 0);
166            *bit = ((value as f64) * scale).round() as u8;
167        }
168
169        T::from_bits(&bits)
170    }
171
172    pub fn noise<S: DataRef, BE: Backend>(
173        &self,
174        module: &Module<BE>,
175        sk: &GLWESecretPrepared<S, BE>,
176        want: T,
177        scratch: &mut Scratch<BE>,
178    ) -> Vec<f64>
179    where
180        Module<BE>: VecZnxDftBytesOf
181            + VecZnxBigBytesOf
182            + VecZnxDftApply<BE>
183            + SvpApplyDftToDftInplace<BE>
184            + VecZnxIdftApplyConsume<BE>
185            + VecZnxBigAddInplace<BE>
186            + VecZnxBigAddSmallInplace<BE>
187            + VecZnxBigNormalize<BE>
188            + VecZnxNormalizeTmpBytes
189            + VecZnxSubInplace
190            + VecZnxNormalizeInplace<BE>,
191        Scratch<BE>: ScratchTakeCore<BE>,
192    {
193        #[cfg(debug_assertions)]
194        {
195            assert!(module.n().is_multiple_of(T::WORD_SIZE));
196            assert_eq!(self.n(), module.n() as u32);
197            assert_eq!(sk.n(), module.n() as u32);
198        }
199
200        let pt_infos = GLWEPlaintextLayout {
201            n: self.n(),
202            base2k: self.base2k(),
203            k: 1_usize.into(),
204        };
205
206        let (mut pt_want, scratch_1) = scratch.take_glwe_pt(&pt_infos);
207
208        let mut noise: Vec<f64> = vec![0f64; T::WORD_SIZE];
209
210        for (i, noise_i) in noise.iter_mut().enumerate().take(T::WORD_SIZE) {
211            pt_want.encode_coeff_i64(want.bit(i) as i64, TorusPrecision(2), 0);
212            *noise_i = self.blocks[i].noise(module, sk, &pt_want, scratch_1);
213        }
214
215        noise
216    }
217}