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