1use poulpy_core::{
  2    GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, 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    // pub fn copy_byte<D0, D1, BRA, M, BE: Backend>(
174    // &mut self,
175    // module: &M,
176    // byte_self: usize,
177    // byte_a: usize,
178    // a: &FheUint<D1, T>,
179    // keys: &BDDKeyPrepared<D0, BRA, BE>,
180    // scratch: &mut Scratch<BE>,
181    // ) where
182    // D0: DataRef,
183    // D1: DataRef,
184    // BRA: BlindRotationAlgo,
185    // M:ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd,
186    // Scratch<BE>: ScratchTakeBDD<T, BE>,
187    // {
188    // let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(a);
189    //
190    //
191    // let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
192    // module.glwe_rotate(-((T::bit_index(byte_a << 3) << log_gap) as i64), tmp_fhe_uint_byte, self);
193    // module.glwe_trace_inplace(&mut tmp_fhe_uint_byte, module.log_n() - 3, module.log_n(),&keys.cbt.atk, scratch);
194    //
195    // let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
196    // let rot: i64 = (T::bit_index(byte_self << 3) << log_gap) as i64;
197    //
198    // Move starting byte index to first coefficient
199    // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch);
200    //
201    // Stores this byte (everything else zeroed) into tmp_trace
202    // let (mut tmp_trace, scratch_1) = scratch.take_glwe(self);
203    // module.glwe_trace(
204    // &mut tmp_trace,
205    // module.log_n() - 3,
206    // module.log_n(),
207    // self,
208    // &keys.cbt.atk,
209    // scratch_1,
210    // );
211    //
212    // Subtracts the byte
213    // module.glwe_sub_inplace(&mut self.bits, &tmp_trace);
214    //
215    // module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte);
216    //
217    // Moves back into the original position
218    // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch);
219    //
220    // }
221}
222
223impl<D: DataMut, T: UnsignedInteger> GLWEToMut for FheUint<D, T> {
224    fn to_mut(&mut self) -> GLWE<&mut [u8]> {
225        self.bits.to_mut()
226    }
227}
228
229pub trait ScratchTakeBDD<T: UnsignedInteger, BE: Backend>
230where
231    Self: ScratchTakeCore<BE>,
232{
233    fn take_fhe_uint<A>(&mut self, infos: &A) -> (FheUint<&mut [u8], T>, &mut Self)
234    where
235        A: GLWEInfos,
236    {
237        let (glwe, scratch) = self.take_glwe(infos);
238        (
239            FheUint {
240                bits: glwe,
241                _phantom: PhantomData,
242            },
243            scratch,
244        )
245    }
246}
247
248impl<T: UnsignedInteger, BE: Backend> ScratchTakeBDD<T, BE> for Scratch<BE> where Self: ScratchTakeCore<BE> {}
249
250impl<D: DataRef, T: UnsignedInteger> FheUint<D, T> {
251    pub fn get_bit<R, K, M, BE: Backend>(&self, module: &M, bit: usize, res: &mut R, ks: &K, scratch: &mut Scratch<BE>)
252    where
253        R: LWEToMut,
254        K: GGLWEPreparedToRef<BE> + GGLWEInfos,
255        M: ModuleLogN + LWEFromGLWE<BE> + GLWERotate<BE>,
256        Scratch<BE>: ScratchTakeCore<BE>,
257    {
258        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
259        res.to_mut()
260            .from_glwe(module, self, T::bit_index(bit) << log_gap, ks, scratch);
261    }
262}
263
264impl<D: DataRef, T: UnsignedInteger> GLWEToRef for FheUint<D, T> {
265    fn to_ref(&self) -> GLWE<&[u8]> {
266        self.bits.to_ref()
267    }
268}