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