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