1use poulpy_core::{
  2    GLWEOperations,
  3    layouts::{
  4        GGLWEAutomorphismKeyLayout, GGLWETensorKeyLayout, GGSWCiphertext, GGSWCiphertextLayout, GLWECiphertext,
  5        GLWECiphertextLayout, GLWEPlaintext, GLWESecret, LWECiphertext, LWECiphertextLayout, LWEInfos, LWEPlaintext, LWESecret,
  6        prepared::{GGSWCiphertextPrepared, GLWESecretPrepared, PrepareAlloc},
  7    },
  8};
  9use std::time::Instant;
 10
 11#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 12use poulpy_backend::FFT64Avx as BackendImpl;
 13
 14#[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
 15use poulpy_backend::FFT64Ref as BackendImpl;
 16
 17use poulpy_hal::{
 18    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, ZnNormalizeInplace},
 19    layouts::{Module, ScalarZnx, ScratchOwned, ZnxView, ZnxViewMut},
 20    source::Source,
 21};
 22
 23use poulpy_schemes::tfhe::{
 24    blind_rotation::{BlindRotationKeyLayout, CGGI},
 25    circuit_bootstrapping::{
 26        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
 27        CircuitBootstrappingKeyPrepared, CirtuitBootstrappingExecute,
 28    },
 29};
 30
 31fn main() {
 32    // GLWE ring degree
 33    let n_glwe: usize = 1024;
 34
 35    // Module provides access to the backend arithmetic
 36    let module: Module<BackendImpl> = Module::<BackendImpl>::new(n_glwe as u64);
 37
 38    // Base 2 loga
 39    let base2k: usize = 13;
 40
 41    // Lookup table extension factor
 42    let extension_factor: usize = 1;
 43
 44    // GLWE rank
 45    let rank: usize = 1;
 46
 47    // LWE degree
 48    let n_lwe: usize = 574;
 49
 50    // LWE plaintext modulus
 51    let k_lwe_pt: usize = 1;
 52
 53    // LWE ciphertext modulus
 54    let k_lwe_ct: usize = 13;
 55
 56    // LWE block binary key block size
 57    let block_size: usize = 7;
 58
 59    // GGSW output number of dnum
 60    let rows_ggsw_res: usize = 2;
 61
 62    // GGSW output modulus
 63    let k_ggsw_res: usize = (rows_ggsw_res + 1) * base2k;
 64
 65    // Blind rotation key GGSW number of dnum
 66    let rows_brk: usize = rows_ggsw_res + 1;
 67
 68    // Blind rotation key GGSW modulus
 69    let k_brk: usize = (rows_brk + 1) * base2k;
 70
 71    // GGLWE automorphism keys number of dnum
 72    let rows_trace: usize = rows_ggsw_res + 1;
 73
 74    // GGLWE automorphism keys modulus
 75    let k_trace: usize = (rows_trace + 1) * base2k;
 76
 77    // GGLWE tensor key number of dnum
 78    let rows_tsk: usize = rows_ggsw_res + 1;
 79
 80    // GGLWE tensor key modulus
 81    let k_tsk: usize = (rows_tsk + 1) * base2k;
 82
 83    let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
 84        layout_brk: BlindRotationKeyLayout {
 85            n_glwe: n_glwe.into(),
 86            n_lwe: n_lwe.into(),
 87            base2k: base2k.into(),
 88            k: k_brk.into(),
 89            dnum: rows_brk.into(),
 90            rank: rank.into(),
 91        },
 92        layout_atk: GGLWEAutomorphismKeyLayout {
 93            n: n_glwe.into(),
 94            base2k: base2k.into(),
 95            k: k_trace.into(),
 96            dnum: rows_trace.into(),
 97            dsize: 1_u32.into(),
 98            rank: rank.into(),
 99        },
100        layout_tsk: GGLWETensorKeyLayout {
101            n: n_glwe.into(),
102            base2k: base2k.into(),
103            k: k_tsk.into(),
104            dnum: rows_tsk.into(),
105            dsize: 1_u32.into(),
106            rank: rank.into(),
107        },
108    };
109
110    let ggsw_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
111        n: n_glwe.into(),
112        base2k: base2k.into(),
113        k: k_ggsw_res.into(),
114        dnum: rows_ggsw_res.into(),
115        dsize: 1_u32.into(),
116        rank: rank.into(),
117    };
118
119    let lwe_infos = LWECiphertextLayout {
120        n: n_lwe.into(),
121        k: k_lwe_ct.into(),
122        base2k: base2k.into(),
123    };
124
125    // Scratch space (4MB)
126    let mut scratch: ScratchOwned<BackendImpl> = ScratchOwned::alloc(1 << 22);
127
128    // Secret key sampling source
129    let mut source_xs: Source = Source::new([1u8; 32]);
130
131    // Public randomness sampling source
132    let mut source_xa: Source = Source::new([1u8; 32]);
133
134    // Noise sampling source
135    let mut source_xe: Source = Source::new([1u8; 32]);
136
137    // LWE secret
138    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
139    sk_lwe.fill_binary_block(block_size, &mut source_xs);
140    sk_lwe.fill_zero();
141
142    // GLWE secret
143    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n_glwe.into(), rank.into());
144    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
145    // sk_glwe.fill_zero();
146
147    // GLWE secret prepared (opaque backend dependant write only struct)
148    let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BackendImpl> = sk_glwe.prepare_alloc(&module, scratch.borrow());
149
150    // Plaintext value to circuit bootstrap
151    let data: i64 = 1 % (1 << k_lwe_pt);
152
153    // LWE plaintext
154    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_with(base2k.into(), k_lwe_pt.into());
155
156    // LWE plaintext(data * 2^{- (k_lwe_pt - 1)})
157    pt_lwe.encode_i64(data, (k_lwe_pt + 1).into()); // +1 for padding bit
158
159    // Normalize plaintext to nicely print coefficients
160    module.zn_normalize_inplace(
161        pt_lwe.n().into(),
162        base2k,
163        pt_lwe.data_mut(),
164        0,
165        scratch.borrow(),
166    );
167    println!("pt_lwe: {pt_lwe}");
168
169    // LWE ciphertext
170    let mut ct_lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&lwe_infos);
171
172    // Encrypt LWE Plaintext
173    ct_lwe.encrypt_sk(&module, &pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe);
174
175    let now: Instant = Instant::now();
176
177    // Circuit bootstrapping evaluation key
178    let cbt_key: CircuitBootstrappingKey<Vec<u8>, CGGI> = CircuitBootstrappingKey::encrypt_sk(
179        &module,
180        &sk_lwe,
181        &sk_glwe,
182        &cbt_infos,
183        &mut source_xa,
184        &mut source_xe,
185        scratch.borrow(),
186    );
187    println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
188
189    // Output GGSW
190    let mut res: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&ggsw_infos);
191
192    // Circuit bootstrapping key prepared (opaque backend dependant write only struct)
193    let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, CGGI, BackendImpl> =
194        cbt_key.prepare_alloc(&module, scratch.borrow());
195
196    // Apply circuit bootstrapping: LWE(data * 2^{- (k_lwe_pt + 2)}) -> GGSW(data)
197    let now: Instant = Instant::now();
198    cbt_prepared.execute_to_constant(
199        &module,
200        &mut res,
201        &ct_lwe,
202        k_lwe_pt,
203        extension_factor,
204        scratch.borrow(),
205    );
206    println!("CBT: {} ms", now.elapsed().as_millis());
207
208    // Allocate "ideal" GGSW(data) plaintext
209    let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
210    pt_ggsw.at_mut(0, 0)[0] = data;
211
212    // Prints noise of GGSW(data)
213    res.print_noise(&module, &sk_glwe_prepared, &pt_ggsw);
214
215    // Tests RLWE(1) * GGSW(data)
216
217    let glwe_infos: GLWECiphertextLayout = GLWECiphertextLayout {
218        n: n_glwe.into(),
219        base2k: base2k.into(),
220        k: (k_ggsw_res - base2k).into(),
221        rank: rank.into(),
222    };
223
224    // GLWE ciphertext modulus
225    let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&glwe_infos);
226
227    // Some GLWE plaintext with signed data
228    let k_glwe_pt: usize = 3;
229    let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&glwe_infos);
230    let mut data_vec: Vec<i64> = vec![0i64; n_glwe];
231    data_vec
232        .iter_mut()
233        .enumerate()
234        .for_each(|(x, y)| *y = (x % (1 << (k_glwe_pt - 1))) as i64 - (1 << (k_glwe_pt - 2)));
235
236    pt_glwe.encode_vec_i64(&data_vec, (k_lwe_pt + 2).into());
237    pt_glwe.normalize_inplace(&module, scratch.borrow());
238
239    println!("{}", pt_glwe);
240
241    // Encrypt
242    ct_glwe.encrypt_sk(
243        &module,
244        &pt_glwe,
245        &sk_glwe_prepared,
246        &mut source_xa,
247        &mut source_xe,
248        scratch.borrow(),
249    );
250
251    // Prepare GGSW output of circuit bootstrapping (opaque backend dependant write only struct)
252    let res_prepared: GGSWCiphertextPrepared<Vec<u8>, BackendImpl> = res.prepare_alloc(&module, scratch.borrow());
253
254    // Apply GLWE x GGSW
255    ct_glwe.external_product_inplace(&module, &res_prepared, scratch.borrow());
256
257    // Decrypt
258    let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&glwe_infos);
259    ct_glwe.decrypt(&module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
260
261    println!("pt_res: {:?}", &pt_res.data.at(0, 0)[..64]);
262}