1#[cfg(test)]
  2use crate::tfhe::bdd_arithmetic::FheUintBlocksPrepDebug;
  3use crate::tfhe::{
  4    bdd_arithmetic::{FheUintBlocks, FheUintBlocksPrep, UnsignedInteger},
  5    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk},
  6    circuit_bootstrapping::{
  7        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
  8        CircuitBootstrappingKeyPrepared, CirtuitBootstrappingExecute,
  9    },
 10};
 11use poulpy_core::{
 12    TakeGGSW, TakeGLWECt,
 13    layouts::{
 14        GLWESecret, GLWEToLWEKeyLayout, GLWEToLWESwitchingKey, LWECiphertext, LWESecret,
 15        prepared::{GLWEToLWESwitchingKeyPrepared, Prepare, PrepareAlloc},
 16    },
 17};
 18use poulpy_hal::{
 19    api::{
 20        ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx,
 21        TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
 22        VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
 23        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize,
 24        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing,
 25        VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPrepare,
 26    },
 27    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 28    source::Source,
 29};
 30
 31pub trait BDDKeyInfos {
 32    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout;
 33    fn ks_infos(&self) -> GLWEToLWEKeyLayout;
 34}
 35
 36#[derive(Debug, Clone, Copy)]
 37pub struct BDDKeyLayout {
 38    pub cbt: CircuitBootstrappingKeyLayout,
 39    pub ks: GLWEToLWEKeyLayout,
 40}
 41
 42impl BDDKeyInfos for BDDKeyLayout {
 43    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
 44        self.cbt
 45    }
 46
 47    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
 48        self.ks
 49    }
 50}
 51
 52pub struct BDDKey<CBT, LWE, BRA>
 53where
 54    CBT: Data,
 55    LWE: Data,
 56    BRA: BlindRotationAlgo,
 57{
 58    cbt: CircuitBootstrappingKey<CBT, BRA>,
 59    ks: GLWEToLWESwitchingKey<LWE>,
 60}
 61
 62impl<BRA: BlindRotationAlgo> BDDKey<Vec<u8>, Vec<u8>, BRA> {
 63    pub fn encrypt_sk<DLwe, DGlwe, A, BE: Backend>(
 64        module: &Module<BE>,
 65        sk_lwe: &LWESecret<DLwe>,
 66        sk_glwe: &GLWESecret<DGlwe>,
 67        infos: &A,
 68        source_xa: &mut Source,
 69        source_xe: &mut Source,
 70        scratch: &mut Scratch<BE>,
 71    ) -> Self
 72    where
 73        A: BDDKeyInfos,
 74        DLwe: DataRef,
 75        DGlwe: DataRef,
 76        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<BE>,
 77        Module<BE>: SvpApplyDftToDft<BE>
 78            + VecZnxIdftApplyTmpA<BE>
 79            + VecZnxAddScalarInplace
 80            + VecZnxDftAllocBytes
 81            + VecZnxBigNormalize<BE>
 82            + VecZnxDftApply<BE>
 83            + SvpApplyDftToDftInplace<BE>
 84            + VecZnxIdftApplyConsume<BE>
 85            + VecZnxNormalizeTmpBytes
 86            + VecZnxFillUniform
 87            + VecZnxSubInplace
 88            + VecZnxAddInplace
 89            + VecZnxNormalizeInplace<BE>
 90            + VecZnxAddNormal
 91            + VecZnxNormalize<BE>
 92            + VecZnxSub
 93            + SvpPrepare<BE>
 94            + VecZnxSwitchRing
 95            + SvpPPolAllocBytes
 96            + SvpPPolAlloc<BE>
 97            + VecZnxAutomorphism
 98            + VecZnxAutomorphismInplace<BE>,
 99        Scratch<BE>: TakeVecZnxDft<BE> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<BE> + TakeVecZnxBig<BE>,
100    {
101        let mut ks: GLWEToLWESwitchingKey<Vec<u8>> = GLWEToLWESwitchingKey::alloc(&infos.ks_infos());
102        ks.encrypt_sk(module, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
103
104        Self {
105            cbt: CircuitBootstrappingKey::encrypt_sk(
106                module,
107                sk_lwe,
108                sk_glwe,
109                &infos.cbt_infos(),
110                source_xa,
111                source_xe,
112                scratch,
113            ),
114            ks,
115        }
116    }
117}
118
119pub struct BDDKeyPrepared<CBT, LWE, BRA, BE>
120where
121    CBT: Data,
122    LWE: Data,
123    BRA: BlindRotationAlgo,
124    BE: Backend,
125{
126    cbt: CircuitBootstrappingKeyPrepared<CBT, BRA, BE>,
127    ks: GLWEToLWESwitchingKeyPrepared<LWE, BE>,
128}
129
130impl<CBT: DataMut, LWE: DataMut, BRA: BlindRotationAlgo, BE: Backend> PrepareAlloc<BE, BDDKeyPrepared<CBT, LWE, BRA, BE>>
131    for BDDKey<CBT, LWE, BRA>
132where
133    CircuitBootstrappingKey<CBT, BRA>: PrepareAlloc<BE, CircuitBootstrappingKeyPrepared<CBT, BRA, BE>>,
134    GLWEToLWESwitchingKey<LWE>: PrepareAlloc<BE, GLWEToLWESwitchingKeyPrepared<LWE, BE>>,
135{
136    fn prepare_alloc(&self, module: &Module<BE>, scratch: &mut Scratch<BE>) -> BDDKeyPrepared<CBT, LWE, BRA, BE> {
137        BDDKeyPrepared {
138            cbt: self.cbt.prepare_alloc(module, scratch),
139            ks: self.ks.prepare_alloc(module, scratch),
140        }
141    }
142}
143
144pub trait FheUintPrepare<BE: Backend, OUT, IN> {
145    fn prepare(&self, module: &Module<BE>, out: &mut OUT, bits: &IN, scratch: &mut Scratch<BE>);
146}
147
148impl<CBT, OUT, IN, LWE, BRA, BE, T> FheUintPrepare<BE, FheUintBlocksPrep<OUT, BE, T>, FheUintBlocks<IN, T>>
149    for BDDKeyPrepared<CBT, LWE, BRA, BE>
150where
151    T: UnsignedInteger,
152    CBT: DataRef,
153    OUT: DataMut,
154    IN: DataRef,
155    LWE: DataRef,
156    BRA: BlindRotationAlgo,
157    BE: Backend,
158    Module<BE>: VmpPrepare<BE>
159        + VecZnxRotate
160        + VecZnxDftAllocBytes
161        + VmpApplyDftToDftTmpBytes
162        + VecZnxBigNormalizeTmpBytes
163        + VmpApplyDftToDft<BE>
164        + VmpApplyDftToDftAdd<BE>
165        + VecZnxDftApply<BE>
166        + VecZnxIdftApplyConsume<BE>
167        + VecZnxBigAddSmallInplace<BE>
168        + VecZnxBigNormalize<BE>
169        + VecZnxNormalize<BE>
170        + VecZnxNormalizeTmpBytes,
171    Scratch<BE>: ScratchAvailable + TakeVecZnxDft<BE> + TakeGLWECt + TakeVecZnx + TakeGGSW,
172    CircuitBootstrappingKeyPrepared<CBT, BRA, BE>: CirtuitBootstrappingExecute<BE>,
173{
174    fn prepare(
175        &self,
176        module: &Module<BE>,
177        out: &mut FheUintBlocksPrep<OUT, BE, T>,
178        bits: &FheUintBlocks<IN, T>,
179        scratch: &mut Scratch<BE>,
180    ) {
181        #[cfg(debug_assertions)]
182        {
183            assert_eq!(out.blocks.len(), bits.blocks.len());
184        }
185        let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&bits.blocks[0]); //TODO: add TakeLWE
186        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(out);
187        for (dst, src) in out.blocks.iter_mut().zip(bits.blocks.iter()) {
188            lwe.from_glwe(module, src, &self.ks, scratch_1);
189            self.cbt
190                .execute_to_constant(module, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
191            dst.prepare(module, &tmp_ggsw, scratch_1);
192        }
193    }
194}
195
196#[cfg(test)]
197impl<CBT, OUT, IN, LWE, BRA, BE, T> FheUintPrepare<BE, FheUintBlocksPrepDebug<OUT, T>, FheUintBlocks<IN, T>>
198    for BDDKeyPrepared<CBT, LWE, BRA, BE>
199where
200    T: UnsignedInteger,
201    CBT: DataRef,
202    OUT: DataMut,
203    IN: DataRef,
204    LWE: DataRef,
205    BRA: BlindRotationAlgo,
206    BE: Backend,
207    Module<BE>: VmpPrepare<BE>
208        + VecZnxRotate
209        + VecZnxDftAllocBytes
210        + VmpApplyDftToDftTmpBytes
211        + VecZnxBigNormalizeTmpBytes
212        + VmpApplyDftToDft<BE>
213        + VmpApplyDftToDftAdd<BE>
214        + VecZnxDftApply<BE>
215        + VecZnxIdftApplyConsume<BE>
216        + VecZnxBigAddSmallInplace<BE>
217        + VecZnxBigNormalize<BE>
218        + VecZnxNormalize<BE>
219        + VecZnxNormalizeTmpBytes,
220    Scratch<BE>: ScratchAvailable + TakeVecZnxDft<BE> + TakeGLWECt + TakeVecZnx + TakeGGSW,
221    CircuitBootstrappingKeyPrepared<CBT, BRA, BE>: CirtuitBootstrappingExecute<BE>,
222{
223    fn prepare(
224        &self,
225        module: &Module<BE>,
226        out: &mut FheUintBlocksPrepDebug<OUT, T>,
227        bits: &FheUintBlocks<IN, T>,
228        scratch: &mut Scratch<BE>,
229    ) {
230        #[cfg(debug_assertions)]
231        {
232            assert_eq!(out.blocks.len(), bits.blocks.len());
233        }
234        let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&bits.blocks[0]); //TODO: add TakeLWE
235        for (dst, src) in out.blocks.iter_mut().zip(bits.blocks.iter()) {
236            lwe.from_glwe(module, src, &self.ks, scratch);
237            self.cbt
238                .execute_to_constant(module, dst, &lwe, 1, 1, scratch);
239        }
240    }
241}