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