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