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