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