1use itertools::Itertools;
  2use poulpy_core::{
  3    layouts::{
  4        prepared::{GLWEAutomorphismKeyPrepared, GLWESecretPrepared}, GLWEInfos, GLWEPlaintextLayout, LWEInfos, TorusPrecision, GLWE
  5    }, ScratchTakeCore,
  6};
  7use poulpy_hal::{
  8    api::{
  9        ScratchAvailable, SvpApplyDftToDftInplace, VecZnxAddInplace, VecZnxAddNormal,
 10        VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
 11        VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy,
 12        VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA,
 13        VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace,
 14        VecZnxRshInplace, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd,
 15        VmpApplyDftToDftTmpBytes,
 16    },
 17    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 18    source::Source,
 19};
 20use std::{collections::HashMap, marker::PhantomData};
 21
 22use crate::tfhe::bdd_arithmetic::{FromBits, ToBits, UnsignedInteger};
 23
 24/// A FHE ciphertext encrypting a [UnsignedInteger].
 25pub struct FheUintWord<D: Data, T: UnsignedInteger>(pub(crate) GLWE<D>, pub(crate) PhantomData<T>);
 26
 27impl<D: DataMut, T: UnsignedInteger> FheUintWord<D, T> {
 28    #[allow(dead_code)]
 29    fn post_process<ATK, BE: Backend>(
 30        &mut self,
 31        module: &Module<BE>,
 32        mut tmp_res: Vec<GLWE<&mut [u8]>>,
 33        auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<ATK, BE>>,
 34        scratch: &mut Scratch<BE>,
 35    ) where
 36        ATK: DataRef,
 37        Module<BE>: VecZnxSub
 38            + VecZnxCopy
 39            + VecZnxNegateInplace
 40            + VecZnxDftBytesOf
 41            + VecZnxAddInplace
 42            + VmpApplyDftToDftTmpBytes
 43            + VecZnxNormalizeTmpBytes
 44            + VecZnxDftApply<BE>
 45            + VmpApplyDftToDft<BE>
 46            + VmpApplyDftToDftAdd<BE>
 47            + VecZnxIdftApplyConsume<BE>
 48            + VecZnxBigNormalize<BE>
 49            + VecZnxNormalize<BE>
 50            + VecZnxRotateInplace<BE>
 51            + VecZnxNormalizeInplace<BE>
 52            + VecZnxSwitchRing
 53            + VecZnxBigAutomorphismInplace<BE>
 54            + VecZnxRshInplace<BE>
 55            + VecZnxDftCopy<BE>
 56            + VecZnxIdftApplyTmpA<BE>
 57            + VecZnxSubInplace
 58            + VecZnxBigNormalizeTmpBytes
 59            + VecZnxBigAddSmallInplace<BE>
 60            + VecZnxAutomorphismInplace<BE>
 61            + VecZnxBigSubSmallNegateInplace<BE>
 62            + VecZnxRotate,
 63        Scratch<BE>: ScratchTakeCore<BE>,
 64    {
 65        // Repacks the GLWE ciphertexts bits
 66        let gap: usize = module.n() / T::WORD_SIZE;
 67        let log_gap: usize = (usize::BITS - (gap - 1).leading_zeros()) as usize;
 68        let mut cts: HashMap<usize, &mut GLWE<&mut [u8]>> = HashMap::new();
 69        for (i, ct) in tmp_res.iter_mut().enumerate().take(T::WORD_SIZE) {
 70            cts.insert(i * gap, ct);
 71        }
 72        glwe_packing(module, &mut cts, log_gap, auto_keys, scratch);
 73
 74        // And copies the repacked ciphertext on the receiver.
 75        self.0.copy(module, cts.remove(&0).unwrap())
 76    }
 77}
 78
 79impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintWord<D, T> {
 80    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 81        self.0.base2k()
 82    }
 83
 84    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 85        self.0.k()
 86    }
 87
 88    fn n(&self) -> poulpy_core::layouts::Degree {
 89        self.0.n()
 90    }
 91}
 92
 93impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintWord<D, T> {
 94    fn rank(&self) -> poulpy_core::layouts::Rank {
 95        self.0.rank()
 96    }
 97}
 98
 99impl<D: DataMut, T: UnsignedInteger + ToBits> FheUintWord<D, T> {
100    pub fn encrypt_sk<S: DataRef, BE: Backend>(
101        &mut self,
102        module: &Module<BE>,
103        data: T,
104        sk: &GLWESecretPrepared<S, BE>,
105        source_xa: &mut Source,
106        source_xe: &mut Source,
107        scratch: &mut Scratch<BE>,
108    ) where
109        Module<BE>: VecZnxAddScalarInplace
110            + VecZnxDftBytesOf
111            + VecZnxBigNormalize<BE>
112            + VecZnxDftApply<BE>
113            + SvpApplyDftToDftInplace<BE>
114            + VecZnxIdftApplyConsume<BE>
115            + VecZnxNormalizeTmpBytes
116            + VecZnxFillUniform
117            + VecZnxSubInplace
118            + VecZnxAddInplace
119            + VecZnxNormalizeInplace<BE>
120            + VecZnxAddNormal
121            + VecZnxNormalize<BE>
122            + VecZnxSub,
123        Scratch<BE>: ScratchTakeCore<BE>,
124    {
125        #[cfg(debug_assertions)]
126        {
127            assert!(module.n().is_multiple_of(T::WORD_SIZE));
128            assert_eq!(self.n(), module.n() as u32);
129            assert_eq!(sk.n(), module.n() as u32);
130        }
131
132        let gap: usize = module.n() / T::WORD_SIZE;
133
134        let mut data_bits: Vec<i64> = vec![0i64; module.n()];
135
136        for i in 0..T::WORD_SIZE {
137            data_bits[i * gap] = data.bit(i) as i64
138        }
139
140        let pt_infos = GLWEPlaintextLayout {
141            n: self.n(),
142            base2k: self.base2k(),
143            k: 1_usize.into(),
144        };
145
146        let (mut pt, scratch_1) = scratch.take_glwe_pt(&pt_infos);
147
148        pt.encode_vec_i64(&data_bits, TorusPrecision(1));
149        self.0
150            .encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
151    }
152}
153
154impl<D: DataRef, T: UnsignedInteger + FromBits> FheUintWord<D, T> {
155    pub fn decrypt<S: DataRef, BE: Backend>(
156        &self,
157        module: &Module<BE>,
158        sk: &GLWESecretPrepared<S, BE>,
159        scratch: &mut Scratch<BE>,
160    ) -> T
161    where
162        Module<BE>: VecZnxDftApply<BE>
163            + SvpApplyDftToDftInplace<BE>
164            + VecZnxIdftApplyConsume<BE>
165            + VecZnxBigAddInplace<BE>
166            + VecZnxBigAddSmallInplace<BE>
167            + VecZnxBigNormalize<BE>,
168        Scratch<BE>: ScratchTakeCore<BE>,
169    {
170        #[cfg(debug_assertions)]
171        {
172            assert!(module.n().is_multiple_of(T::WORD_SIZE));
173            assert_eq!(self.n(), module.n() as u32);
174            assert_eq!(sk.n(), module.n() as u32);
175        }
176
177        let gap: usize = module.n() / T::WORD_SIZE;
178
179        let pt_infos = GLWEPlaintextLayout {
180            n: self.n(),
181            base2k: self.base2k(),
182            k: 1_usize.into(),
183        };
184
185        let (mut pt, scratch_1) = scratch.take_glwe_pt(&pt_infos);
186
187        self.0.decrypt(module, &mut pt, sk, scratch_1);
188
189        let mut data: Vec<i64> = vec![0i64; module.n()];
190
191        pt.decode_vec_i64(&mut data, TorusPrecision(1));
192
193        let bits: Vec<u8> = data.iter().step_by(gap).map(|c| *c as u8).collect_vec();
194        T::from_bits(&bits)
195    }
196}