1use std::marker::PhantomData;
  2
  3use crate::tfhe::bdd_arithmetic::{BDDKeyPrepared, FheUintBlockDebugPrepare, ToBits};
  4use crate::tfhe::{
  5    bdd_arithmetic::{FheUintBlocks, UnsignedInteger},
  6    blind_rotation::BlindRotationAlgo,
  7    circuit_bootstrapping::CirtuitBootstrappingExecute,
  8};
  9use poulpy_core::GGSWNoise;
 10#[cfg(test)]
 11use poulpy_core::layouts::{Base2K, Dnum, Dsize, Rank, TorusPrecision};
 12use poulpy_core::layouts::{GGSW, GLWESecretPreparedToRef};
 13use poulpy_core::{
 14    LWEFromGLWE, ScratchTakeCore,
 15    layouts::{GGSWInfos, GGSWPreparedFactory, GLWEInfos, LWE, LWEInfos},
 16};
 17#[cfg(test)]
 18use poulpy_hal::api::ModuleN;
 19use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
 20
 21pub(crate) struct FheUintBlocksPreparedDebug<D: Data, T: UnsignedInteger> {
 22    pub(crate) blocks: Vec<GGSW<D>>,
 23    pub(crate) _base: u8,
 24    pub(crate) _phantom: PhantomData<T>,
 25}
 26
 27#[cfg(test)]
 28impl<T: UnsignedInteger> FheUintBlocksPreparedDebug<Vec<u8>, T> {
 29    #[allow(dead_code)]
 30    pub(crate) fn alloc<A, M>(module: &M, infos: &A) -> Self
 31    where
 32        M: ModuleN,
 33        A: GGSWInfos,
 34    {
 35        Self::alloc_with(
 36            module,
 37            infos.base2k(),
 38            infos.k(),
 39            infos.dnum(),
 40            infos.dsize(),
 41            infos.rank(),
 42        )
 43    }
 44
 45    #[allow(dead_code)]
 46    pub(crate) fn alloc_with<M>(module: &M, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
 47    where
 48        M: ModuleN,
 49    {
 50        Self {
 51            blocks: (0..T::WORD_SIZE)
 52                .map(|_| GGSW::alloc(module.n().into(), base2k, k, rank, dnum, dsize))
 53                .collect(),
 54            _base: 1,
 55            _phantom: PhantomData,
 56        }
 57    }
 58}
 59
 60impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintBlocksPreparedDebug<D, T> {
 61    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 62        self.blocks[0].base2k()
 63    }
 64
 65    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 66        self.blocks[0].k()
 67    }
 68
 69    fn n(&self) -> poulpy_core::layouts::Degree {
 70        self.blocks[0].n()
 71    }
 72}
 73
 74impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintBlocksPreparedDebug<D, T> {
 75    fn rank(&self) -> poulpy_core::layouts::Rank {
 76        self.blocks[0].rank()
 77    }
 78}
 79
 80impl<D: DataRef, T: UnsignedInteger> GGSWInfos for FheUintBlocksPreparedDebug<D, T> {
 81    fn dsize(&self) -> poulpy_core::layouts::Dsize {
 82        self.blocks[0].dsize()
 83    }
 84
 85    fn dnum(&self) -> poulpy_core::layouts::Dnum {
 86        self.blocks[0].dnum()
 87    }
 88}
 89
 90impl<D: DataRef, T: UnsignedInteger + ToBits> FheUintBlocksPreparedDebug<D, T> {
 91    pub(crate) fn noise<S, M, BE: Backend>(&self, module: &M, sk: &S, want: T)
 92    where
 93        S: GLWESecretPreparedToRef<BE>,
 94        M: GGSWNoise<BE>,
 95    {
 96        for (i, ggsw) in self.blocks.iter().enumerate() {
 97            use poulpy_hal::layouts::{ScalarZnx, ZnxViewMut};
 98            let mut pt_want = ScalarZnx::alloc(self.n().into(), 1);
 99            pt_want.at_mut(0, 0)[0] = want.bit(i) as i64;
100            ggsw.print_noise(module, sk, &pt_want);
101        }
102    }
103}
104
105impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintBlockDebugPrepare<BRA, T, BE> for Module<BE>
106where
107    Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
108    Scratch<BE>: ScratchTakeCore<BE>,
109{
110    fn fhe_uint_block_debug_prepare<DM, DR0, DR1>(
111        &self,
112        res: &mut FheUintBlocksPreparedDebug<DM, T>,
113        bits: &FheUintBlocks<DR0, T>,
114        key: &BDDKeyPrepared<DR1, BRA, BE>,
115        scratch: &mut Scratch<BE>,
116    ) where
117        DM: DataMut,
118        DR0: DataRef,
119        DR1: DataRef,
120    {
121        assert_eq!(res.blocks.len(), bits.blocks.len());
122
123        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&bits.blocks[0]); //TODO: add TakeLWE
124        for (dst, src) in res.blocks.iter_mut().zip(bits.blocks.iter()) {
125            lwe.from_glwe(self, src, &key.ks, scratch);
126            key.cbt.execute_to_constant(self, dst, &lwe, 1, 1, scratch);
127        }
128    }
129}
130
131impl<D: DataMut, T: UnsignedInteger> FheUintBlocksPreparedDebug<D, T> {
132    pub fn prepare<BRA, M, O, K, BE: Backend>(
133        &mut self,
134        module: &M,
135        other: &FheUintBlocks<O, T>,
136        key: &BDDKeyPrepared<K, BRA, BE>,
137        scratch: &mut Scratch<BE>,
138    ) where
139        BRA: BlindRotationAlgo,
140        O: DataRef,
141        K: DataRef,
142        M: FheUintBlockDebugPrepare<BRA, T, BE>,
143        Scratch<BE>: ScratchTakeCore<BE>,
144    {
145        module.fhe_uint_block_debug_prepare(self, other, key, scratch);
146    }
147}