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