1use bin_rs::*;
 2use itertools::Itertools;
 3use rand::{thread_rng, RngCore};
 4
 5fn plain_circuit(a: u8, b: u8, c: u8) -> u8 {
 6    (a + b) * c
 7}
 8
 9fn fhe_circuit(fhe_a: &FheUint8, fhe_b: &FheUint8, fhe_c: &FheUint8) -> FheUint8 {
10    &(fhe_a + fhe_b) * fhe_c
11}
12
13fn main() {
14    set_parameter_set(ParameterSelector::InteractiveLTE2Party);
15    let no_of_parties = 2;
16    let client_keys = (0..no_of_parties)
17        .into_iter()
18        .map(|_| gen_client_key())
19        .collect_vec();
20
21    // set Multi-Party seed
22    let mut seed = [0u8; 32];
23    thread_rng().fill_bytes(&mut seed);
24    set_mp_seed(seed);
25
26    // multi-party key gen round 1
27    let pk_shares = client_keys
28        .iter()
29        .map(|k| gen_mp_keys_phase1(k))
30        .collect_vec();
31
32    // create public key
33    let public_key = aggregate_public_key_shares(&pk_shares);
34
35    // multi-party key gen round 2
36    let server_key_shares = client_keys
37        .iter()
38        .enumerate()
39        .map(|(user_id, k)| gen_mp_keys_phase2(k, user_id, no_of_parties, &public_key))
40        .collect_vec();
41
42    // server aggregates server key shares and sets it
43    let server_key = aggregate_server_key_shares(&server_key_shares);
44    server_key.set_server_key();
45
46    // private inputs
47    let a = 4u8;
48    let b = 6u8;
49    let c = 128u8;
50    let fhe_a = public_key.encrypt(&a);
51    let fhe_b = public_key.encrypt(&b);
52    let fhe_c = public_key.encrypt(&c);
53
54    // fhe evaluation
55    let now = std::time::Instant::now();
56    let fhe_out = fhe_circuit(&fhe_a, &fhe_b, &fhe_c);
57    println!("Circuit time: {:?}", now.elapsed());
58
59    // plain evaluation
60    let out = plain_circuit(a, b, c);
61
62    // generate decryption shares to decrypt ciphertext fhe_out
63    let decryption_shares = client_keys
64        .iter()
65        .map(|k| k.gen_decryption_share(&fhe_out))
66        .collect_vec();
67
68    // decrypt fhe_out using decryption shares
69    let got_out = client_keys[0].aggregate_decryption_shares(&fhe_out, &decryption_shares);
70
71    assert_eq!(got_out, out);
72}