1use poulpy_core::{
  2    GLWEAdd, GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, GLWESub, GLWETrace, LWEFromGLWE, ScratchTakeCore,
  3    layouts::{
  4        Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToMut,
  5        GLWEToRef, LWEInfos, LWEToMut, Rank, TorusPrecision,
  6    },
  7};
  8use poulpy_hal::{
  9    api::ModuleLogN,
 10    layouts::{Backend, Data, DataMut, DataRef, Scratch},
 11    source::Source,
 12};
 13use std::{collections::HashMap, marker::PhantomData};
 14
 15use crate::tfhe::{
 16    bdd_arithmetic::{BDDKeyPrepared, FromBits, ToBits, UnsignedInteger},
 17    blind_rotation::BlindRotationAlgo,
 18};
 19
 20/// An FHE ciphertext encrypting the bits of an [UnsignedInteger].
 21pub struct FheUint<D: Data, T: UnsignedInteger> {
 22    pub(crate) bits: GLWE<D>,
 23    pub(crate) _phantom: PhantomData<T>,
 24}
 25
 26impl<T: UnsignedInteger> FheUint<Vec<u8>, T> {
 27    pub fn alloc_from_infos<A>(infos: &A) -> Self
 28    where
 29        A: GLWEInfos,
 30    {
 31        Self::alloc(infos.n(), infos.base2k(), infos.k(), infos.rank())
 32    }
 33
 34    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 35        Self {
 36            bits: GLWE::alloc(n, base2k, k, rank),
 37            _phantom: PhantomData,
 38        }
 39    }
 40}
 41
 42impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUint<D, T> {
 43    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 44        self.bits.base2k()
 45    }
 46
 47    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 48        self.bits.k()
 49    }
 50
 51    fn n(&self) -> poulpy_core::layouts::Degree {
 52        self.bits.n()
 53    }
 54}
 55
 56impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUint<D, T> {
 57    fn rank(&self) -> poulpy_core::layouts::Rank {
 58        self.bits.rank()
 59    }
 60}
 61
 62impl<D: DataMut, T: UnsignedInteger + ToBits> FheUint<D, T> {
 63    pub fn encrypt_sk<S, M, BE: Backend>(
 64        &mut self,
 65        module: &M,
 66        data: T,
 67        sk: &S,
 68        source_xa: &mut Source,
 69        source_xe: &mut Source,
 70        scratch: &mut Scratch<BE>,
 71    ) where
 72        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 73        M: ModuleLogN + GLWEEncryptSk<BE>,
 74        Scratch<BE>: ScratchTakeCore<BE>,
 75    {
 76        #[cfg(debug_assertions)]
 77        {
 78            assert!(module.n().is_multiple_of(T::BITS as usize));
 79            assert_eq!(self.n(), module.n() as u32);
 80            assert_eq!(sk.n(), module.n() as u32);
 81        }
 82
 83        let mut data_bits: Vec<i64> = vec![0i64; module.n()];
 84
 85        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
 86
 87        // Interleaves bytes
 88        for i in 0..T::BITS as usize {
 89            data_bits[T::bit_index(i) << log_gap] = data.bit(i) as i64
 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
100        pt.encode_vec_i64(&data_bits, TorusPrecision(2));
101        self.bits
102            .encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
103    }
104}
105
106impl<D: DataRef, T: UnsignedInteger + FromBits> FheUint<D, T> {
107    pub fn decrypt<S, M, BE: Backend>(&self, module: &M, sk: &S, scratch: &mut Scratch<BE>) -> T
108    where
109        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
110        M: ModuleLogN + GLWEDecrypt<BE>,
111        Scratch<BE>: ScratchTakeCore<BE>,
112    {
113        #[cfg(debug_assertions)]
114        {
115            assert!(module.n().is_multiple_of(T::BITS as usize));
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: 1_usize.into(),
124        };
125
126        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
127
128        self.bits.decrypt(module, &mut pt, sk, scratch_1);
129
130        let mut data_bits: Vec<i64> = vec![0i64; module.n()];
131        pt.decode_vec_i64(&mut data_bits, TorusPrecision(2));
132
133        let mut bits: Vec<u8> = vec![0u8; T::BITS as usize];
134
135        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
136
137        // Retrives from interleaved bytes
138        for i in 0..T::BITS as usize {
139            bits[i] = data_bits[T::bit_index(i) << log_gap] as u8
140        }
141
142        T::from_bits(&bits)
143    }
144}
145
146impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
147    /// Packs Vec<GLWE(bit[i])> into [FheUint].
148    pub fn pack<G, D1, M, BRA: BlindRotationAlgo, BE: Backend>(
149        &mut self,
150        module: &M,
151        mut bits: Vec<G>,
152        key: &BDDKeyPrepared<D1, BRA, BE>,
153        scratch: &mut Scratch<BE>,
154    ) where
155        G: GLWEToMut + GLWEToRef + GLWEInfos,
156        D1: DataRef,
157        M: ModuleLogN + GLWEPacking<BE> + GLWECopy,
158        Scratch<BE>: ScratchTakeCore<BE>,
159    {
160        // Repacks the GLWE ciphertexts bits
161        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
162        let mut cts: HashMap<usize, &mut G> = HashMap::new();
163        for (i, ct) in bits.iter_mut().enumerate().take(T::BITS as usize) {
164            cts.insert(T::bit_index(i) << log_gap, ct);
165        }
166
167        module.glwe_pack(&mut cts, log_gap, &key.cbt.atk, scratch);
168
169        // And copies the repacked ciphertext on the receiver.
170        module.glwe_copy(&mut self.bits, cts.remove(&0).unwrap());
171    }
172
173    #[allow(clippy::too_many_arguments)]
174    pub fn splice_u16<D0, A, B, BRA, M, BE: Backend>(
175        &mut self,
176        module: &M,
177        dst: usize,
178        src: usize,
179        a: &A,
180        b: &B,
181        keys: &BDDKeyPrepared<D0, BRA, BE>,
182        scratch: &mut Scratch<BE>,
183    ) where
184        D0: DataRef,
185        A: GLWEToRef + GLWEInfos,
186        B: GLWEToRef + GLWEInfos,
187        BRA: BlindRotationAlgo,
188        M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
189        Scratch<BE>: ScratchTakeBDD<T, BE>,
190    {
191        assert!(dst < (T::BITS >> 4) as usize);
192        assert!(src < (T::BITS >> 4) as usize);
193
194        let (mut tmp, scratch_1) = scratch.take_fhe_uint(self);
195        tmp.splice_u8(module, dst << 1, src << 1, a, b, keys, scratch_1);
196        self.splice_u8(
197            module,
198            (dst << 1) + 1,
199            (src << 1) + 1,
200            &tmp,
201            b,
202            keys,
203            scratch_1,
204        );
205    }
206
207    #[allow(clippy::too_many_arguments)]
208    // Store on the receiver a where the byte_a-th byte of a has been replaced by byte_src2 of src2.
209    pub fn splice_u8<D0, A, B, BRA, M, BE: Backend>(
210        &mut self,
211        module: &M,
212        dst: usize,
213        src: usize,
214        a: &A,
215        b: &B,
216        keys: &BDDKeyPrepared<D0, BRA, BE>,
217        scratch: &mut Scratch<BE>,
218    ) where
219        D0: DataRef,
220        A: GLWEToRef + GLWEInfos,
221        B: GLWEToRef + GLWEInfos,
222        BRA: BlindRotationAlgo,
223        M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
224        Scratch<BE>: ScratchTakeBDD<T, BE>,
225    {
226        assert!(dst < (T::BITS >> 3) as usize);
227        assert!(src < (T::BITS >> 3) as usize);
228
229        // 1) Zero the byte receiver
230        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
231        let trace_start = (T::LOG_BITS - T::LOG_BYTES) as usize;
232        let rot: i64 = (T::bit_index(dst << 3) << log_gap) as i64;
233
234        // Move a to self and align byte
235        module.glwe_rotate(-rot, &mut self.bits, a);
236
237        // Stores this byte (everything else zeroed) into tmp_trace
238        let (mut tmp_trace, scratch_1) = scratch.take_glwe(a);
239        module.glwe_trace(
240            &mut tmp_trace,
241            trace_start,
242            module.log_n(),
243            self,
244            &keys.cbt.atk,
245            scratch_1,
246        );
247
248        // Subtracts to self to zero it
249        module.glwe_sub_inplace(&mut self.bits, &tmp_trace);
250
251        // Isolate the byte to transfer from a
252        let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(b);
253
254        // Move a[byte_a] into a[0]
255        module.glwe_rotate(
256            -((T::bit_index(src << 3) << log_gap) as i64),
257            &mut tmp_fhe_uint_byte,
258            b,
259        );
260
261        // Zeroes all other bytes
262        module.glwe_trace_inplace(
263            &mut tmp_fhe_uint_byte,
264            trace_start,
265            module.log_n(),
266            &keys.cbt.atk,
267            scratch_1,
268        );
269
270        // Add self[0] += a[0]
271        module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte);
272
273        // Moves back self[0] to self[byte_tg]
274        module.glwe_rotate_inplace(rot, &mut self.bits, scratch);
275    }
276}
277
278impl<D: DataMut, T: UnsignedInteger> GLWEToMut for FheUint<D, T> {
279    fn to_mut(&mut self) -> GLWE<&mut [u8]> {
280        self.bits.to_mut()
281    }
282}
283
284pub trait ScratchTakeBDD<T: UnsignedInteger, BE: Backend>
285where
286    Self: ScratchTakeCore<BE>,
287{
288    fn take_fhe_uint<A>(&mut self, infos: &A) -> (FheUint<&mut [u8], T>, &mut Self)
289    where
290        A: GLWEInfos,
291    {
292        let (glwe, scratch) = self.take_glwe(infos);
293        (
294            FheUint {
295                bits: glwe,
296                _phantom: PhantomData,
297            },
298            scratch,
299        )
300    }
301}
302
303impl<T: UnsignedInteger, BE: Backend> ScratchTakeBDD<T, BE> for Scratch<BE> where Self: ScratchTakeCore<BE> {}
304
305impl<D: DataRef, T: UnsignedInteger> FheUint<D, T> {
306    pub fn get_bit<R, K, M, BE: Backend>(&self, module: &M, bit: usize, res: &mut R, ks: &K, scratch: &mut Scratch<BE>)
307    where
308        R: LWEToMut,
309        K: GGLWEPreparedToRef<BE> + GGLWEInfos,
310        M: ModuleLogN + LWEFromGLWE<BE> + GLWERotate<BE>,
311        Scratch<BE>: ScratchTakeCore<BE>,
312    {
313        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
314        res.to_mut()
315            .from_glwe(module, self, T::bit_index(bit) << log_gap, ks, scratch);
316    }
317}
318
319impl<D: DataRef, T: UnsignedInteger> GLWEToRef for FheUint<D, T> {
320    fn to_ref(&self) -> GLWE<&[u8]> {
321        self.bits.to_ref()
322    }
323}