1use std::marker::PhantomData;
  2
  3use crate::bin_fhe::bdd_arithmetic::{BDDKeyPrepared, FheUint, FheUintPrepareDebug, ToBits};
  4use crate::bin_fhe::{
  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(module, infos.base2k(), infos.k(), infos.dnum(), infos.dsize(), infos.rank())
 31    }
 32
 33    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
 34    where
 35        M: ModuleN,
 36    {
 37        Self {
 38            bits: (0..T::BITS)
 39                .map(|_| GGSW::alloc(module.n().into(), base2k, k, rank, dnum, dsize))
 40                .collect(),
 41            _phantom: PhantomData,
 42        }
 43    }
 44}
 45
 46impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUintPreparedDebug<D, T> {
 47    fn base2k(&self) -> poulpy_core::layouts::Base2K {
 48        self.bits[0].base2k()
 49    }
 50
 51    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
 52        self.bits[0].k()
 53    }
 54
 55    fn n(&self) -> poulpy_core::layouts::Degree {
 56        self.bits[0].n()
 57    }
 58}
 59
 60impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUintPreparedDebug<D, T> {
 61    fn rank(&self) -> poulpy_core::layouts::Rank {
 62        self.bits[0].rank()
 63    }
 64}
 65
 66impl<D: DataRef, T: UnsignedInteger> GGSWInfos for FheUintPreparedDebug<D, T> {
 67    fn dsize(&self) -> poulpy_core::layouts::Dsize {
 68        self.bits[0].dsize()
 69    }
 70
 71    fn dnum(&self) -> poulpy_core::layouts::Dnum {
 72        self.bits[0].dnum()
 73    }
 74}
 75
 76impl<D: DataRef, T: UnsignedInteger + ToBits> FheUintPreparedDebug<D, T> {
 77    pub fn noise<S, M, BE: Backend>(
 78        &self,
 79        module: &M,
 80        row: usize,
 81        col: usize,
 82        want: T,
 83        sk: &S,
 84        scratch: &mut Scratch<BE>,
 85    ) -> Vec<Stats>
 86    where
 87        S: GLWESecretPreparedToRef<BE>,
 88        M: GGSWNoise<BE>,
 89        Scratch<BE>: ScratchTakeCore<BE>,
 90    {
 91        let mut stats = Vec::new();
 92        for (i, ggsw) in self.bits.iter().enumerate() {
 93            use poulpy_hal::layouts::ZnxViewMut;
 94            let (mut pt_want, scratch_1) = scratch.take_scalar_znx(self.n().into(), 1);
 95            pt_want.zero();
 96            pt_want.at_mut(0, 0)[0] = want.bit(i) as i64;
 97            stats.push(ggsw.noise(module, row, col, &pt_want, sk, scratch_1));
 98        }
 99        stats
100    }
101}
102
103impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintPrepareDebug<BRA, T, BE> for Module<BE>
104where
105    Self: ModuleN + LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
106    Scratch<BE>: ScratchTakeCore<BE>,
107{
108    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
109        &self,
110        res: &mut FheUintPreparedDebug<DM, T>,
111        bits: &FheUint<DR0, T>,
112        key: &BDDKeyPrepared<DR1, BRA, BE>,
113        scratch: &mut Scratch<BE>,
114    ) where
115        DM: DataMut,
116        DR0: DataRef,
117        DR1: DataRef,
118    {
119        let (_, scratch_1) = scratch.take_ggsw(res);
120        let (mut tmp_lwe, scratch_2) = scratch_1.take_lwe(bits);
121        for (bit, dst) in res.bits.iter_mut().enumerate() {
122            bits.get_bit_lwe(self, bit, &mut tmp_lwe, key.ks_glwe.as_ref(), &key.ks_lwe, scratch_2);
123            key.cbt.execute_to_constant(self, dst, &tmp_lwe, 1, 1, scratch_2);
124        }
125    }
126}
127
128impl<D: DataMut, T: UnsignedInteger> FheUintPreparedDebug<D, T> {
129    pub fn prepare<BRA, M, O, K, BE: Backend>(
130        &mut self,
131        module: &M,
132        other: &FheUint<O, T>,
133        key: &BDDKeyPrepared<K, BRA, BE>,
134        scratch: &mut Scratch<BE>,
135    ) where
136        BRA: BlindRotationAlgo,
137        O: DataRef,
138        K: DataRef,
139        M: FheUintPrepareDebug<BRA, T, BE>,
140        Scratch<BE>: ScratchTakeCore<BE>,
141    {
142        module.fhe_uint_debug_prepare(self, other, key, scratch);
143    }
144}