modify rgsw/runtime ot use traits

src/bool/evaluator.rs

@@ -25,15 +25,15 @@ use crate::{
         RandomFillUniformInModulus, RandomGaussianElementInModulus,
     },
     rgsw::{
-        decrypt_rlwe, generate_auto_map, public_key_encrypt_rgsw, rgsw_by_rgsw_inplace, rlwe_auto,
-        secret_key_encrypt_rgsw, seeded_auto_key_gen,
+        decrypt_rlwe, generate_auto_map, public_key_encrypt_rgsw, rgsw_by_rgsw_inplace,
+        rgsw_x_rgsw_scratch_rows, rlwe_auto, secret_key_encrypt_rgsw, seeded_auto_key_gen,
+        RgswCiphertext, RgswCiphertextMutRef, RgswCiphertextRef, RuntimeScratchMutRef,
     },
     utils::{
         encode_x_pow_si_with_emebedding_factor, fill_random_ternary_secret_with_hamming_weight,
         generate_prime, mod_exponent, puncture_p_rng, Global, TryConvertFrom1, WithLocal,
     },
-    Decryptor, Encoder, Encryptor, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor, Row,
-    RowEntity, RowMut, Secret,
+    Encoder, Matrix, MatrixEntity, MatrixMut, RowEntity, RowMut,
 };
 
 use super::{
@@ -45,10 +45,7 @@ use super::{
         SeededMultiPartyServerKey, SeededNonInteractiveMultiPartyServerKey,
         SeededSinglePartyServerKey, SinglePartyClientKey,
     },
-    parameters::{
-        BoolParameters, CiphertextModulus, DecompositionCount, DecompostionLogBase,
-        DoubleDecomposerParams,
-    },
+    parameters::{BoolParameters, CiphertextModulus, DecompositionCount, DoubleDecomposerParams},
 };
 
 /// Common reference seed used for Interactive multi-party,
@@ -1146,12 +1143,12 @@ where
 
         // rgsw ciphertext (most expensive part!)
         let rgsw_cts = {
-            let rgsw_by_rgsw_decomposer =
+            let rgsw_x_rgsw_decomposer =
                 parameters.rgsw_rgsw_decomposer::<DefaultDecomposer<M::MatElement>>();
             let rlwe_x_rgsw_decomposer = self.pbs_info().rlwe_rgsw_decomposer();
             let rgsw_x_rgsw_dimension = (
-                rgsw_by_rgsw_decomposer.a().decomposition_count() * 2
-                    + rgsw_by_rgsw_decomposer.b().decomposition_count() * 2,
+                rgsw_x_rgsw_decomposer.a().decomposition_count() * 2
+                    + rgsw_x_rgsw_decomposer.b().decomposition_count() * 2,
                 rlwe_n,
             );
             let rlwe_x_rgsw_dimension = (
@@ -1159,11 +1156,9 @@ where
                     + rlwe_x_rgsw_decomposer.b().decomposition_count() * 2,
                 rlwe_n,
             );
-            let mut rgsw_x_rgsw_scratch_mat = M::zeros(
-                std::cmp::max(
-                    rgsw_by_rgsw_decomposer.a().decomposition_count(),
-                    rgsw_by_rgsw_decomposer.b().decomposition_count(),
-                ) + rlwe_x_rgsw_dimension.0,
+
+            let mut rgsw_x_rgsw_scratch = M::zeros(
+                rgsw_x_rgsw_scratch_rows(rlwe_x_rgsw_decomposer, &rgsw_x_rgsw_decomposer),
                 rlwe_n,
             );
 
@@ -1216,15 +1211,22 @@ where
                                     .for_each(|r| rlweq_nttop.forward(r.as_mut()));
 
                                 rgsw_by_rgsw_inplace(
-                                    &mut rgsw_i,
+                                    &mut RgswCiphertextMutRef::new(
+                                        rgsw_i.as_mut(),
                                         rlwe_x_rgsw_decomposer.a().decomposition_count(),
                                         rlwe_x_rgsw_decomposer.b().decomposition_count(),
-                                    &other_rgsw_i,
-                                    &rgsw_by_rgsw_decomposer,
-                                    &mut rgsw_x_rgsw_scratch_mat,
+                                    ),
+                                    &RgswCiphertextRef::new(
+                                        other_rgsw_i.as_ref(),
+                                        rgsw_x_rgsw_decomposer.a().decomposition_count(),
+                                        rgsw_x_rgsw_decomposer.b().decomposition_count(),
+                                    ),
+                                    rlwe_x_rgsw_decomposer,
+                                    &rgsw_x_rgsw_decomposer,
+                                    &mut RuntimeScratchMutRef::new(rgsw_x_rgsw_scratch.as_mut()),
                                     rlweq_nttop,
                                     rlweq_modop,
-                                )
+                                );
                             });
 
                         rgsw_cts.push(rgsw_i);
@@ -1370,11 +1372,7 @@ where
             };
 
             let mut scratch_rgsw_x_rgsw = M::zeros(
-                std::cmp::max(
-                    rgsw_x_rgsw_decomposer.a().decomposition_count(),
-                    rgsw_x_rgsw_decomposer.b().decomposition_count(),
-                ) + rlwe_x_rgsw_decomposer.a().decomposition_count() * 2
-                    + rlwe_x_rgsw_decomposer.b().decomposition_count() * 2,
+                rgsw_x_rgsw_scratch_rows(&rlwe_x_rgsw_decomposer, &rgsw_x_rgsw_decomposer),
                 self.parameters().rlwe_n().0,
             );
 
@@ -1534,7 +1532,7 @@ where
                             (0..total_users)
                                 .filter(|i| *i != user_id)
                                 .for_each(|other_user_id| {
-                                    let other_rgsw_i = produce_rgsw_ciphertext_from_ni_rgsw(
+                                    let mut other_rgsw_i = produce_rgsw_ciphertext_from_ni_rgsw(
                                         key_shares[other_user_id]
                                             .ni_rgsw_cts_for_self_not_leader_lwe_index(lwe_index),
                                         &ni_rgsw_zero_encs,
@@ -1551,12 +1549,21 @@ where
                                     );
 
                                     rgsw_by_rgsw_inplace(
-                                        &mut rgsw_i,
+                                        &mut RgswCiphertextMutRef::new(
+                                            rgsw_i.as_mut(),
                                             rlwe_x_rgsw_decomposer.a().decomposition_count(),
                                             rlwe_x_rgsw_decomposer.b().decomposition_count(),
-                                        &other_rgsw_i,
+                                        ),
+                                        &RgswCiphertextRef::new(
+                                            other_rgsw_i.as_ref(),
+                                            rgsw_x_rgsw_decomposer.a().decomposition_count(),
+                                            rgsw_x_rgsw_decomposer.b().decomposition_count(),
+                                        ),
+                                        &rlwe_x_rgsw_decomposer,
                                         &rgsw_x_rgsw_decomposer,
-                                        &mut scratch_rgsw_x_rgsw,
+                                        &mut RuntimeScratchMutRef::new(
+                                            scratch_rgsw_x_rgsw.as_mut(),
+                                        ),
                                         nttop,
                                         rlwe_modop,
                                     )
@@ -2096,9 +2103,7 @@ where
         });
 
         let e = DefaultSecureRng::with_local_mut(|rng| {
-            let mut e =
-                RandomGaussianElementInModulus::random(rng, self.pbs_info.parameters.rlwe_q());
-            e
+            RandomGaussianElementInModulus::random(rng, self.pbs_info.parameters.rlwe_q())
         });
         let share = modop.add(&neg_sa, &e);
 

src/bool/keys.rs

@@ -4,7 +4,6 @@ use crate::{
     backend::{ModInit, VectorOps},
     pbs::WithShoupRepr,
     random::{NewWithSeed, RandomFillUniformInModulus},
-    rgsw::RlweSecret,
     utils::{ToShoup, WithLocal},
     Decryptor, Encryptor, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor, RowEntity, RowMut,
 };

src/bool/print_noise.rs

@@ -10,7 +10,10 @@ use crate::{
     lwe::{decrypt_lwe, lwe_key_switch},
     parameters::{BoolParameters, CiphertextModulus},
     random::{DefaultSecureRng, RandomFillUniformInModulus},
-    rgsw::{decrypt_rlwe, rlwe_auto, IsTrivial, RlweCiphertext},
+    rgsw::{
+        decrypt_rlwe, rlwe_auto, rlwe_auto_scratch_rows, RlweCiphertextMutRef, RlweKskRef,
+        RuntimeScratchMutRef,
+    },
     utils::{encode_x_pow_si_with_emebedding_factor, tests::Stats, TryConvertFrom1},
     ArithmeticOps, ClientKey, Decomposer, MatrixEntity, MatrixMut, ModInit, Ntt, NttInit,
     RowEntity, RowMut, VectorOps,
@@ -223,7 +226,8 @@ where
         let br_q = parameters.br_q();
         let g_dlogs = parameters.auto_element_dlogs();
         let auto_decomposer = parameters.auto_decomposer::<D>();
-        let mut scratch_matrix = M::zeros(auto_decomposer.decomposition_count() + 2, rlwe_n);
+        let mut scratch_matrix = M::zeros(rlwe_auto_scratch_rows(&auto_decomposer), rlwe_n);
+        let mut scratch_matrix_ref = RuntimeScratchMutRef::new(scratch_matrix.as_mut());
 
         g_dlogs.iter().for_each(|k| {
             let g_pow_k = if *k == 0 {
@@ -279,19 +283,22 @@ where
             // RLWE auto sends part A, A(X), of RLWE to A(X^{g^k}) and then multiplies it
             // with -s(X^{g^k}) using auto key. Deliberately set RLWE = (0, m(X))
             // (ie. m in part A) to get back RLWE(-m(X^{g^k})s(X^{g^k}))
-            let mut rlwe = RlweCiphertext::<_, DefaultSecureRng>::new_trivial(M::zeros(2, rlwe_n));
-            rlwe.data.get_row_mut(0).copy_from_slice(m.as_ref());
-            rlwe.set_not_trivial();
+            let mut rlwe = M::zeros(2, rlwe_n);
+            rlwe.get_row_mut(0).copy_from_slice(m.as_ref());
 
             rlwe_auto(
-                &mut rlwe,
-                server_key.galois_key_for_auto(*k),
-                &mut scratch_matrix,
+                &mut RlweCiphertextMutRef::new(rlwe.as_mut()),
+                &RlweKskRef::new(
+                    server_key.galois_key_for_auto(*k).as_ref(),
+                    auto_decomposer.decomposition_count(),
+                ),
+                &mut scratch_matrix_ref,
                 &auto_index_map,
                 &auto_sign_map,
                 &rlwe_modop,
                 &rlwe_nttop,
                 &auto_decomposer,
+                false,
             );
 
             // decrypt RLWE(-m(X)s(X^{g^k]}))
@@ -430,7 +437,7 @@ mod tests {
         set_parameter_set(crate::ParameterSelector::NonInteractiveLTE2Party);
         set_common_reference_seed(NonInteractiveMultiPartyCrs::random().seed);
         let parties = 2;
-        let cks = (0..parties).map(|i| gen_client_key()).collect_vec();
+        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();
         let server_key_shares = cks
             .iter()
             .enumerate()

src/pbs.rs

@@ -1,14 +1,16 @@
-use std::{fmt::Display, marker::PhantomData};
+use std::fmt::Display;
 
 use num_traits::{FromPrimitive, One, PrimInt, ToPrimitive, Zero};
 
 use crate::{
     backend::{ArithmeticOps, Modulus, ShoupMatrixFMA, VectorOps},
-    decomposer::Decomposer,
+    decomposer::{Decomposer, RlweDecomposer},
     lwe::lwe_key_switch,
     ntt::Ntt,
-    random::DefaultSecureRng,
-    rgsw::{galois_auto_shoup, rlwe_by_rgsw_shoup, IsTrivial, RlweCiphertext},
+    rgsw::{
+        rlwe_auto_shoup, rlwe_by_rgsw_shoup, RgswCiphertextRef, RlweCiphertextMutRef, RlweKskRef,
+        RuntimeScratchMutRef,
+    },
     Matrix, MatrixEntity, MatrixMut, RowMut,
 };
 pub(crate) trait PbsKey {
@@ -215,7 +217,8 @@ pub(crate) fn pbs<
 /// gk_to_si: [g^0, ..., g^{q/2-1}, -g^0, -g^1, .., -g^{q/2-1}]
 fn blind_rotation<
     Mmut: MatrixMut,
-    D: Decomposer<Element = Mmut::MatElement>,
+    RlweD: RlweDecomposer<Element = Mmut::MatElement>,
+    AutoD: Decomposer<Element = Mmut::MatElement>,
     NttOp: Ntt<Element = Mmut::MatElement>,
     ModOp: ArithmeticOps<Element = Mmut::MatElement> + ShoupMatrixFMA<Mmut::R>,
     MShoup: WithShoupRepr<M = Mmut>,
@@ -228,8 +231,8 @@ fn blind_rotation<
     w: usize,
     q: usize,
     gk_to_si: &[Vec<usize>],
-    rlwe_rgsw_decomposer: &(D, D),
-    auto_decomposer: &D,
+    rlwe_rgsw_decomposer: &RlweD,
+    auto_decomposer: &AutoD,
     ntt_op: &NttOp,
     mod_op: &ModOp,
     parameters: &P,
@@ -239,6 +242,11 @@ fn blind_rotation<
     Mmut::MatElement: Copy + Zero,
 {
     let mut is_trivial = true;
+    let mut scratch_matrix = RuntimeScratchMutRef::new(scratch_matrix.as_mut());
+    let mut rlwe = RlweCiphertextMutRef::new(trivial_rlwe_test_poly.as_mut());
+    let d_a = rlwe_rgsw_decomposer.a().decomposition_count();
+    let d_b = rlwe_rgsw_decomposer.b().decomposition_count();
+    let d_auto = auto_decomposer.decomposition_count();
 
     let q_by_4 = q >> 2;
     let mut count = 0;
@@ -252,10 +260,10 @@ fn blind_rotation<
             // let new = std::time::Instant::now();
             let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
             rlwe_by_rgsw_shoup(
-                trivial_rlwe_test_poly,
-                ct.as_ref(),
-                ct.shoup_repr(),
-                scratch_matrix,
+                &mut rlwe,
+                &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
+                &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
+                &mut scratch_matrix,
                 rlwe_rgsw_decomposer,
                 ntt_op,
                 mod_op,
@@ -271,11 +279,11 @@ fn blind_rotation<
 
             // let now = std::time::Instant::now();
             let auto_key = pbs_key.galois_key_for_auto(v);
-            galois_auto_shoup(
-                trivial_rlwe_test_poly,
-                auto_key.as_ref(),
-                auto_key.shoup_repr(),
-                scratch_matrix,
+            rlwe_auto_shoup(
+                &mut rlwe,
+                &RlweKskRef::new(auto_key.as_ref().as_ref(), d_auto),
+                &RlweKskRef::new(auto_key.shoup_repr().as_ref(), d_auto),
+                &mut scratch_matrix,
                 &auto_map_index,
                 &auto_map_sign,
                 mod_op,
@@ -295,10 +303,10 @@ fn blind_rotation<
         gk_to_si[q_by_4].iter().for_each(|s_index| {
             let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
             rlwe_by_rgsw_shoup(
-                trivial_rlwe_test_poly,
-                ct.as_ref(),
-                ct.shoup_repr(),
-                scratch_matrix,
+                &mut rlwe,
+                &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
+                &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
+                &mut scratch_matrix,
                 rlwe_rgsw_decomposer,
                 ntt_op,
                 mod_op,
@@ -309,11 +317,11 @@ fn blind_rotation<
 
         let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(0);
         let auto_key = pbs_key.galois_key_for_auto(0);
-        galois_auto_shoup(
-            trivial_rlwe_test_poly,
-            auto_key.as_ref(),
-            auto_key.shoup_repr(),
-            scratch_matrix,
+        rlwe_auto_shoup(
+            &mut rlwe,
+            &RlweKskRef::new(auto_key.as_ref().as_ref(), d_auto),
+            &RlweKskRef::new(auto_key.shoup_repr().as_ref(), d_auto),
+            &mut scratch_matrix,
             &auto_map_index,
             &auto_map_sign,
             mod_op,
@@ -331,10 +339,10 @@ fn blind_rotation<
         s_indices.iter().for_each(|s_index| {
             let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
             rlwe_by_rgsw_shoup(
-                trivial_rlwe_test_poly,
-                ct.as_ref(),
-                ct.shoup_repr(),
-                scratch_matrix,
+                &mut rlwe,
+                &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
+                &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
+                &mut scratch_matrix,
                 rlwe_rgsw_decomposer,
                 ntt_op,
                 mod_op,
@@ -347,11 +355,11 @@ fn blind_rotation<
         if gk_to_si[i - 1].len() != 0 || v == w || i == 1 {
             let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(v);
             let auto_key = pbs_key.galois_key_for_auto(v);
-            galois_auto_shoup(
-                trivial_rlwe_test_poly,
-                auto_key.as_ref(),
-                auto_key.shoup_repr(),
-                scratch_matrix,
+            rlwe_auto_shoup(
+                &mut rlwe,
+                &RlweKskRef::new(auto_key.as_ref().as_ref(), d_auto),
+                &RlweKskRef::new(auto_key.shoup_repr().as_ref(), d_auto),
+                &mut scratch_matrix,
                 &auto_map_index,
                 &auto_map_sign,
                 mod_op,
@@ -368,10 +376,10 @@ fn blind_rotation<
     gk_to_si[0].iter().for_each(|s_index| {
         let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
         rlwe_by_rgsw_shoup(
-            trivial_rlwe_test_poly,
-            ct.as_ref(),
-            ct.shoup_repr(),
-            scratch_matrix,
+            &mut rlwe,
+            &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
+            &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
+            &mut scratch_matrix,
             rlwe_rgsw_decomposer,
             ntt_op,
             mod_op,

src/rgsw/keygen.rs

@@ -624,16 +624,15 @@ pub(crate) fn decrypt_rlwe<
     mod_op.elwise_add_mut(m_out.as_mut(), rlwe_ct.get_row_slice(1));
 }
 
-// Measures noise in degree 1 RLWE ciphertext against encoded ideal message
-// encoded_m
-pub(crate) fn measure_noise<
+// Measures maximum noise in degree 1 RLWE ciphertext against message `want_m`
+pub(crate) fn measure_max_noise<
     Mmut: MatrixMut + Matrix,
     ModOp: VectorOps<Element = Mmut::MatElement> + GetModulus<Element = Mmut::MatElement>,
     NttOp: Ntt<Element = Mmut::MatElement>,
     S,
 >(
     rlwe_ct: &Mmut,
-    encoded_m_ideal: &Mmut::R,
+    want_m: &Mmut::R,
     ntt_op: &NttOp,
     mod_op: &ModOp,
     s: &[S],
@@ -645,7 +644,7 @@ where
 {
     let ring_size = s.len();
     assert!(rlwe_ct.dimension() == (2, ring_size));
-    assert!(encoded_m_ideal.as_ref().len() == ring_size);
+    assert!(want_m.as_ref().len() == ring_size);
 
     // -(s * a)
     let q = mod_op.modulus();
@@ -663,11 +662,11 @@ where
     mod_op.elwise_add_mut(m_plus_e.as_mut(), rlwe_ct.get_row_slice(1));
 
     // difference
-    mod_op.elwise_sub_mut(m_plus_e.as_mut(), encoded_m_ideal.as_ref());
+    mod_op.elwise_sub_mut(m_plus_e.as_mut(), want_m.as_ref());
 
     let mut max_diff_bits = f64::MIN;
     m_plus_e.as_ref().iter().for_each(|v| {
-        let bits = (q.map_element_to_i64(v).to_f64().unwrap()).log2();
+        let bits = (q.map_element_to_i64(v).to_f64().unwrap().abs()).log2();
 
         if max_diff_bits < bits {
             max_diff_bits = bits;