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<D: Data, T: UnsignedInteger> FheUintBlocks<D, T> {
 43    pub fn new(blocks: Vec<GLWE<D>>) -> Self {
 44        assert_eq!(blocks.len(), T::WORD_SIZE);
 45        Self {
 46            blocks,
 47            _base: 1,
 48            _phantom: PhantomData,
 49        }
 50    }
 51}
 52
 53impl<T: UnsignedInteger> FheUintBlocks<Vec<u8>, T> {
 54    pub fn alloc_from_infos<A, BE: Backend>(module: &Module<BE>, infos: &A) -> Self
 55    where
 56        A: GLWEInfos,
 57    {
 58        Self::alloc(module, infos.base2k(), infos.k(), infos.rank())
 59    }
 60
 61    pub fn alloc<BE: Backend>(module: &Module<BE>, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 62        Self {
 63            blocks: (0..T::WORD_SIZE)
 64                .map(|_| GLWE::alloc(module.n().into(), base2k, k, rank))
 65                .collect(),
 66            _base: 1,
 67            _phantom: PhantomData,
 68        }
 69    }
 70}
 71
 72impl<D: DataMut, T: UnsignedInteger + ToBits> FheUintBlocks<D, T> {
 73    pub fn encrypt_sk<S, BE: Backend>(
 74        &mut self,
 75        module: &Module<BE>,
 76        value: T,
 77        sk: &S,
 78        source_xa: &mut Source,
 79        source_xe: &mut Source,
 80        scratch: &mut Scratch<BE>,
 81    ) where
 82        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
 83        Module<BE>: GLWEEncryptSk<BE>,
 84        Scratch<BE>: ScratchTakeCore<BE>,
 85    {
 86        use poulpy_core::layouts::GLWEPlaintextLayout;
 87        use poulpy_hal::layouts::ZnxZero;
 88
 89        #[cfg(debug_assertions)]
 90        {
 91            assert!(module.n().is_multiple_of(T::WORD_SIZE));
 92            assert_eq!(self.n(), module.n() as u32);
 93            assert_eq!(sk.n(), module.n() as u32);
 94        }
 95
 96        let pt_infos = GLWEPlaintextLayout {
 97            n: self.n(),
 98            base2k: self.base2k(),
 99            k: 1_usize.into(),
100        };
101
102        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
103        pt.data.zero();
104
105        for i in 0..T::WORD_SIZE {
106            pt.encode_coeff_i64(value.bit(i) as i64, TorusPrecision(2), 0);
107            self.blocks[i].encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
108        }
109    }
110}
111
112impl<D: DataRef, T: UnsignedInteger + FromBits + ToBits> FheUintBlocks<D, T> {
113    pub fn decrypt<S, BE: Backend>(&self, module: &Module<BE>, sk: &S, scratch: &mut Scratch<BE>) -> T
114    where
115        Module<BE>: GLWEDecrypt<BE>,
116        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
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 pt_infos = GLWEPlaintextLayout {
127            n: self.n(),
128            base2k: self.base2k(),
129            k: self.k(),
130        };
131
132        let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
133
134        let mut bits: Vec<u8> = vec![0u8; T::WORD_SIZE];
135
136        let base2k: usize = self.base2k().into();
137        let scale: f64 = 4.0 / ((1 << base2k) as f64);
138
139        for (i, bit) in bits.iter_mut().enumerate().take(T::WORD_SIZE) {
140            self.blocks[i].decrypt(module, &mut pt, sk, scratch_1);
141            let value: i64 = pt.decode_coeff_i64(base2k.into(), 0);
142            *bit = ((value as f64) * scale).round() as u8;
143        }
144
145        T::from_bits(&bits)
146    }
147
148    pub fn noise<S, BE: Backend>(&self, module: &Module<BE>, sk: &S, want: T, scratch: &mut Scratch<BE>) -> Vec<f64>
149    where
150        Module<BE>: GLWENoise<BE>,
151        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
152        Scratch<BE>: ScratchTakeCore<BE>,
153    {
154        #[cfg(debug_assertions)]
155        {
156            assert!(module.n().is_multiple_of(T::WORD_SIZE));
157            assert_eq!(self.n(), module.n() as u32);
158            assert_eq!(sk.n(), module.n() as u32);
159        }
160
161        let pt_infos = GLWEPlaintextLayout {
162            n: self.n(),
163            base2k: self.base2k(),
164            k: 1_usize.into(),
165        };
166
167        let (mut pt_want, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
168        pt_want.data.zero();
169
170        let mut noise: Vec<f64> = vec![0f64; T::WORD_SIZE];
171
172        for (i, noise_i) in noise.iter_mut().enumerate().take(T::WORD_SIZE) {
173            pt_want.encode_coeff_i64(want.bit(i) as i64, TorusPrecision(2), 0);
174            *noise_i = self.blocks[i].noise(module, sk, &pt_want, scratch_1);
175        }
176
177        noise
178    }
179}