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