1use std::marker::PhantomData;
  2
  3use poulpy_core::layouts::{
  4    Base2K, Dnum, Dsize, GGSWInfos, GGSWPreparedFactory, GLWEInfos, LWEInfos, Rank, TorusPrecision, prepared::GGSWPrepared,
  5};
  6#[cfg(test)]
  7use poulpy_core::{GGSWEncryptSk, ScratchTakeCore, layouts::GLWESecretPreparedToRef};
  8use poulpy_hal::layouts::{Backend, Data, DataRef, Module};
  9#[cfg(test)]
 10use poulpy_hal::{
 11    api::ModuleN,
 12    layouts::{DataMut, Scratch},
 13    source::Source,
 14};
 15
 16#[cfg(test)]
 17use crate::tfhe::bdd_arithmetic::ToBits;
 18use crate::tfhe::bdd_arithmetic::UnsignedInteger;
 19
 20/// A prepared FHE ciphertext encrypting the bits of an [UnsignedInteger].
 21pub struct FheUintBlocksPrepared<D: Data, T: UnsignedInteger, B: Backend> {
 22    pub(crate) blocks: Vec<GGSWPrepared<D, B>>,
 23    pub(crate) _base: u8,
 24    pub(crate) _phantom: PhantomData<T>,
 25}
 26
 27impl<T: UnsignedInteger, BE: Backend> FheUintBlocksPreparedFactory<T, BE> for Module<BE> where
 28    Self: Sized + GGSWPreparedFactory<BE>
 29{
 30}
 31
 32pub trait FheUintBlocksPreparedFactory<T: UnsignedInteger, BE: Backend>
 33where
 34    Self: Sized + GGSWPreparedFactory<BE>,
 35{
 36    fn alloc_fhe_uint_blocks_prepared(
 37        &self,
 38        base2k: Base2K,
 39        k: TorusPrecision,
 40        dnum: Dnum,
 41        dsize: Dsize,
 42        rank: Rank,
 43    ) -> FheUintBlocksPrepared<Vec<u8>, T, BE> {
 44        FheUintBlocksPrepared {
 45            blocks: (0..T::WORD_SIZE)
 46                .map(|_| GGSWPrepared::alloc(self, base2k, k, dnum, dsize, rank))
 47                .collect(),
 48            _base: 1,
 49            _phantom: PhantomData,
 50        }
 51    }
 52
 53    fn alloc_fhe_uint_blocks_prepared_from_infos<A>(&self, infos: &A) -> FheUintBlocksPrepared<Vec<u8>, T, BE>
 54    where
 55        A: GGSWInfos,
 56    {
 57        self.alloc_fhe_uint_blocks_prepared(
 58            infos.base2k(),
 59            infos.k(),
 60            infos.dnum(),
 61            infos.dsize(),
 62            infos.rank(),
 63        )
 64    }
 65}
 66
 67impl<T: UnsignedInteger, BE: Backend> FheUintBlocksPrepared<Vec<u8>, T, BE> {
 68    #[allow(dead_code)]
 69    pub(crate) fn alloc<A, M>(module: &M, infos: &A) -> Self
 70    where
 71        A: GGSWInfos,
 72        M: FheUintBlocksPreparedFactory<T, BE>,
 73    {
 74        module.alloc_fhe_uint_blocks_prepared_from_infos(infos)
 75    }
 76
 77    #[allow(dead_code)]
 78    pub(crate) fn alloc_with<M>(module: &M, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
 79    where
 80        M: FheUintBlocksPreparedFactory<T, BE>,
 81    {
 82        module.alloc_fhe_uint_blocks_prepared(base2k, k, dnum, dsize, rank)
 83    }
 84}
 85
 86#[cfg(test)]
 87impl<T: UnsignedInteger + ToBits, BE: Backend> FheUintBlocksPreparedEncryptSk<T, BE> for Module<BE> where
 88    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>
 89{
 90}
 91
 92#[cfg(test)]
 93pub trait FheUintBlocksPreparedEncryptSk<T: UnsignedInteger + ToBits, BE: Backend>
 94where
 95    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
 96{
 97    fn fhe_uint_blocks_prepared_encrypt_sk<DM, S>(
 98        &self,
 99        res: &mut FheUintBlocksPrepared<DM, T, BE>,
100        value: T,
101        sk: &S,
102        source_xa: &mut Source,
103        source_xe: &mut Source,
104        scratch: &mut Scratch<BE>,
105    ) where
106        DM: DataMut,
107        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
108        Scratch<BE>: ScratchTakeCore<BE>,
109    {
110        use poulpy_hal::{api::ScratchTakeBasic, layouts::ZnxZero};
111
112        assert!(self.n().is_multiple_of(T::WORD_SIZE));
113        assert_eq!(res.n(), self.n() as u32);
114        assert_eq!(sk.n(), self.n() as u32);
115
116        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
117        let (mut pt, scratch_2) = scratch_1.take_scalar_znx(self.n(), 1);
118        pt.zero();
119
120        for i in 0..T::WORD_SIZE {
121            use poulpy_hal::layouts::ZnxViewMut;
122            pt.at_mut(0, 0)[0] = value.bit(i) as i64;
123            tmp_ggsw.encrypt_sk(self, &pt, sk, source_xa, source_xe, scratch_2);
124            res.blocks[i].prepare(self, &tmp_ggsw, scratch_2);
125        }
126    }
127}
128
129#[cfg(test)]
130impl<D: DataMut, T: UnsignedInteger + ToBits, BE: Backend> FheUintBlocksPrepared<D, T, BE> {
131    pub(crate) fn encrypt_sk<M, S>(
132        &mut self,
133        module: &M,
134        value: T,
135        sk: &S,
136        source_xa: &mut Source,
137        source_xe: &mut Source,
138        scratch: &mut Scratch<BE>,
139    ) where
140        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
141        M: FheUintBlocksPreparedEncryptSk<T, BE>,
142        Scratch<BE>: ScratchTakeCore<BE>,
143    {
144        module.fhe_uint_blocks_prepared_encrypt_sk(self, value, sk, source_xa, source_xe, scratch);
145    }
146}
147
148impl<D: DataRef, T: UnsignedInteger, B: Backend> LWEInfos for FheUintBlocksPrepared<D, T, B> {
149    fn base2k(&self) -> poulpy_core::layouts::Base2K {
150        self.blocks[0].base2k()
151    }
152
153    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
154        self.blocks[0].k()
155    }
156
157    fn n(&self) -> poulpy_core::layouts::Degree {
158        self.blocks[0].n()
159    }
160}
161
162impl<D: DataRef, T: UnsignedInteger, B: Backend> GLWEInfos for FheUintBlocksPrepared<D, T, B> {
163    fn rank(&self) -> poulpy_core::layouts::Rank {
164        self.blocks[0].rank()
165    }
166}
167
168impl<D: DataRef, T: UnsignedInteger, B: Backend> GGSWInfos for FheUintBlocksPrepared<D, T, B> {
169    fn dsize(&self) -> poulpy_core::layouts::Dsize {
170        self.blocks[0].dsize()
171    }
172
173    fn dnum(&self) -> poulpy_core::layouts::Dnum {
174        self.blocks[0].dnum()
175    }
176}