1use itertools::Itertools;
  2use poulpy_core::{
  3    GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, ScratchTakeCore,
  4    layouts::{
  5        GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, LWEInfos, TorusPrecision,
  6        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] in compressed format (ideal for decryption/serialization).
 19pub struct FheUintCompressed<D: Data, T: UnsignedInteger>(pub(crate) GLWE<D>, pub(crate) PhantomData<T>);
 20
 21impl<D: DataMut, T: UnsignedInteger> FheUintCompressed<D, T> {
 22    #[allow(dead_code)]
 23    fn post_process<ATK, M, BE: Backend>(
 24        &mut self,
 25        module: &M,
 26        mut tmp_res: Vec<GLWE<&mut [u8]>>,
 27        auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<ATK, BE>>,
 28        scratch: &mut Scratch<BE>,
 29    ) where
 30        ATK: DataRef,
 31        M: GLWEPacking<BE> + GLWECopy,
 32        Scratch<BE>: ScratchTakeCore<BE>,
 33    {
 34        // Repacks the GLWE ciphertexts bits
 35        let gap: usize = module.n() / T::WORD_SIZE;
 36        let log_gap: usize = (usize::BITS - (gap - 1).leading_zeros()) as usize;
 37        let mut cts: HashMap<usize, &mut GLWE<&mut [u8]>> = HashMap::new();
 38        for (i, ct) in tmp_res.iter_mut().enumerate().take(T::WORD_SIZE) {
 39            cts.insert(i * gap, ct);
 40        }
 41
 42        module.glwe_pack(&mut cts, log_gap, auto_keys, scratch);
 43
 44        // And copies the repacked ciphertext on the receiver.
 45        module.glwe_copy(&mut self.0, cts.remove(&0).unwrap());
 46    }
 47}
 48
 49impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintCompressed<D, T> {
 50    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 51        self.0.base2k()
 52    }
 53
 54    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 55        self.0.k()
 56    }
 57
 58    fn n(&self) -> poulpy_core::layouts::Degree {
 59        self.0.n()
 60    }
 61}
 62
 63impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintCompressed<D, T> {
 64    fn rank(&self) -> poulpy_core::layouts::Rank {
 65        self.0.rank()
 66    }
 67}
 68
 69impl<D: DataMut, T: UnsignedInteger + ToBits> FheUintCompressed<D, T> {
 70    pub fn encrypt_sk<S, M, BE: Backend>(
 71        &mut self,
 72        module: &M,
 73        data: T,
 74        sk: &S,
 75        source_xa: &mut Source,
 76        source_xe: &mut Source,
 77        scratch: &mut Scratch<BE>,
 78    ) where
 79        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 80        M: ModuleN + GLWEEncryptSk<BE>,
 81        Scratch<BE>: ScratchTakeCore<BE>,
 82    {
 83        #[cfg(debug_assertions)]
 84        {
 85            assert!(module.n().is_multiple_of(T::WORD_SIZE));
 86            assert_eq!(self.n(), module.n() as u32);
 87            assert_eq!(sk.n(), module.n() as u32);
 88        }
 89
 90        let gap: usize = module.n() / T::WORD_SIZE;
 91
 92        let mut data_bits: Vec<i64> = vec![0i64; module.n()];
 93
 94        for i in 0..T::WORD_SIZE {
 95            data_bits[i * gap] = data.bit(i) as i64
 96        }
 97
 98        let pt_infos = GLWEPlaintextLayout {
 99            n: self.n(),
100            base2k: self.base2k(),
101            k: 1_usize.into(),
102        };
103
104        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
105
106        pt.encode_vec_i64(&data_bits, TorusPrecision(1));
107        self.0
108            .encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
109    }
110}
111
112impl<D: DataRef, T: UnsignedInteger + FromBits> FheUintCompressed<D, T> {
113    pub fn decrypt<S, M, BE: Backend>(&self, module: &M, sk: &S, scratch: &mut Scratch<BE>) -> T
114    where
115        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
116        M: GLWEDecrypt<BE>,
117        Scratch<BE>: ScratchTakeCore<BE>,
118    {
119        #[cfg(debug_assertions)]
120        {
121            assert!(module.n().is_multiple_of(T::WORD_SIZE));
122            assert_eq!(self.n(), module.n() as u32);
123            assert_eq!(sk.n(), module.n() as u32);
124        }
125
126        let gap: usize = module.n() / T::WORD_SIZE;
127
128        let pt_infos = GLWEPlaintextLayout {
129            n: self.n(),
130            base2k: self.base2k(),
131            k: 1_usize.into(),
132        };
133
134        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
135
136        self.0.decrypt(module, &mut pt, sk, scratch_1);
137
138        let mut data: Vec<i64> = vec![0i64; module.n()];
139
140        pt.decode_vec_i64(&mut data, TorusPrecision(1));
141
142        let bits: Vec<u8> = data.iter().step_by(gap).map(|c| *c as u8).collect_vec();
143        T::from_bits(&bits)
144    }
145}