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