1use bin_rs::*;
 2use itertools::Itertools;
 3use rand::{thread_rng, Rng, RngCore};
 4
 5fn circuit(a: u8, b: u8, c: u8, d: u8) -> u8 {
 6    ((a + b) * c) * d
 7}
 8
 9fn fhe_circuit(a: &FheUint8, b: &FheUint8, c: &FheUint8, d: &FheUint8) -> FheUint8 {
10    &(&(a + b) * c) * d
11}
12
13fn main() {
14    set_parameter_set(ParameterSelector::NonInteractiveLTE2Party);
15
16    // set CRS
17    let mut seed = [0u8; 32];
18    thread_rng().fill_bytes(&mut seed);
19    set_common_reference_seed(seed);
20
21    let no_of_parties = 2;
22
23    // Generate client keys
24    let cks = (0..no_of_parties).map(|_| gen_client_key()).collect_vec();
25
26    // client 0 encrypts private input
27    let c0_a = thread_rng().gen::<u8>();
28    let c0_b = thread_rng().gen::<u8>();
29    let c0_batched_to_send = cks[0].encrypt(vec![c0_a, c0_b].as_slice());
30
31    // client 1 encrypts private input
32    let c1_a = thread_rng().gen::<u8>();
33    let c1_b = thread_rng().gen::<u8>();
34    let c1_batch_to_send = cks[1].encrypt(vec![c1_a, c1_b].as_slice());
35
36    // Both client indenpendently generate their server key shares
37    let server_key_shares = cks
38        .iter()
39        .enumerate()
40        .map(|(id, k)| gen_server_key_share(id, no_of_parties, k))
41        .collect_vec();
42
43    // Server side
44
45    // aggregates shares and generates server key
46    let server_key = aggregate_server_key_shares(&server_key_shares);
47    server_key.set_server_key();
48
49    // extract a and b from client0 inputs
50    let (ct_c0_a, ct_c0_b) = {
51        let ct = c0_batched_to_send.unseed::<Vec<Vec<u64>>>().key_switch(0);
52        (ct.extract(0), ct.extract(1))
53    };
54
55    // extract a and b from client1 inputs
56    let (ct_c1_a, ct_c1_b) = {
57        let ct = c1_batch_to_send.unseed::<Vec<Vec<u64>>>().key_switch(1);
58        (ct.extract(0), ct.extract(1))
59    };
60
61    let now = std::time::Instant::now();
62    let c_out = fhe_circuit(&ct_c0_a, &ct_c1_a, &ct_c0_b, &ct_c1_b);
63    println!("Circuit Time: {:?}", now.elapsed());
64
65    // decrypt c_out
66    let decryption_shares = cks
67        .iter()
68        .map(|k| k.gen_decryption_share(&c_out))
69        .collect_vec();
70    let m_out = cks[0].aggregate_decryption_shares(&c_out, &decryption_shares);
71    let m_expected = circuit(c0_a, c1_a, c0_b, c1_b);
72    assert!(m_expected == m_out);
73}