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 now = std::time::Instant::now();
51    let (ct_c0_a, ct_c0_b) = {
52        let ct = c0_batched_to_send.unseed::<Vec<Vec<u64>>>().key_switch(0);
53        (ct.extract_at(0), ct.extract_at(1))
54    };
55    // println!(
56    //     "Time to unseed, key switch, and extract 2 ciphertexts: {:?}",
57    //     now.elapsed()
58    // );
59
60    // extract a and b from client1 inputs
61    let (ct_c1_a, ct_c1_b) = {
62        let ct = c1_batch_to_send.unseed::<Vec<Vec<u64>>>().key_switch(1);
63        (ct.extract_at(0), ct.extract_at(1))
64    };
65
66    let now = std::time::Instant::now();
67    let c_out = fhe_circuit(&ct_c0_a, &ct_c1_a, &ct_c0_b, &ct_c1_b);
68    println!("Circuit Time: {:?}", now.elapsed());
69
70    // decrypt c_out
71    let decryption_shares = cks
72        .iter()
73        .map(|k| k.gen_decryption_share(&c_out))
74        .collect_vec();
75    let m_out = cks[0].aggregate_decryption_shares(&c_out, &decryption_shares);
76    let m_expected = circuit(c0_a, c1_a, c0_b, c1_b);
77    assert!(m_expected == m_out);
78}