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