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    pub fn alloc<A, M>(module: &M, infos: &A) -> Self
 97    where
 98        A: GGSWInfos,
 99        M: FheUintBlocksPreparedFactory<T, BE>,
100    {
101        module.alloc_fhe_uint_blocks_prepared_from_infos(infos)
102    }
103
104    pub fn alloc_with<M>(module: &M, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
105    where
106        M: FheUintBlocksPreparedFactory<T, BE>,
107    {
108        module.alloc_fhe_uint_blocks_prepared(base2k, k, dnum, dsize, rank)
109    }
110}
111
112impl<T: UnsignedInteger + ToBits, BE: Backend> FheUintBlocksPreparedEncryptSk<T, BE> for Module<BE> where
113    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>
114{
115}
116
117pub trait FheUintBlocksPreparedEncryptSk<T: UnsignedInteger + ToBits, BE: Backend>
118where
119    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
120{
121    fn fhe_uint_blocks_prepared_encrypt_sk<DM, S>(
122        &self,
123        res: &mut FheUintBlocksPrepared<DM, T, BE>,
124        value: T,
125        sk: &S,
126        source_xa: &mut Source,
127        source_xe: &mut Source,
128        scratch: &mut Scratch<BE>,
129    ) where
130        DM: DataMut,
131        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
132        Scratch<BE>: ScratchTakeCore<BE>,
133    {
134        use poulpy_hal::{api::ScratchTakeBasic, layouts::ZnxZero};
135
136        assert!(self.n().is_multiple_of(T::WORD_SIZE));
137        assert_eq!(res.n(), self.n() as u32);
138        assert_eq!(sk.n(), self.n() as u32);
139
140        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
141        let (mut pt, scratch_2) = scratch_1.take_scalar_znx(self.n(), 1);
142        pt.zero();
143
144        for i in 0..T::WORD_SIZE {
145            use poulpy_hal::layouts::ZnxViewMut;
146            pt.at_mut(0, 0)[0] = value.bit(i) as i64;
147            tmp_ggsw.encrypt_sk(self, &pt, sk, source_xa, source_xe, scratch_2);
148            res.blocks[i].prepare(self, &tmp_ggsw, scratch_2);
149        }
150    }
151}
152
153impl<D: DataMut, T: UnsignedInteger + ToBits, BE: Backend> FheUintBlocksPrepared<D, T, BE> {
154    pub fn encrypt_sk<M, S>(
155        &mut self,
156        module: &M,
157        value: T,
158        sk: &S,
159        source_xa: &mut Source,
160        source_xe: &mut Source,
161        scratch: &mut Scratch<BE>,
162    ) where
163        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
164        M: FheUintBlocksPreparedEncryptSk<T, BE>,
165        Scratch<BE>: ScratchTakeCore<BE>,
166    {
167        module.fhe_uint_blocks_prepared_encrypt_sk(self, value, sk, source_xa, source_xe, scratch);
168    }
169}
170
171impl<D: DataRef, T: UnsignedInteger, B: Backend> LWEInfos for FheUintBlocksPrepared<D, T, B> {
172    fn base2k(&self) -> poulpy_core::layouts::Base2K {
173        self.blocks[0].base2k()
174    }
175
176    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
177        self.blocks[0].k()
178    }
179
180    fn n(&self) -> poulpy_core::layouts::Degree {
181        self.blocks[0].n()
182    }
183}
184
185impl<D: DataRef, T: UnsignedInteger, B: Backend> GLWEInfos for FheUintBlocksPrepared<D, T, B> {
186    fn rank(&self) -> poulpy_core::layouts::Rank {
187        self.blocks[0].rank()
188    }
189}
190
191impl<D: DataRef, T: UnsignedInteger, B: Backend> GGSWInfos for FheUintBlocksPrepared<D, T, B> {
192    fn dsize(&self) -> poulpy_core::layouts::Dsize {
193        self.blocks[0].dsize()
194    }
195
196    fn dnum(&self) -> poulpy_core::layouts::Dnum {
197        self.blocks[0].dnum()
198    }
199}