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