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