1use std::cell::RefCell;
  2
  3mod test {
  4    use itertools::{izip, Itertools};
  5
  6    use crate::{
  7        backend::{ArithmeticOps, ModularOpsU64, Modulus},
  8        bool::{
  9            set_parameter_set, BoolEncoding, BoolEvaluator, BooleanGates, CiphertextModulus,
 10            ClientKey, PublicKey, ServerKeyEvaluationDomain, MP_BOOL_PARAMS, SMALL_MP_BOOL_PARAMS,
 11        },
 12        lwe::{decrypt_lwe, LweSecret},
 13        ntt::NttBackendU64,
 14        pbs::PbsInfo,
 15        random::DefaultSecureRng,
 16        rgsw::RlweSecret,
 17        utils::Stats,
 18        Secret,
 19    };
 20
 21    #[test]
 22    fn mp_noise() {
 23        // set_parameter_set(&SMALL_MP_BOOL_PARAMS);
 24        let mut evaluator = BoolEvaluator::<
 25            Vec<Vec<u64>>,
 26            NttBackendU64,
 27            ModularOpsU64<CiphertextModulus<u64>>,
 28            ModularOpsU64<CiphertextModulus<u64>>,
 29        >::new(SMALL_MP_BOOL_PARAMS);
 30
 31        let parties = 16;
 32
 33        let mut rng = DefaultSecureRng::new();
 34        let mut pk_cr_seed = [0u8; 32];
 35        let mut bk_cr_seed = [0u8; 32];
 36        rng.fill_bytes(&mut pk_cr_seed);
 37        rng.fill_bytes(&mut bk_cr_seed);
 38
 39        let cks = (0..parties)
 40            .into_iter()
 41            .map(|_| evaluator.client_key())
 42            .collect_vec();
 43
 44        // construct ideal rlwe sk for meauring noise
 45        let ideal_client_key = {
 46            let mut ideal_rlwe_sk = vec![0i32; evaluator.parameters().rlwe_n().0];
 47            cks.iter().for_each(|k| {
 48                izip!(ideal_rlwe_sk.iter_mut(), k.sk_rlwe().values()).for_each(|(ideal_i, s_i)| {
 49                    *ideal_i = *ideal_i + s_i;
 50                });
 51            });
 52            let mut ideal_lwe_sk = vec![0i32; evaluator.parameters().lwe_n().0];
 53            cks.iter().for_each(|k| {
 54                izip!(ideal_lwe_sk.iter_mut(), k.sk_lwe().values()).for_each(|(ideal_i, s_i)| {
 55                    *ideal_i = *ideal_i + s_i;
 56                });
 57            });
 58
 59            ClientKey::new(
 60                RlweSecret {
 61                    values: ideal_rlwe_sk,
 62                },
 63                LweSecret {
 64                    values: ideal_lwe_sk,
 65                },
 66            )
 67        };
 68
 69        // round 1
 70        let pk_shares = cks
 71            .iter()
 72            .map(|c| evaluator.multi_party_public_key_share(pk_cr_seed, c))
 73            .collect_vec();
 74
 75        // public key
 76        let pk = PublicKey::<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>>::from(
 77            pk_shares.as_slice(),
 78        );
 79
 80        // round 2
 81        let server_key_shares = cks
 82            .iter()
 83            .map(|c| evaluator.multi_party_server_key_share(bk_cr_seed, &pk.key(), c))
 84            .collect_vec();
 85
 86        let server_key = evaluator.aggregate_multi_party_server_key_shares(&server_key_shares);
 87        let server_key_eval_domain = ServerKeyEvaluationDomain::from(&server_key);
 88
 89        let mut m0 = false;
 90        let mut m1 = true;
 91
 92        let mut c_m0 = evaluator.pk_encrypt(pk.key(), m0);
 93        let mut c_m1 = evaluator.pk_encrypt(pk.key(), m1);
 94
 95        // println!(
 96        //     "Evaluator ciphertext modulus: {:?}",
 97        //     evaluator.parameters().rlwe_q()
 98        // );
 99
100        // {
101        //     // let v = LWE_VEC.take();
102        //     let added = decrypt_lwe(
103        //         &c_m1,
104        //         ideal_client_key.sk_rlwe().values(),
105        //         evaluator.pbs_info().modop_rlweq(),
106        //     );
107        //     println!("Of interest {added}");
108        // }
109
110        let true_el_encoded = evaluator.parameters().rlwe_q().true_el();
111        let false_el_encoded = evaluator.parameters().rlwe_q().false_el();
112
113        let mut stats = Stats::new();
114
115        for _ in 0..1000 {
116            let c_out = evaluator.xor(&c_m0, &c_m1, &server_key_eval_domain);
117
118            // mp decrypt
119            let decryption_shares = cks
120                .iter()
121                .map(|c| evaluator.multi_party_decryption_share(&c_out, c))
122                .collect_vec();
123            let m_out = evaluator.multi_party_decrypt(&decryption_shares, &c_out);
124            let m_expected = (m0 ^ m1);
125            assert_eq!(m_expected, m_out, "Expected {m_expected} but got {m_out}");
126
127            // find noise update
128            {
129                let out = decrypt_lwe(
130                    &c_out,
131                    ideal_client_key.sk_rlwe().values(),
132                    evaluator.pbs_info().modop_rlweq(),
133                );
134
135                let out_want = {
136                    if m_expected == true {
137                        true_el_encoded
138                    } else {
139                        false_el_encoded
140                    }
141                };
142                let diff = evaluator.pbs_info().modop_rlweq().sub(&out, &out_want);
143
144                stats.add_more(&vec![evaluator
145                    .pbs_info()
146                    .rlwe_q()
147                    .map_element_to_i64(&diff)]);
148            }
149
150            m1 = m0;
151            m0 = m_out;
152
153            c_m1 = c_m0;
154            c_m0 = c_out;
155        }
156
157        println!("log2 std dev {}", stats.std_dev().abs().log2());
158    }
159}