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, ScratchTakeBasic};
 17use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, Stats, ZnxZero};
 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 noise<S, M, BE: Backend>(
 85        &self,
 86        module: &M,
 87        row: usize,
 88        col: usize,
 89        want: T,
 90        sk: &S,
 91        scratch: &mut Scratch<BE>,
 92    ) -> Vec<Stats>
 93    where
 94        S: GLWESecretPreparedToRef<BE>,
 95        M: GGSWNoise<BE>,
 96        Scratch<BE>: ScratchTakeCore<BE>,
 97    {
 98        let mut stats = Vec::new();
 99        for (i, ggsw) in self.bits.iter().enumerate() {
100            use poulpy_hal::layouts::ZnxViewMut;
101            let (mut pt_want, scratch_1) = scratch.take_scalar_znx(self.n().into(), 1);
102            pt_want.zero();
103            pt_want.at_mut(0, 0)[0] = want.bit(i) as i64;
104            stats.push(ggsw.noise(module, row, col, &pt_want, sk, scratch_1));
105        }
106        stats
107    }
108}
109
110impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintPrepareDebug<BRA, T, BE> for Module<BE>
111where
112    Self: ModuleN + LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
113    Scratch<BE>: ScratchTakeCore<BE>,
114{
115    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
116        &self,
117        res: &mut FheUintPreparedDebug<DM, T>,
118        bits: &FheUint<DR0, T>,
119        key: &BDDKeyPrepared<DR1, BRA, BE>,
120        scratch: &mut Scratch<BE>,
121    ) where
122        DM: DataMut,
123        DR0: DataRef,
124        DR1: DataRef,
125    {
126        let (_, scratch_1) = scratch.take_ggsw(res);
127        let (mut tmp_lwe, scratch_2) = scratch_1.take_lwe(bits);
128        for (bit, dst) in res.bits.iter_mut().enumerate() {
129            bits.get_bit_lwe(
130                self,
131                bit,
132                &mut tmp_lwe,
133                key.ks_glwe.as_ref(),
134                &key.ks_lwe,
135                scratch_2,
136            );
137            key.cbt
138                .execute_to_constant(self, dst, &tmp_lwe, 1, 1, scratch_2);
139        }
140    }
141}
142
143impl<D: DataMut, T: UnsignedInteger> FheUintPreparedDebug<D, T> {
144    pub fn prepare<BRA, M, O, K, BE: Backend>(
145        &mut self,
146        module: &M,
147        other: &FheUint<O, T>,
148        key: &BDDKeyPrepared<K, BRA, BE>,
149        scratch: &mut Scratch<BE>,
150    ) where
151        BRA: BlindRotationAlgo,
152        O: DataRef,
153        K: DataRef,
154        M: FheUintPrepareDebug<BRA, T, BE>,
155        Scratch<BE>: ScratchTakeCore<BE>,
156    {
157        module.fhe_uint_debug_prepare(self, other, key, scratch);
158    }
159}