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