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};
 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
 84        #[cfg(debug_assertions)]
 85        {
 86            assert!(module.n().is_multiple_of(T::WORD_SIZE));
 87            assert_eq!(self.n(), module.n() as u32);
 88            assert_eq!(sk.n(), module.n() as u32);
 89        }
 90
 91        let pt_infos = GLWEPlaintextLayout {
 92            n: self.n(),
 93            base2k: self.base2k(),
 94            k: 1_usize.into(),
 95        };
 96
 97        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
 98
 99        for i in 0..T::WORD_SIZE {
100            pt.encode_coeff_i64(value.bit(i) as i64, TorusPrecision(2), 0);
101            self.blocks[i].encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
102        }
103    }
104}
105
106impl<D: DataRef, T: UnsignedInteger + FromBits + ToBits> FheUintBlocks<D, T> {
107    pub fn decrypt<S, BE: Backend>(&self, module: &Module<BE>, sk: &S, scratch: &mut Scratch<BE>) -> T
108    where
109        Module<BE>: GLWEDecrypt<BE>,
110        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
111        Scratch<BE>: ScratchTakeCore<BE>,
112    {
113        #[cfg(debug_assertions)]
114        {
115            assert!(module.n().is_multiple_of(T::WORD_SIZE));
116            assert_eq!(self.n(), module.n() as u32);
117            assert_eq!(sk.n(), module.n() as u32);
118        }
119
120        let pt_infos = GLWEPlaintextLayout {
121            n: self.n(),
122            base2k: self.base2k(),
123            k: self.k(),
124        };
125
126        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
127
128        let mut bits: Vec<u8> = vec![0u8; T::WORD_SIZE];
129
130        let base2k: usize = self.base2k().into();
131        let scale: f64 = 4.0 / ((1 << base2k) as f64);
132
133        for (i, bit) in bits.iter_mut().enumerate().take(T::WORD_SIZE) {
134            self.blocks[i].decrypt(module, &mut pt, sk, scratch_1);
135            let value: i64 = pt.decode_coeff_i64(base2k.into(), 0);
136            *bit = ((value as f64) * scale).round() as u8;
137        }
138
139        T::from_bits(&bits)
140    }
141
142    pub fn noise<S, BE: Backend>(&self, module: &Module<BE>, sk: &S, want: T, scratch: &mut Scratch<BE>) -> Vec<f64>
143    where
144        Module<BE>: GLWENoise<BE>,
145        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
146        Scratch<BE>: ScratchTakeCore<BE>,
147    {
148        #[cfg(debug_assertions)]
149        {
150            assert!(module.n().is_multiple_of(T::WORD_SIZE));
151            assert_eq!(self.n(), module.n() as u32);
152            assert_eq!(sk.n(), module.n() as u32);
153        }
154
155        let pt_infos = GLWEPlaintextLayout {
156            n: self.n(),
157            base2k: self.base2k(),
158            k: 1_usize.into(),
159        };
160
161        let (mut pt_want, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
162
163        let mut noise: Vec<f64> = vec![0f64; T::WORD_SIZE];
164
165        for (i, noise_i) in noise.iter_mut().enumerate().take(T::WORD_SIZE) {
166            pt_want.encode_coeff_i64(want.bit(i) as i64, TorusPrecision(2), 0);
167            *noise_i = self.blocks[i].noise(module, sk, &pt_want, scratch_1);
168        }
169
170        noise
171    }
172}