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