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