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