failing bivariate example

failing_bivariate_example.rs

use tfhe::shortint::prelude::*;
use rand::Rng;

pub fn shortint_mux(
    server_key: &ServerKey,
    control_bit: &mut Ciphertext,
    val_true: &mut Ciphertext,
    val_false: &mut Ciphertext,
) -> Ciphertext {
    let mut res = server_key.smart_sub(val_true, val_false);

    let mut bivariate_pbs_possible = server_key.is_functional_bivariate_pbs_possible(&res, control_bit);
    println!("bivariate_pbs_is_pos: {}", bivariate_pbs_possible);
    if !bivariate_pbs_possible {
        server_key.message_extract_assign(&mut res);
    }
    bivariate_pbs_possible = server_key.is_functional_bivariate_pbs_possible(&res, control_bit);
    println!("bivariate_pbs_is_pos: {}", bivariate_pbs_possible);
    if !bivariate_pbs_possible {
        server_key.message_extract_assign(control_bit);
    }
    bivariate_pbs_possible = server_key.is_functional_bivariate_pbs_possible(&res, control_bit);
    println!("bivariate_pbs_is_pos: {}", bivariate_pbs_possible);

    let acc = server_key
        .generate_accumulator_bivariate(|x, y| (if y == 0 { 0 } else { x }));
    server_key.keyswitch_programmable_bootstrap_bivariate_assign(&mut res, control_bit, &acc);
    // server_key.unchecked_mul_lsb_small_carry_assign(&mut res, &mut server_key.smart_neg(control_bit));
    server_key.smart_add_assign(&mut res, val_false);
    res
}

#[test]
fn test_shortint_mux1() {
    let keys = tfhe::shortint::keycache::KEY_CACHE
        .get_from_param(tfhe::shortint::parameters::PARAM_MESSAGE_3_CARRY_3);
    let client_key = keys.client_key();
    let server_key = keys.server_key();

    let mut zero = client_key.encrypt(0);
    let mut one = client_key.encrypt(1);
    let mut x1 = client_key.encrypt(3);
    let mut y1 = client_key.encrypt(5);
    let res1 = shortint_mux(server_key, &mut one, &mut x1, &mut y1);
    let res2 = shortint_mux(server_key, &mut zero, &mut x1, &mut y1);
    assert_eq!(client_key.decrypt(&res1), 3);
    assert_eq!(client_key.decrypt(&res2), 5);
    
    let message_modulus = server_key.message_modulus.0 as u64;
    let mut x_clear = rand::thread_rng().gen_range(0..8);
    let mut y_clear = rand::thread_rng().gen_range(0..8);
    let mut x = client_key.encrypt(x_clear);
    let mut y = client_key.encrypt(y_clear);
    const ROUNDS: usize = 10;
    println!("x_clear: {}, y_clear: {} ", x_clear, y_clear);
    for _ in 0..ROUNDS {
        let y_add_clear = rand::thread_rng().gen_range(0..8);
        let mut y_add = client_key.encrypt(y_add_clear);
        let beginning_y_clear = y_clear;
        server_key.smart_add_assign(&mut y, &mut y_add);
        y_clear = client_key.decrypt(&y);
        println!("beginning_y_clear: {}, y_clear: {}, y_add_clear: {} ", beginning_y_clear, y_clear, y_add_clear);
        assert_eq!(y_clear, (beginning_y_clear + y_add_clear) % message_modulus);

        let mut control = server_key.smart_less(&mut x, &mut y);
        let control_clear = client_key.decrypt(&control);
        println!("control_clear: {}, x_clear: {}, y_clear: {} ", control_clear, x_clear, y_clear);
        assert_eq!(control_clear, if x_clear < y_clear { 1 } else { 0 });
        
        let min = shortint_mux(server_key, &mut control, &mut x, &mut y);
        let min_clear = client_key.decrypt(&min);
        println!("min_clear: {}, x_clear: {}, y_clear: {} ", min_clear, x_clear, y_clear);
        assert_eq!(min_clear, if x_clear < y_clear { x_clear } else { y_clear });
    }
}