1use itertools::Itertools;
  2use phantom_zone::*;
  3use rand::{thread_rng, Rng, RngCore};
  4
  5/// Code that runs when conditional branch is `True`
  6fn circuit_branch_true(a: &FheUint8, b: &FheUint8) -> FheUint8 {
  7    a + b
  8}
  9
 10/// Code that runs when conditional branch is `False`
 11fn circuit_branch_false(a: &FheUint8, b: &FheUint8) -> FheUint8 {
 12    a * b
 13}
 14
 15// Conditional branching (ie. If and else) are generally expensive in encrypted
 16// domain. The code must execute all the branches, and, as apparent, the
 17// runtime cost grows exponentially with no. of conditional branches.
 18//
 19// In general we recommend to write branchless code. In case the code cannot be
 20// modified to be branchless, the code must execute all branches and use a
 21// muxer to select correct output at the end.
 22//
 23// Below we showcase example of a single conditional branch in encrypted domain.
 24// The code executes both the branches (i.e. program runs both If and Else) and
 25// selects output of one of the branches with a mux.
 26fn main() {
 27    set_parameter_set(ParameterSelector::NonInteractiveLTE2Party);
 28
 29    // set application's common reference seed
 30    let mut seed = [0u8; 32];
 31    thread_rng().fill_bytes(&mut seed);
 32    set_common_reference_seed(seed);
 33
 34    let no_of_parties = 2;
 35
 36    // Generate client keys
 37    let cks = (0..no_of_parties).map(|_| gen_client_key()).collect_vec();
 38
 39    // Generate server key shares
 40    let server_key_shares = cks
 41        .iter()
 42        .enumerate()
 43        .map(|(id, k)| gen_server_key_share(id, no_of_parties, k))
 44        .collect_vec();
 45
 46    // Aggregate server key shares and set the server key
 47    let server_key = aggregate_server_key_shares(&server_key_shares);
 48    server_key.set_server_key();
 49
 50    // -------
 51
 52    // User 0 encrypts their private input `v_a` and User 1 encrypts their
 53    // private input `v_b`. We want to execute:
 54    //
 55    // if v_a < v_b:
 56    //      return v_a + v_b
 57    // else:
 58    //      return v_a * v_b
 59    //
 60    // We define two functions
 61    //      (1) `circuit_branch_true`: which executes v_a + v_b in encrypted domain.
 62    //      (2) `circuit_branch_false`: which executes v_a * v_b in encrypted
 63    //                                  domain.
 64    //
 65    // The circuit runs both `circuit_branch_true` and `circuit_branch_false` and
 66    // then selects the output of `circuit_branch_true` if `v_a < v_b == TRUE`
 67    // otherwise selects the output of `circuit_branch_false` if `v_a < v_b ==
 68    // FALSE` using mux.
 69
 70    // Clients private inputs
 71    let v_a = thread_rng().gen::<u8>();
 72    let v_b = thread_rng().gen::<u8>();
 73    let v_a_enc = cks[0]
 74        .encrypt(vec![v_a].as_slice())
 75        .unseed::<Vec<Vec<u64>>>()
 76        .key_switch(0)
 77        .extract_at(0);
 78    let v_b_enc = cks[1]
 79        .encrypt(vec![v_b].as_slice())
 80        .unseed::<Vec<Vec<u64>>>()
 81        .key_switch(1)
 82        .extract_at(0);
 83
 84    // Run both branches
 85    let out_true_enc = circuit_branch_true(&v_a_enc, &v_b_enc);
 86    let out_false_enc = circuit_branch_false(&v_a_enc, &v_b_enc);
 87
 88    // define condition select v_a < v_b
 89    let selector_bit = v_a_enc.lt(&v_b_enc);
 90
 91    // select output of `circuit_branch_true` if selector_bit == TRUE otherwise
 92    // select output of `circuit_branch_false`
 93    let out_enc = out_true_enc.mux(&out_false_enc, &selector_bit);
 94
 95    let out = cks[0].aggregate_decryption_shares(
 96        &out_enc,
 97        &cks.iter()
 98            .map(|k| k.gen_decryption_share(&out_enc))
 99            .collect_vec(),
100    );
101    let want_out = if v_a < v_b {
102        v_a.wrapping_add(v_b)
103    } else {
104        v_a.wrapping_mul(v_b)
105    };
106    assert_eq!(out, want_out);
107}