arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/bool/print_noise.rs
src/bool/print_noise.rsRaw
  1use std::{fmt::Debug, iter::Sum};
  2
  3use itertools::izip;
  4use num_traits::{FromPrimitive, PrimInt, Zero};
  5use rand_distr::uniform::SampleUniform;
  6
  7use crate::{
  8    backend::{GetModulus, Modulus},
  9    decomposer::{Decomposer, NumInfo, RlweDecomposer},
 10    lwe::{decrypt_lwe, lwe_key_switch},
 11    parameters::{BoolParameters, CiphertextModulus},
 12    random::{DefaultSecureRng, RandomFillUniformInModulus},
 13    rgsw::{
 14        decrypt_rlwe, rlwe_auto, rlwe_auto_scratch_rows, RlweCiphertextMutRef, RlweKskRef,
 15        RuntimeScratchMutRef,
 16    },
 17    utils::{encode_x_pow_si_with_emebedding_factor, tests::Stats, TryConvertFrom1},
 18    ArithmeticOps, ClientKey, MatrixEntity, MatrixMut, ModInit, Ntt, NttInit, RowEntity, RowMut,
 19    VectorOps,
 20};
 21
 22use super::keys::tests::{ideal_sk_lwe, ideal_sk_rlwe};
 23
 24pub(crate) trait CollectRuntimeServerKeyStats {
 25    type M;
 26    /// RGSW ciphertext X^{s[s_index]} in evaluation domain where s the LWE
 27    /// secret
 28    fn rgsw_cts_lwe_si(&self, s_index: usize) -> &Self::M;
 29    /// Auto key in evaluation domain for automorphism g^k. For auto key for
 30    /// automorphism corresponding to -g, set k = 0
 31    fn galois_key_for_auto(&self, k: usize) -> &Self::M;
 32    /// LWE key switching key
 33    fn lwe_ksk(&self) -> &Self::M;
 34}
 35
 36struct ServerKeyStats<T> {
 37    brk_rgsw_cts: (Stats<T>, Stats<T>),
 38    post_1_auto: Stats<T>,
 39    post_lwe_key_switch: Stats<T>,
 40}
 41
 42impl<T: PrimInt + FromPrimitive + Debug + Sum> ServerKeyStats<T>
 43where
 44    T: for<'a> Sum<&'a T>,
 45{
 46    fn new() -> Self {
 47        ServerKeyStats {
 48            brk_rgsw_cts: (Stats::default(), Stats::default()),
 49            post_1_auto: Stats::default(),
 50            post_lwe_key_switch: Stats::default(),
 51        }
 52    }
 53
 54    fn add_noise_brk_rgsw_cts_nsm(&mut self, noise: &[T]) {
 55        self.brk_rgsw_cts.0.add_more(noise);
 56    }
 57
 58    fn add_noise_brk_rgsw_cts_m(&mut self, noise: &[T]) {
 59        self.brk_rgsw_cts.1.add_more(noise);
 60    }
 61
 62    fn add_noise_post_1_auto(&mut self, noise: &[T]) {
 63        self.post_1_auto.add_more(&noise);
 64    }
 65
 66    fn add_noise_post_kwe_key_switch(&mut self, noise: &[T]) {
 67        self.post_lwe_key_switch.add_more(&noise);
 68    }
 69}
 70
 71fn collect_server_key_stats<
 72    M: MatrixEntity + MatrixMut,
 73    D: Decomposer<Element = M::MatElement>,
 74    NttOp: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
 75    ModOp: VectorOps<Element = M::MatElement>
 76        + ArithmeticOps<Element = M::MatElement>
 77        + ModInit<M = CiphertextModulus<M::MatElement>>
 78        + GetModulus<M = CiphertextModulus<M::MatElement>, Element = M::MatElement>,
 79    S: CollectRuntimeServerKeyStats<M = M>,
 80>(
 81    parameters: BoolParameters<M::MatElement>,
 82    client_keys: &[ClientKey],
 83    server_key: &S,
 84) -> ServerKeyStats<i64>
 85where
 86    M::R: RowMut + RowEntity + TryConvertFrom1<[i32], CiphertextModulus<M::MatElement>> + Clone,
 87    M::MatElement: Copy + PrimInt + FromPrimitive + SampleUniform + Zero + Debug + NumInfo,
 88{
 89    let ideal_sk_rlwe = ideal_sk_rlwe(client_keys);
 90    let ideal_sk_lwe = ideal_sk_lwe(client_keys);
 91
 92    let embedding_factor = (2 * parameters.rlwe_n().0) / parameters.br_q();
 93    let rlwe_n = parameters.rlwe_n().0;
 94    let rlwe_q = parameters.rlwe_q();
 95    let lwe_q = parameters.lwe_q();
 96    let rlwe_modop = ModOp::new(rlwe_q.clone());
 97    let rlwe_nttop = NttOp::new(rlwe_q, rlwe_n);
 98    let lwe_modop = ModOp::new(*parameters.lwe_q());
 99
100    let rlwe_x_rgsw_decomposer = parameters.rlwe_rgsw_decomposer::<D>();
101    let (rlwe_x_rgsw_gadget_a, rlwe_x_rgsw_gadget_b) = (
102        rlwe_x_rgsw_decomposer.a().gadget_vector(),
103        rlwe_x_rgsw_decomposer.b().gadget_vector(),
104    );
105
106    let lwe_ks_decomposer = parameters.lwe_decomposer::<D>();
107
108    let mut server_key_stats = ServerKeyStats::new();
109
110    let mut rng = DefaultSecureRng::new();
111
112    // RGSW ciphertext noise
113    // Check noise in RGSW ciphertexts of ideal LWE secret elements
114    {
115        ideal_sk_lwe.iter().enumerate().for_each(|(s_index, s_i)| {
116            let rgsw_ct_i = server_key.rgsw_cts_lwe_si(s_index);
117
118            // X^{s[i]}
119            let m_si = encode_x_pow_si_with_emebedding_factor::<M::R, _>(
120                *s_i,
121                embedding_factor,
122                rlwe_n,
123                rlwe_q,
124            );
125
126            // RLWE'(-sm)
127            let mut neg_s_eval = M::R::try_convert_from(ideal_sk_rlwe.as_slice(), rlwe_q);
128            rlwe_modop.elwise_neg_mut(neg_s_eval.as_mut());
129            rlwe_nttop.forward(neg_s_eval.as_mut());
130
131            for j in 0..rlwe_x_rgsw_decomposer.a().decomposition_count().0 {
132                // RLWE(B^{j} * -s[X]*X^{s_lwe[i]})
133
134                // -s[X]*X^{s_lwe[i]}*B_j
135                let mut m_ideal = m_si.clone();
136                rlwe_nttop.forward(m_ideal.as_mut());
137                rlwe_modop.elwise_mul_mut(m_ideal.as_mut(), neg_s_eval.as_ref());
138                rlwe_nttop.backward(m_ideal.as_mut());
139                rlwe_modop.elwise_scalar_mul_mut(m_ideal.as_mut(), &rlwe_x_rgsw_gadget_a[j]);
140
141                // RLWE(-s*X^{s_lwe[i]}*B_j)
142                let mut rlwe_ct = M::zeros(2, rlwe_n);
143                rlwe_ct
144                    .get_row_mut(0)
145                    .copy_from_slice(rgsw_ct_i.get_row_slice(j));
146                rlwe_ct.get_row_mut(1).copy_from_slice(
147                    rgsw_ct_i.get_row_slice(j + rlwe_x_rgsw_decomposer.a().decomposition_count().0),
148                );
149                // RGSW ciphertexts are in eval domain. We put RLWE ciphertexts back in
150                // coefficient domain
151                rlwe_ct
152                    .iter_rows_mut()
153                    .for_each(|r| rlwe_nttop.backward(r.as_mut()));
154
155                let mut m_back = M::R::zeros(rlwe_n);
156                decrypt_rlwe(
157                    &rlwe_ct,
158                    &ideal_sk_rlwe,
159                    &mut m_back,
160                    &rlwe_nttop,
161                    &rlwe_modop,
162                );
163
164                // diff
165                rlwe_modop.elwise_sub_mut(m_back.as_mut(), m_ideal.as_ref());
166                server_key_stats.add_noise_brk_rgsw_cts_nsm(&Vec::<i64>::try_convert_from(
167                    m_back.as_ref(),
168                    rlwe_q,
169                ));
170            }
171
172            // RLWE'(m)
173            for j in 0..rlwe_x_rgsw_decomposer.b().decomposition_count().0 {
174                // RLWE(B^{j} * X^{s_lwe[i]})
175
176                // X^{s_lwe[i]}*B_j
177                let mut m_ideal = m_si.clone();
178                rlwe_modop.elwise_scalar_mul_mut(m_ideal.as_mut(), &rlwe_x_rgsw_gadget_b[j]);
179
180                // RLWE(X^{s_lwe[i]}*B_j)
181                let mut rlwe_ct = M::zeros(2, rlwe_n);
182                rlwe_ct.get_row_mut(0).copy_from_slice(
183                    rgsw_ct_i.get_row_slice(
184                        j + (2 * rlwe_x_rgsw_decomposer.a().decomposition_count().0),
185                    ),
186                );
187                rlwe_ct
188                    .get_row_mut(1)
189                    .copy_from_slice(rgsw_ct_i.get_row_slice(
190                        j + (2 * rlwe_x_rgsw_decomposer.a().decomposition_count().0)
191                            + rlwe_x_rgsw_decomposer.b().decomposition_count().0,
192                    ));
193                rlwe_ct
194                    .iter_rows_mut()
195                    .for_each(|r| rlwe_nttop.backward(r.as_mut()));
196
197                let mut m_back = M::R::zeros(rlwe_n);
198                decrypt_rlwe(
199                    &rlwe_ct,
200                    &ideal_sk_rlwe,
201                    &mut m_back,
202                    &rlwe_nttop,
203                    &rlwe_modop,
204                );
205
206                // diff
207                rlwe_modop.elwise_sub_mut(m_back.as_mut(), m_ideal.as_ref());
208                server_key_stats.add_noise_brk_rgsw_cts_m(&Vec::<i64>::try_convert_from(
209                    m_back.as_ref(),
210                    rlwe_q,
211                ));
212            }
213        });
214    }
215
216    // Noise in ciphertext after 1 auto
217    // For each auto key g^k. Sample random polynomial m(X) and multiply with
218    // -s(X^{g^k}) using key corresponding to auto g^k. Then check the noise in
219    // resutling RLWE(m(X) * -s(X^{g^k}))
220    {
221        let neg_s = {
222            let mut s = M::R::try_convert_from(ideal_sk_rlwe.as_slice(), rlwe_q);
223            rlwe_modop.elwise_neg_mut(s.as_mut());
224            s
225        };
226        let g = parameters.g();
227        let br_q = parameters.br_q();
228        let g_dlogs = parameters.auto_element_dlogs();
229        let auto_decomposer = parameters.auto_decomposer::<D>();
230        let mut scratch_matrix = M::zeros(rlwe_auto_scratch_rows(&auto_decomposer), rlwe_n);
231        let mut scratch_matrix_ref = RuntimeScratchMutRef::new(scratch_matrix.as_mut());
232
233        g_dlogs.iter().for_each(|k| {
234            let g_pow_k = if *k == 0 {
235                -(g as isize)
236            } else {
237                (g.pow(*k as u32) % br_q) as isize
238            };
239
240            // Send s(X) -> s(X^{g^k})
241            let (auto_index_map, auto_sign_map) = crate::rgsw::generate_auto_map(rlwe_n, g_pow_k);
242            let mut neg_s_g_k = M::R::zeros(rlwe_n);
243            izip!(
244                neg_s.as_ref().iter(),
245                auto_index_map.iter(),
246                auto_sign_map.iter()
247            )
248            .for_each(|(el, to_index, to_sign)| {
249                if !to_sign {
250                    neg_s_g_k.as_mut()[*to_index] = rlwe_modop.neg(el);
251                } else {
252                    neg_s_g_k.as_mut()[*to_index] = *el;
253                }
254            });
255
256            let mut m = M::R::zeros(rlwe_n);
257            RandomFillUniformInModulus::random_fill(&mut rng, rlwe_q, m.as_mut());
258
259            // We want -m(X^{g^k})s(X^{g^k}) after key switch
260            let want_m = {
261                let mut m_g_k_eval = M::R::zeros(rlwe_n);
262                // send m(X) -> m(X^{g^k})
263                izip!(
264                    m.as_ref().iter(),
265                    auto_index_map.iter(),
266                    auto_sign_map.iter()
267                )
268                .for_each(|(el, to_index, to_sign)| {
269                    if !to_sign {
270                        m_g_k_eval.as_mut()[*to_index] = rlwe_modop.neg(el);
271                    } else {
272                        m_g_k_eval.as_mut()[*to_index] = *el;
273                    }
274                });
275
276                rlwe_nttop.forward(m_g_k_eval.as_mut());
277                let mut s_g_k = neg_s_g_k.clone();
278                rlwe_nttop.forward(s_g_k.as_mut());
279                rlwe_modop.elwise_mul_mut(m_g_k_eval.as_mut(), s_g_k.as_ref());
280                rlwe_nttop.backward(m_g_k_eval.as_mut());
281                m_g_k_eval
282            };
283
284            // RLWE auto sends part A, A(X), of RLWE to A(X^{g^k}) and then multiplies it
285            // with -s(X^{g^k}) using auto key. Deliberately set RLWE = (0, m(X))
286            // (ie. m in part A) to get back RLWE(-m(X^{g^k})s(X^{g^k}))
287            let mut rlwe = M::zeros(2, rlwe_n);
288            rlwe.get_row_mut(0).copy_from_slice(m.as_ref());
289
290            rlwe_auto(
291                &mut RlweCiphertextMutRef::new(rlwe.as_mut()),
292                &RlweKskRef::new(
293                    server_key.galois_key_for_auto(*k).as_ref(),
294                    auto_decomposer.decomposition_count().0,
295                ),
296                &mut scratch_matrix_ref,
297                &auto_index_map,
298                &auto_sign_map,
299                &rlwe_modop,
300                &rlwe_nttop,
301                &auto_decomposer,
302                false,
303            );
304
305            // decrypt RLWE(-m(X)s(X^{g^k]}))
306            let mut back_m = M::R::zeros(rlwe_n);
307            decrypt_rlwe(&rlwe, &ideal_sk_rlwe, &mut back_m, &rlwe_nttop, &rlwe_modop);
308
309            // check difference
310            let mut diff = back_m;
311            rlwe_modop.elwise_sub_mut(diff.as_mut(), want_m.as_ref());
312            server_key_stats
313                .add_noise_post_1_auto(&Vec::<i64>::try_convert_from(diff.as_ref(), rlwe_q));
314        });
315
316        // sample random m
317
318        // key switch
319    }
320
321    // LWE Key switch
322    // LWE key switches LWE_in = LWE_{Q_ks,N, s}(m) = (b, a_0, ... a_N) -> LWE_out =
323    // LWE_{Q_{ks}, n, z}(m) = (b', a'_0, ..., a'n)
324    // If LWE_in = (0, a = {a_0, ..., a_N}), then LWE_out = LWE(-a \cdot s_{rlwe})
325    for _ in 0..10 {
326        let mut lwe_in = M::R::zeros(rlwe_n + 1);
327        RandomFillUniformInModulus::random_fill(&mut rng, lwe_q, &mut lwe_in.as_mut()[1..]);
328
329        // Key switch
330        let mut lwe_out = M::R::zeros(parameters.lwe_n().0 + 1);
331        lwe_key_switch(
332            &mut lwe_out,
333            &lwe_in,
334            server_key.lwe_ksk(),
335            &lwe_modop,
336            &lwe_ks_decomposer,
337        );
338
339        // -a \cdot s
340        let mut want_m = M::MatElement::zero();
341        izip!(lwe_in.as_ref().iter().skip(1), ideal_sk_rlwe.iter()).for_each(|(a, b)| {
342            want_m = lwe_modop.add(
343                &want_m,
344                &lwe_modop.mul(a, &lwe_q.map_element_from_i64(*b as i64)),
345            );
346        });
347        want_m = lwe_modop.neg(&want_m);
348
349        // decrypt lwe out
350        let back_m = decrypt_lwe(&lwe_out, &ideal_sk_lwe, &lwe_modop);
351
352        let noise = lwe_modop.sub(&want_m, &back_m);
353        server_key_stats.add_noise_post_kwe_key_switch(&vec![lwe_q.map_element_to_i64(&noise)]);
354    }
355
356    server_key_stats
357    // Auto keys noise
358
359    // Ksk noise
360}
361
362#[cfg(test)]
363mod tests {
364    use itertools::Itertools;
365
366    use super::collect_server_key_stats;
367
368    #[test]
369    #[cfg(feature = "interactive_mp")]
370    fn qwerty() {
371        use crate::{
372            aggregate_public_key_shares, aggregate_server_key_shares,
373            bool::{
374                evaluator::InteractiveMultiPartyCrs,
375                keys::{key_size::KeySize, ServerKeyEvaluationDomain},
376            },
377            gen_client_key, gen_mp_keys_phase2, interactive_multi_party_round1_share,
378            parameters::CiphertextModulus,
379            random::DefaultSecureRng,
380            set_common_reference_seed, set_parameter_set,
381            utils::WithLocal,
382            BoolEvaluator, DefaultDecomposer, ModularOpsU64, Ntt, NttBackendU64,
383        };
384
385        set_parameter_set(crate::ParameterSelector::InteractiveLTE2Party);
386        set_common_reference_seed(InteractiveMultiPartyCrs::random().seed);
387        let parties = 2;
388        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();
389        let pk_shares = cks
390            .iter()
391            .map(|k| interactive_multi_party_round1_share(k))
392            .collect_vec();
393
394        let pk = aggregate_public_key_shares(&pk_shares);
395        let server_key_shares = cks
396            .iter()
397            .enumerate()
398            .map(|(index, k)| gen_mp_keys_phase2(k, index, parties, &pk))
399            .collect_vec();
400
401        // println!("Size: {}", server_key_shares[0].size());
402        let seeded_server_key = aggregate_server_key_shares(&server_key_shares);
403        let server_key_eval =
404            ServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(
405                &seeded_server_key,
406            );
407
408        let parameters = BoolEvaluator::with_local(|e| e.parameters().clone());
409        let server_key_stats = collect_server_key_stats::<
410            _,
411            DefaultDecomposer<u64>,
412            NttBackendU64,
413            ModularOpsU64<CiphertextModulus<u64>>,
414            _,
415        >(parameters, &cks, &server_key_eval);
416
417        println!(
418            "Rgsw nsm std log2 {}",
419            server_key_stats.brk_rgsw_cts.0.std_dev().abs().log2()
420        );
421        println!(
422            "Rgsw m std log2 {}",
423            server_key_stats.brk_rgsw_cts.1.std_dev().abs().log2()
424        );
425        println!(
426            "rlwe post 1 auto std log2 {}",
427            server_key_stats.post_1_auto.std_dev().abs().log2()
428        );
429        println!(
430            "key switching noise rlwe secret s to lwe secret z std log2 {}",
431            server_key_stats.post_lwe_key_switch.std_dev().abs().log2()
432        );
433    }
434
435    #[test]
436    #[cfg(feature = "non_interactive_mp")]
437    fn querty2() {
438        use crate::{
439            aggregate_server_key_shares,
440            bool::{
441                evaluator::NonInteractiveMultiPartyCrs,
442                keys::{key_size::KeySize, NonInteractiveServerKeyEvaluationDomain},
443            },
444            decomposer::DefaultDecomposer,
445            gen_client_key, gen_server_key_share,
446            parameters::CiphertextModulus,
447            random::DefaultSecureRng,
448            set_common_reference_seed, set_parameter_set,
449            utils::WithLocal,
450            BoolEvaluator, ModularOpsU64, NttBackendU64,
451        };
452
453        set_parameter_set(crate::ParameterSelector::NonInteractiveLTE2Party);
454        set_common_reference_seed(NonInteractiveMultiPartyCrs::random().seed);
455        let parties = 2;
456        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();
457        let server_key_shares = cks
458            .iter()
459            .enumerate()
460            .map(|(user_id, k)| gen_server_key_share(user_id, parties, k))
461            .collect_vec();
462
463        let server_key = aggregate_server_key_shares(&server_key_shares);
464
465        let server_key_eval =
466            NonInteractiveServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(
467                &server_key,
468            );
469
470        let parameters = BoolEvaluator::with_local(|e| e.parameters().clone());
471        let server_key_stats = collect_server_key_stats::<
472            _,
473            DefaultDecomposer<u64>,
474            NttBackendU64,
475            ModularOpsU64<CiphertextModulus<u64>>,
476            _,
477        >(parameters, &cks, &server_key_eval);
478
479        println!(
480            "Common reference seeded server key share key size size: {} Bits",
481            server_key_shares[0].size()
482        );
483        println!(
484            "Rgsw nsm std log2 {}",
485            server_key_stats.brk_rgsw_cts.0.std_dev().abs().log2()
486        );
487        println!(
488            "Rgsw m std log2 {}",
489            server_key_stats.brk_rgsw_cts.1.std_dev().abs().log2()
490        );
491        println!(
492            "rlwe post 1 auto std log2 {}",
493            server_key_stats.post_1_auto.std_dev().abs().log2()
494        );
495        println!(
496            "key switching noise rlwe secret s to lwe secret z std log2 {}",
497            server_key_stats.post_lwe_key_switch.std_dev().abs().log2()
498        );
499    }
500
501    #[test]
502    #[cfg(feature = "non_interactive_mp")]
503    fn enc_under_sk_and_key_switch() {
504        use rand::{thread_rng, Rng};
505
506        use crate::{
507            aggregate_server_key_shares,
508            bool::{keys::tests::ideal_sk_rlwe, ni_mp_api::NonInteractiveBatchedFheBools},
509            gen_client_key, gen_server_key_share,
510            rgsw::decrypt_rlwe,
511            set_common_reference_seed, set_parameter_set,
512            utils::{tests::Stats, TryConvertFrom1, WithLocal},
513            BoolEvaluator, Encoder, Encryptor, KeySwitchWithId, ModInit, ModularOpsU64,
514            NttBackendU64, NttInit, ParameterSelector, VectorOps,
515        };
516
517        set_parameter_set(ParameterSelector::NonInteractiveLTE2Party);
518        set_common_reference_seed([2; 32]);
519
520        let parties = 2;
521
522        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();
523
524        let key_shares = cks
525            .iter()
526            .enumerate()
527            .map(|(user_index, ck)| gen_server_key_share(user_index, parties, ck))
528            .collect_vec();
529
530        let seeded_server_key = aggregate_server_key_shares(&key_shares);
531        seeded_server_key.set_server_key();
532
533        let parameters = BoolEvaluator::with_local(|e| e.parameters().clone());
534        let nttop = NttBackendU64::new(parameters.rlwe_q(), parameters.rlwe_n().0);
535        let rlwe_q_modop = ModularOpsU64::new(*parameters.rlwe_q());
536
537        let m = (0..parameters.rlwe_n().0)
538            .map(|_| thread_rng().gen_bool(0.5))
539            .collect_vec();
540        let ct: NonInteractiveBatchedFheBools<_> = cks[0].encrypt(m.as_slice());
541        let ct = ct.key_switch(0);
542
543        let ideal_rlwe_sk = ideal_sk_rlwe(&cks);
544
545        let message = m
546            .iter()
547            .map(|b| parameters.rlwe_q().encode(*b))
548            .collect_vec();
549
550        let mut m_out = vec![0u64; parameters.rlwe_n().0];
551        decrypt_rlwe(
552            &ct.data[0],
553            &ideal_rlwe_sk,
554            &mut m_out,
555            &nttop,
556            &rlwe_q_modop,
557        );
558
559        let mut diff = m_out;
560        rlwe_q_modop.elwise_sub_mut(diff.as_mut_slice(), message.as_ref());
561
562        let mut stats = Stats::new();
563        stats.add_more(&Vec::<i64>::try_convert_from(
564            diff.as_slice(),
565            parameters.rlwe_q(),
566        ));
567        println!("Noise std log2: {}", stats.std_dev().abs().log2());
568    }
569}