1use itertools::Itertools;
  2use poulpy_core::{
  3    GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, LWEFromGLWE, ScratchTakeCore,
  4    layouts::{
  5        Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToRef,
  6        LWEInfos, LWEToMut, Rank, TorusPrecision, prepared::GLWEAutomorphismKeyPrepared,
  7    },
  8};
  9use poulpy_hal::{
 10    api::ModuleN,
 11    layouts::{Backend, Data, DataMut, DataRef, Scratch},
 12    source::Source,
 13};
 14use std::{collections::HashMap, marker::PhantomData};
 15
 16use crate::tfhe::bdd_arithmetic::{FromBits, ToBits, UnsignedInteger};
 17
 18/// An FHE ciphertext encrypting the bits of an [UnsignedInteger].
 19pub struct FheUint<D: Data, T: UnsignedInteger> {
 20    pub(crate) bits: GLWE<D>,
 21    pub(crate) _phantom: PhantomData<T>,
 22}
 23
 24impl<T: UnsignedInteger> FheUint<Vec<u8>, T> {
 25    pub fn alloc_from_infos<A>(infos: &A) -> Self
 26    where
 27        A: GLWEInfos,
 28    {
 29        Self::alloc(infos.n(), infos.base2k(), infos.k(), infos.rank())
 30    }
 31
 32    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 33        Self {
 34            bits: GLWE::alloc(n, base2k, k, rank),
 35            _phantom: PhantomData,
 36        }
 37    }
 38}
 39
 40impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUint<D, T> {
 41    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 42        self.bits.base2k()
 43    }
 44
 45    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 46        self.bits.k()
 47    }
 48
 49    fn n(&self) -> poulpy_core::layouts::Degree {
 50        self.bits.n()
 51    }
 52}
 53
 54impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUint<D, T> {
 55    fn rank(&self) -> poulpy_core::layouts::Rank {
 56        self.bits.rank()
 57    }
 58}
 59
 60impl<D: DataMut, T: UnsignedInteger + ToBits> FheUint<D, T> {
 61    pub fn encrypt_sk<S, M, BE: Backend>(
 62        &mut self,
 63        module: &M,
 64        data: T,
 65        sk: &S,
 66        source_xa: &mut Source,
 67        source_xe: &mut Source,
 68        scratch: &mut Scratch<BE>,
 69    ) where
 70        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 71        M: ModuleN + GLWEEncryptSk<BE>,
 72        Scratch<BE>: ScratchTakeCore<BE>,
 73    {
 74        #[cfg(debug_assertions)]
 75        {
 76            assert!(module.n().is_multiple_of(T::WORD_SIZE));
 77            assert_eq!(self.n(), module.n() as u32);
 78            assert_eq!(sk.n(), module.n() as u32);
 79        }
 80
 81        let gap: usize = module.n() / T::WORD_SIZE;
 82
 83        let mut data_bits: Vec<i64> = vec![0i64; module.n()];
 84
 85        for i in 0..T::WORD_SIZE {
 86            data_bits[i * 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: GLWEDecrypt<BE>,
108        Scratch<BE>: ScratchTakeCore<BE>,
109    {
110        #[cfg(debug_assertions)]
111        {
112            assert!(module.n().is_multiple_of(T::WORD_SIZE));
113            assert_eq!(self.n(), module.n() as u32);
114            assert_eq!(sk.n(), module.n() as u32);
115        }
116
117        let gap: usize = module.n() / T::WORD_SIZE;
118
119        let pt_infos = GLWEPlaintextLayout {
120            n: self.n(),
121            base2k: self.base2k(),
122            k: 1_usize.into(),
123        };
124
125        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
126
127        self.bits.decrypt(module, &mut pt, sk, scratch_1);
128
129        let mut data: Vec<i64> = vec![0i64; module.n()];
130
131        pt.decode_vec_i64(&mut data, TorusPrecision(2));
132
133        let bits: Vec<u8> = data.iter().step_by(gap).map(|c| *c as u8).collect_vec();
134        T::from_bits(&bits)
135    }
136}
137
138impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
139    #[allow(dead_code)]
140    pub(crate) fn pack<D1, ATK, M, BE: Backend>(
141        &mut self,
142        module: &M,
143        mut tmp_res: Vec<GLWE<D1>>,
144        auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<ATK, BE>>,
145        scratch: &mut Scratch<BE>,
146    ) where
147        D1: DataMut,
148        ATK: DataRef,
149        M: GLWEPacking<BE> + GLWECopy,
150        Scratch<BE>: ScratchTakeCore<BE>,
151    {
152        // Repacks the GLWE ciphertexts bits
153        let gap: usize = module.n() / T::WORD_SIZE;
154        let log_gap: usize = (usize::BITS - (gap - 1).leading_zeros()) as usize;
155        let mut cts: HashMap<usize, &mut GLWE<D1>> = HashMap::new();
156        for (i, ct) in tmp_res.iter_mut().enumerate().take(T::WORD_SIZE) {
157            cts.insert(i * gap, ct);
158        }
159
160        module.glwe_pack(&mut cts, log_gap, auto_keys, scratch);
161
162        // And copies the repacked ciphertext on the receiver.
163        module.glwe_copy(&mut self.bits, cts.remove(&0).unwrap());
164    }
165}
166
167impl<D: DataRef, T: UnsignedInteger> FheUint<D, T> {
168    pub fn get_bit<L, K, M, BE: Backend>(&self, module: &M, bit: usize, res: &mut L, ks: &K, scratch: &mut Scratch<BE>)
169    where
170        L: LWEToMut,
171        K: GGLWEPreparedToRef<BE> + GGLWEInfos,
172        M: LWEFromGLWE<BE> + GLWERotate<BE>,
173        Scratch<BE>: ScratchTakeCore<BE>,
174    {
175        let gap: usize = module.n() / T::WORD_SIZE;
176        res.to_mut().from_glwe(module, self, bit * gap, ks, scratch);
177    }
178}
179
180impl<D: DataRef, T: UnsignedInteger> GLWEToRef for FheUint<D, T> {
181    fn to_ref(&self) -> GLWE<&[u8]> {
182        self.bits.to_ref()
183    }
184}