1use std::marker::PhantomData;
  2
  3use poulpy_core::{
  4    GLWEDecrypt, GLWENoise,
  5    layouts::{Base2K, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, LWEInfos, Rank, TorusPrecision},
  6};
  7
  8use poulpy_core::GLWEEncryptSk;
  9use poulpy_core::ScratchTakeCore;
 10use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, ZnxZero};
 11use poulpy_hal::source::Source;
 12
 13use crate::tfhe::bdd_arithmetic::{FromBits, ToBits, UnsignedInteger};
 14
 15/// An FHE ciphertext encrypting the bits of an [UnsignedInteger].
 16pub struct FheUintBlocks<D: Data, T: UnsignedInteger> {
 17    pub(crate) blocks: Vec<GLWE<D>>,
 18    pub(crate) _base: u8,
 19    pub(crate) _phantom: PhantomData<T>,
 20}
 21
 22impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintBlocks<D, T> {
 23    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 24        self.blocks[0].base2k()
 25    }
 26
 27    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 28        self.blocks[0].k()
 29    }
 30
 31    fn n(&self) -> poulpy_core::layouts::Degree {
 32        self.blocks[0].n()
 33    }
 34}
 35
 36impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintBlocks<D, T> {
 37    fn rank(&self) -> poulpy_core::layouts::Rank {
 38        self.blocks[0].rank()
 39    }
 40}
 41
 42impl<T: UnsignedInteger> FheUintBlocks<Vec<u8>, T> {
 43    pub fn alloc_from_infos<A, BE: Backend>(module: &Module<BE>, infos: &A) -> Self
 44    where
 45        A: GLWEInfos,
 46    {
 47        Self::alloc(module, infos.base2k(), infos.k(), infos.rank())
 48    }
 49
 50    pub fn alloc<BE: Backend>(module: &Module<BE>, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 51        Self {
 52            blocks: (0..T::WORD_SIZE)
 53                .map(|_| GLWE::alloc(module.n().into(), base2k, k, rank))
 54                .collect(),
 55            _base: 1,
 56            _phantom: PhantomData,
 57        }
 58    }
 59}
 60
 61impl<D: DataMut, T: UnsignedInteger + ToBits> FheUintBlocks<D, T> {
 62    pub fn encrypt_sk<S, BE: Backend>(
 63        &mut self,
 64        module: &Module<BE>,
 65        value: T,
 66        sk: &S,
 67        source_xa: &mut Source,
 68        source_xe: &mut Source,
 69        scratch: &mut Scratch<BE>,
 70    ) where
 71        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 72        Module<BE>: GLWEEncryptSk<BE>,
 73        Scratch<BE>: ScratchTakeCore<BE>,
 74    {
 75        use poulpy_core::layouts::GLWEPlaintextLayout;
 76        use poulpy_hal::layouts::ZnxZero;
 77
 78        #[cfg(debug_assertions)]
 79        {
 80            assert!(module.n().is_multiple_of(T::WORD_SIZE));
 81            assert_eq!(self.n(), module.n() as u32);
 82            assert_eq!(sk.n(), module.n() as u32);
 83        }
 84
 85        let pt_infos = GLWEPlaintextLayout {
 86            n: self.n(),
 87            base2k: self.base2k(),
 88            k: 1_usize.into(),
 89        };
 90
 91        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
 92        pt.data.zero();
 93
 94        for i in 0..T::WORD_SIZE {
 95            pt.encode_coeff_i64(value.bit(i) as i64, TorusPrecision(2), 0);
 96            self.blocks[i].encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
 97        }
 98    }
 99}
100
101impl<D: DataRef, T: UnsignedInteger + FromBits + ToBits> FheUintBlocks<D, T> {
102    pub fn decrypt<S, BE: Backend>(&self, module: &Module<BE>, sk: &S, scratch: &mut Scratch<BE>) -> T
103    where
104        Module<BE>: GLWEDecrypt<BE>,
105        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
106        Scratch<BE>: ScratchTakeCore<BE>,
107    {
108        #[cfg(debug_assertions)]
109        {
110            assert!(module.n().is_multiple_of(T::WORD_SIZE));
111            assert_eq!(self.n(), module.n() as u32);
112            assert_eq!(sk.n(), module.n() as u32);
113        }
114
115        let pt_infos = GLWEPlaintextLayout {
116            n: self.n(),
117            base2k: self.base2k(),
118            k: self.k(),
119        };
120
121        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
122
123        let mut bits: Vec<u8> = vec![0u8; T::WORD_SIZE];
124
125        let base2k: usize = self.base2k().into();
126        let scale: f64 = 4.0 / ((1 << base2k) as f64);
127
128        for (i, bit) in bits.iter_mut().enumerate().take(T::WORD_SIZE) {
129            self.blocks[i].decrypt(module, &mut pt, sk, scratch_1);
130            let value: i64 = pt.decode_coeff_i64(base2k.into(), 0);
131            *bit = ((value as f64) * scale).round() as u8;
132        }
133
134        T::from_bits(&bits)
135    }
136
137    pub fn noise<S, BE: Backend>(&self, module: &Module<BE>, sk: &S, want: T, scratch: &mut Scratch<BE>) -> Vec<f64>
138    where
139        Module<BE>: GLWENoise<BE>,
140        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
141        Scratch<BE>: ScratchTakeCore<BE>,
142    {
143        #[cfg(debug_assertions)]
144        {
145            assert!(module.n().is_multiple_of(T::WORD_SIZE));
146            assert_eq!(self.n(), module.n() as u32);
147            assert_eq!(sk.n(), module.n() as u32);
148        }
149
150        let pt_infos = GLWEPlaintextLayout {
151            n: self.n(),
152            base2k: self.base2k(),
153            k: 1_usize.into(),
154        };
155
156        let (mut pt_want, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
157        pt_want.data.zero();
158
159        let mut noise: Vec<f64> = vec![0f64; T::WORD_SIZE];
160
161        for (i, noise_i) in noise.iter_mut().enumerate().take(T::WORD_SIZE) {
162            pt_want.encode_coeff_i64(want.bit(i) as i64, TorusPrecision(2), 0);
163            *noise_i = self.blocks[i].noise(module, sk, &pt_want, scratch_1);
164        }
165
166        noise
167    }
168}