1use poulpy_core::{
  2    GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, LWEFromGLWE, ScratchTakeCore,
  3    layouts::{
  4        Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToRef,
  5        LWEInfos, LWEToMut, Rank, TorusPrecision, prepared::GLWEAutomorphismKeyPrepared,
  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<D: DataRef, T: UnsignedInteger> LWEInfos for FheUint<D, T> {
 40    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 41        self.bits.base2k()
 42    }
 43
 44    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 45        self.bits.k()
 46    }
 47
 48    fn n(&self) -> poulpy_core::layouts::Degree {
 49        self.bits.n()
 50    }
 51}
 52
 53impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUint<D, T> {
 54    fn rank(&self) -> poulpy_core::layouts::Rank {
 55        self.bits.rank()
 56    }
 57}
 58
 59impl<D: DataMut, T: UnsignedInteger + ToBits> FheUint<D, T> {
 60    pub fn encrypt_sk<S, M, BE: Backend>(
 61        &mut self,
 62        module: &M,
 63        data: T,
 64        sk: &S,
 65        source_xa: &mut Source,
 66        source_xe: &mut Source,
 67        scratch: &mut Scratch<BE>,
 68    ) where
 69        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 70        M: ModuleLogN + GLWEEncryptSk<BE>,
 71        Scratch<BE>: ScratchTakeCore<BE>,
 72    {
 73        #[cfg(debug_assertions)]
 74        {
 75            assert!(module.n().is_multiple_of(T::BITS as usize));
 76            assert_eq!(self.n(), module.n() as u32);
 77            assert_eq!(sk.n(), module.n() as u32);
 78        }
 79
 80        let mut data_bits: Vec<i64> = vec![0i64; module.n()];
 81
 82        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
 83
 84        // Interleaves bytes
 85        for i in 0..T::BITS as usize {
 86            data_bits[T::bit_index(i) << log_gap] = data.bit(i) as i64
 87        }
 88
 89        let pt_infos = GLWEPlaintextLayout {
 90            n: self.n(),
 91            base2k: self.base2k(),
 92            k: 1_usize.into(),
 93        };
 94
 95        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
 96
 97        pt.encode_vec_i64(&data_bits, TorusPrecision(2));
 98        self.bits
 99            .encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
100    }
101}
102
103impl<D: DataRef, T: UnsignedInteger + FromBits> FheUint<D, T> {
104    pub fn decrypt<S, M, BE: Backend>(&self, module: &M, sk: &S, scratch: &mut Scratch<BE>) -> T
105    where
106        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
107        M: ModuleLogN + GLWEDecrypt<BE>,
108        Scratch<BE>: ScratchTakeCore<BE>,
109    {
110        #[cfg(debug_assertions)]
111        {
112            assert!(module.n().is_multiple_of(T::BITS as usize));
113            assert_eq!(self.n(), module.n() as u32);
114            assert_eq!(sk.n(), module.n() as u32);
115        }
116
117        let pt_infos = GLWEPlaintextLayout {
118            n: self.n(),
119            base2k: self.base2k(),
120            k: 1_usize.into(),
121        };
122
123        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
124
125        self.bits.decrypt(module, &mut pt, sk, scratch_1);
126
127        let mut data_bits: Vec<i64> = vec![0i64; module.n()];
128        pt.decode_vec_i64(&mut data_bits, TorusPrecision(2));
129
130        let mut bits: Vec<u8> = vec![0u8; T::BITS as usize];
131
132        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
133
134        // Retrives from interleaved bytes
135        for i in 0..T::BITS as usize {
136            bits[i] = data_bits[T::bit_index(i) << log_gap] as u8
137        }
138
139        T::from_bits(&bits)
140    }
141}
142
143impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
144    #[allow(dead_code)]
145    pub(crate) fn pack<D1, ATK, M, BE: Backend>(
146        &mut self,
147        module: &M,
148        mut tmp_res: Vec<GLWE<D1>>,
149        auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<ATK, BE>>,
150        scratch: &mut Scratch<BE>,
151    ) where
152        D1: DataMut,
153        ATK: DataRef,
154        M: ModuleLogN + GLWEPacking<BE> + GLWECopy,
155        Scratch<BE>: ScratchTakeCore<BE>,
156    {
157        // Repacks the GLWE ciphertexts bits
158        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
159        let mut cts: HashMap<usize, &mut GLWE<D1>> = HashMap::new();
160        for (i, ct) in tmp_res.iter_mut().enumerate().take(T::BITS as usize) {
161            cts.insert(T::bit_index(i) << log_gap, ct);
162        }
163
164        module.glwe_pack(&mut cts, log_gap, auto_keys, scratch);
165
166        // And copies the repacked ciphertext on the receiver.
167        module.glwe_copy(&mut self.bits, cts.remove(&0).unwrap());
168    }
169}
170
171impl<D: DataRef, T: UnsignedInteger> FheUint<D, T> {
172    pub fn get_bit<L, K, M, BE: Backend>(&self, module: &M, bit: usize, res: &mut L, ks: &K, scratch: &mut Scratch<BE>)
173    where
174        L: LWEToMut,
175        K: GGLWEPreparedToRef<BE> + GGLWEInfos,
176        M: ModuleLogN + LWEFromGLWE<BE> + GLWERotate<BE>,
177        Scratch<BE>: ScratchTakeCore<BE>,
178    {
179        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
180        res.to_mut()
181            .from_glwe(module, self, T::bit_index(bit) << log_gap, ks, scratch);
182    }
183}
184
185impl<D: DataRef, T: UnsignedInteger> GLWEToRef for FheUint<D, T> {
186    fn to_ref(&self) -> GLWE<&[u8]> {
187        self.bits.to_ref()
188    }
189}