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