add more comments

benches/modulus.rs

@@ -49,7 +49,7 @@ fn benchmark_decomposer(c: &mut Criterion) {
                     "decompose",
                     format!(
                         "q={prime}/N={ring_size}/logB={logb}/d={}",
-                        decomposer.decomposition_count()
+                        *decomposer.decomposition_count().as_ref()
                     ),
                 ),
                 |b| {
@@ -57,7 +57,10 @@ fn benchmark_decomposer(c: &mut Criterion) {
                         || {
                             (
                                 a.clone(),
-                                vec![vec![0u64; ring_size]; decomposer.decomposition_count()],
+                                vec![
+                                    vec![0u64; ring_size];
+                                    *decomposer.decomposition_count().as_ref()
+                                ],
                             )
                         },
                         |(r, decomp_r)| (decompose_r(r, decomp_r, &decomposer)),

src/bool/keys.rs

@@ -5,7 +5,7 @@ use crate::{
     pbs::WithShoupRepr,
     random::{NewWithSeed, RandomFillUniformInModulus},
     utils::ToShoup,
-    Matrix, MatrixEntity, MatrixMut, RowEntity, RowMut, SizeInBitsWithLogModulus,
+    Matrix, MatrixEntity, MatrixMut, RowEntity, RowMut,
 };
 
 use super::parameters::{BoolParameters, CiphertextModulus};
@@ -319,8 +319,8 @@ pub struct CommonReferenceSeededInteractiveMultiPartyServerKeyShare
     /// is not the leader RGSW ciphertext is encrypted using RGSW1
     /// decomposer for RGSW0 x RGSW1
     not_self_leader_rgsws: Vec<M>,
-    /// Auto key shares for auto elements [g^{-1}, g, g^2, .., g^{w}] where `w`
-    /// is the window size parameter. Share corresponding to auto element g^{-1}
+    /// Auto key shares for auto elements [-g, g, g^2, .., g^{w}] where `w`
+    /// is the window size parameter. Share corresponding to auto element -g
     /// is stored at key `0` and share corresponding to auto element g^{k} is
     /// stored at key `k`.
     auto_keys: HashMap<usize, M>,
@@ -393,8 +393,8 @@ pub struct SeededInteractiveMultiPartyServerKey {
     /// where `s` is ideal LWE secret key for each LWE secret dimension.
     rgsw_cts: Vec<M>,
     /// Seeded auto keys under ideal RLWE secret for RLWE automorphisms with
-    /// auto elements [g^-1, g, g^2,..., g^{w}]. Auto key corresponidng to
-    /// auto element g^{-1} is stored at key `0` and key corresponding to auto
+    /// auto elements [-g, g, g^2,..., g^{w}]. Auto key corresponidng to
+    /// auto element -g is stored at key `0` and key corresponding to auto
     /// element g^{k} is stored at key `k`
     auto_keys: HashMap<usize, M>,
     /// Seeded LWE key switching key under ideal LWE secret to switch LWE_{q,
@@ -484,7 +484,7 @@ impl SeededSinglePartyServerKey,
 pub(crate) struct ServerKeyEvaluationDomain<M, P, R, N> {
     /// RGSW ciphertext RGSW(X^{s[i]}) for each LWE index in evaluation domain
     rgsw_cts: Vec<M>,
-    /// Auto keys for all auto elements [g^{-1}, g, g^2,..., g^w] in evaluation
+    /// Auto keys for all auto elements [-g, g, g^2,..., g^w] in evaluation
     /// domain
     galois_keys: HashMap<usize, M>,
     /// LWE key switching key to key switch LWE_{q, s}(m) to LWE_{q, z}(m)
@@ -801,7 +801,7 @@ pub(crate) struct NonInteractiveServerKeyEvaluationDomain {
     /// RGSW ciphertexts RGSW(X^{s[i]}) under ideal RLWE secret key in
     /// evaluation domain
     rgsw_cts: Vec<M>,
-    /// Auto keys for all auto elements [g^{-1}, g, g^2, g^w] in evaluation
+    /// Auto keys for all auto elements [-g, g, g^2, g^w] in evaluation
     /// domain
     auto_keys: HashMap<usize, M>,
     /// LWE key switching key to key switch LWE_{q, s}(m) to LWE_{q, z}(m)
@@ -1010,7 +1010,7 @@ pub struct SeededNonInteractiveMultiPartyServerKey {
     ui_to_s_ksks: Vec<M>,
     /// RGSW ciphertexts RGSW(X^{s[i]}) under ideal RLWE secret key
     rgsw_cts: Vec<M>,
-    /// Auto keys for all auto elements [g^{-1}, g, g^2, g^w]
+    /// Auto keys for all auto elements [-g, g, g^2, g^w]
     auto_keys: HashMap<usize, M>,
     /// LWE key switching key to key switch LWE_{q, s}(m) to LWE_{q, z}(m)
     lwe_ksk: M::R,
@@ -1244,8 +1244,8 @@ pub struct CommonReferenceSeededNonInteractiveMultiPartyServerKeyShare
     /// it is stored at index l - 1, where j is self's user_id
     ksk_zero_encs_for_others: Vec<M>,
 
-    /// RLWE auto key shares for auto elements [g^{-1}, g, g^2, g^{w}] where `w`
-    /// is the window size. Auto key share corresponding to auto element g^{-1}
+    /// RLWE auto key shares for auto elements [-g, g, g^2, g^{w}] where `w`
+    /// is the window size. Auto key share corresponding to auto element -g
     /// is stored at key 0 and key share corresponding to auto element g^{k} is
     /// stored at key `k`
     auto_keys_share: HashMap<usize, M>,
@@ -1400,7 +1400,7 @@ impl WithShoupRepr for NormalAndShoup {
 pub(crate) mod key_size {
     use num_traits::{FromPrimitive, PrimInt};
 
-    use crate::{backend::Modulus, decomposer::NumInfo};
+    use crate::{backend::Modulus, decomposer::NumInfo, SizeInBitsWithLogModulus};
 
     use super::*;
 

src/bool/mod.rs

@@ -38,10 +38,13 @@ mod common_mp_enc_dec {
     impl<E> MultiPartyDecryptor<bool, <Mat as Matrix>::R> for super::keys::ClientKey<[u8; 32], E> {
         type DecryptionShare = <Mat as Matrix>::MatElement;
 
+        /// Generate multi-party decryption share for LWE ciphertext `c`
         fn gen_decryption_share(&self, c: &<Mat as Matrix>::R) -> Self::DecryptionShare {
             BoolEvaluator::with_local(|e| e.multi_party_decryption_share(c, self))
         }
 
+        /// Aggregate mult-party decryptions shares of all parties, decrypt LWE
+        /// ciphertext `c`, and return the bool plaintext
         fn aggregate_decryption_shares(
             &self,
             c: &<Mat as Matrix>::R,
@@ -52,6 +55,8 @@ mod common_mp_enc_dec {
     }
 
     impl SampleExtractor<<Mat as Matrix>::R> for Mat {
+        /// Sample extract coefficient at `index` as a LWE ciphertext from RLWE
+        /// ciphertext `Self`
         fn extract(&self, index: usize) -> <Mat as Matrix>::R {
             // input is RLWE ciphertext
             assert!(self.dimension().0 == 2);

src/bool/ni_mp_api.rs

@@ -167,9 +167,26 @@ impl Global for RuntimeServerKey {
     }
 }
 
+/// Batch of bool ciphertexts stored as vector of RLWE ciphertext under user j's
+/// RLWE secret `u_j`
+///
+/// To use the bool ciphertexts in multi-party protocol first key switch the
+/// ciphertexts from u_j to ideal RLWE secret `s` with
+/// `self.key_switch(user_id)` where `user_id` is user j's id. Key switch
+/// returns `BatchedFheBools` that stored key vector of key switched RLWE
+/// ciphertext.
 pub(super) struct NonInteractiveBatchedFheBools<C> {
     data: Vec<C>,
 }
+
+/// Batch of Bool cipphertexts stored as vector of RLWE ciphertexts under the
+/// ideal RLWE secret key `s` of the protocol
+///
+/// Bool ciphertext at `index` can be extracted from the coefficient at `index %
+/// N` of `index / N`th RLWE ciphertext.
+///
+/// To extract bool ciphertext at `index` as LWE ciphertext use
+/// `self.extract(index)`
 pub(super) struct BatchedFheBools<C> {
     pub(in super::super) data: Vec<C>,
 }
@@ -191,6 +208,8 @@ mod impl_enc_dec {
 
     type Mat = Vec<Vec<u64>>;
 
+    // Implement `extract` to extract Bool LWE ciphertext at `index` from
+    // `BatchedFheBools`
     impl<C: MatrixMut<MatElement = u64>> BatchedFheBools<C>
     where
         C::R: RowEntity + RowMut,
@@ -217,6 +236,11 @@ mod impl_enc_dec {
     where
         <M as Matrix>::R: RowMut,
     {
+        /// Derive `NonInteractiveBatchedFheBools` from a vector seeded RLWE
+        /// ciphertexts (Vec<RLWE>, Seed)
+        ///
+        /// Unseed the RLWE ciphertexts and store them as vector RLWE
+        /// ciphertexts in `NonInteractiveBatchedFheBools`
         fn from(value: &(Vec<M::R>, [u8; 32])) -> Self {
             BoolEvaluator::with_local(|e| {
                 let parameters = e.parameters();
@@ -252,17 +276,21 @@ mod impl_enc_dec {
     where
         K: Encryptor<[bool], (Mat, [u8; 32])>,
     {
+        /// Encrypt a vector bool of arbitrary length as vector of unseeded RLWE
+        /// ciphertexts in `NonInteractiveBatchedFheBools`
         fn encrypt(&self, m: &[bool]) -> NonInteractiveBatchedFheBools<Mat> {
             NonInteractiveBatchedFheBools::from(&K::encrypt(&self, m))
         }
     }
 
-    impl<K> Encryptor<[bool], (Mat, [u8; 32])> for K
+    impl<K> Encryptor<[bool], (Vec<<Mat as Matrix>::R>, [u8; 32])> for K
     where
         K: NonInteractiveMultiPartyClientKey,
         <Mat as Matrix>::R:
             TryConvertFrom1<[K::Element], CiphertextModulus<<Mat as Matrix>::MatElement>>,
     {
+        /// Encrypt a vector of bool of arbitrary length as vector of seeded
+        /// RLWE ciphertexts and returns (Vec<RLWE>, Seed)
         fn encrypt(&self, m: &[bool]) -> (Mat, [u8; 32]) {
             BoolEvaluator::with_local(|e| {
                 DefaultSecureRng::with_local_mut(|rng| {
@@ -319,8 +347,13 @@ mod impl_enc_dec {
     }
 
     impl KeySwitchWithId<Mat> for Mat {
+        /// Key switch RLWE ciphertext `Self` from user j's RLWE secret u_j
+        /// to ideal RLWE secret `s` of non-interactive multi-party protocol.
+        ///
+        /// - user_id: user j's user_id in the protocol
         fn key_switch(&self, user_id: usize) -> Mat {
             BoolEvaluator::with_local(|e| {
+                assert!(self.dimension() == (2, e.parameters().rlwe_n().0));
                 let server_key = BOOL_SERVER_KEY.get().unwrap();
                 let ksk = server_key.ui_to_s_ksk(user_id);
                 let decomposer = e.ni_ui_to_s_ks_decomposer().as_ref().unwrap();
@@ -342,6 +375,14 @@ mod impl_enc_dec {
     where
         C: KeySwitchWithId<C>,
     {
+        /// Key switch `Self`'s vector of RLWE ciphertexts from user j's RLWE
+        /// secret u_j to ideal RLWE secret `s` of non-interactive
+        /// multi-party protocol.
+        ///
+        /// Returns vector of key switched RLWE ciphertext as `BatchedFheBools`
+        /// which can then be used to extract individual Bool LWE ciphertexts.
+        ///
+        /// - user_id: user j's user_id in the protocol
         fn key_switch(&self, user_id: usize) -> BatchedFheBools<C> {
             let data = self
                 .data
@@ -355,24 +396,15 @@ mod impl_enc_dec {
 
 #[cfg(test)]
 mod tests {
-    use itertools::{izip, Itertools};
-    use num_traits::{FromPrimitive, PrimInt, ToPrimitive, Zero};
-    use rand::{thread_rng, Rng, RngCore};
+    use itertools::Itertools;
+    use rand::{thread_rng, RngCore};
 
     use crate::{
-        backend::{GetModulus, Modulus},
         bool::{
             evaluator::BooleanGates,
-            keys::{
-                tests::{ideal_sk_rlwe, measure_noise_lwe},
-                SinglePartyClientKey,
-            },
+            keys::tests::{ideal_sk_rlwe, measure_noise_lwe},
         },
-        lwe::decrypt_lwe,
-        rgsw::decrypt_rlwe,
-        utils::{tests::Stats, TryConvertFrom1},
-        ArithmeticOps, Encoder, Encryptor, KeySwitchWithId, ModInit, MultiPartyDecryptor, NttInit,
-        Row, VectorOps,
+        Encoder, Encryptor, KeySwitchWithId, MultiPartyDecryptor,
     };
 
     use super::*;

src/bool/parameters.rs

@@ -144,7 +144,7 @@ pub struct BoolParameters {
     /// and must be supplied only for non-interactive multi-party
     non_interactive_ui_to_s_key_switch_decomposer:
         Option<(DecompostionLogBase, DecompositionCount)>,
-    /// Group generator for Z^*_{2N}
+    /// Group generator for Z^*_{br_q}
     g: usize,
     /// Window size parameter for LMKC++ blind rotation
     w: usize,
@@ -353,14 +353,14 @@ impl BoolParameters {
 }
 
 #[derive(Clone, Copy, PartialEq)]
-pub(crate) struct DecompostionLogBase(pub(crate) usize);
+pub struct DecompostionLogBase(pub(crate) usize);
 impl AsRef<usize> for DecompostionLogBase {
     fn as_ref(&self) -> &usize {
         &self.0
     }
 }
 #[derive(Clone, Copy, PartialEq)]
-pub(crate) struct DecompositionCount(pub(crate) usize);
+pub struct DecompositionCount(pub(crate) usize);
 impl AsRef<usize> for DecompositionCount {
     fn as_ref(&self) -> &usize {
         &self.0

src/bool/print_noise.rs

@@ -393,7 +393,7 @@ mod tests {
             .map(|(index, k)| gen_mp_keys_phase2(k, index, parties, &pk))
             .collect_vec();
 
-        println!("Size: {}", server_key_shares[0].size());
+        // println!("Size: {}", server_key_shares[0].size());
         let seeded_server_key = aggregate_server_key_shares(&server_key_shares);
         let server_key_eval =
             ServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(
@@ -452,7 +452,7 @@ mod tests {
             .enumerate()
             .map(|(user_id, k)| gen_server_key_share(user_id, parties, k))
             .collect_vec();
-        println!("Size: {}", server_key_shares[0].size());
+
         let server_key = aggregate_server_key_shares(&server_key_shares);
 
         let server_key_eval =
@@ -469,6 +469,10 @@ mod tests {
             _,
         >(parameters, &cks, &server_key_eval);
 
+        println!(
+            "Common reference seeded server key share key size size: {} Bits",
+            server_key_shares[0].size()
+        );
         println!(
             "Rgsw nsm std log2 {}",
             server_key_stats.brk_rgsw_cts.0.std_dev().abs().log2()

src/pbs.rs

@@ -32,48 +32,68 @@ pub(crate) trait WithShoupRepr: AsRef {
 }
 
 pub(crate) trait PbsInfo {
+    /// Type of Matrix
     type M: Matrix;
+    /// Type of Ciphertext modulus
     type Modulus: Modulus<Element = <Self::M as Matrix>::MatElement>;
+    /// Type of Ntt Operator for Ring polynomials
     type NttOp: Ntt<Element = <Self::M as Matrix>::MatElement>;
+    /// Type of Signed Decomposer
     type D: Decomposer<Element = <Self::M as Matrix>::MatElement>;
 
-    // Although both types have same bounds, they can be different types. For ex,
-    // type RlweModOp may only support native modulus, where LweModOp may only
-    // support prime modulus, etc.
+    // Although both `RlweModOp` and `LweModOp` types have same bounds, they can be
+    // different types. For ex, type RlweModOp may only support native modulus,
+    // where LweModOp may only support prime modulus, etc.
+
+    /// Type of RLWE Modulus Operator
     type RlweModOp: ArithmeticOps<Element = <Self::M as Matrix>::MatElement>
         + ShoupMatrixFMA<<Self::M as Matrix>::R>;
+    /// Type of LWE Modulus Operator
     type LweModOp: VectorOps<Element = <Self::M as Matrix>::MatElement>
         + ArithmeticOps<Element = <Self::M as Matrix>::MatElement>;
 
+    /// RLWE ciphertext modulus
     fn rlwe_q(&self) -> &Self::Modulus;
+    /// LWE ciphertext modulus
     fn lwe_q(&self) -> &Self::Modulus;
+    /// Blind rotation modulus. It is the modulus to which we switch for blind
+    /// rotation. Since blind rotation decrypts LWE ciphetext in the exponent of
+    /// ring polynmial (which is a ring mod 2N), `br_q <= 2N`
     fn br_q(&self) -> usize;
+    /// Ring polynomial size `N`
     fn rlwe_n(&self) -> usize;
+    /// LWE dimension `n`
     fn lwe_n(&self) -> usize;
     /// Embedding fator for ring X^{q}+1 inside
     fn embedding_factor(&self) -> usize;
-    /// Window size
+    /// Window size parameter LKMC++ blind rotaiton
     fn w(&self) -> usize;
-    /// generator g
+    /// generator `g` for group Z^*_{br_q}
     fn g(&self) -> isize;
-    /// Decomposers
+    /// LWE key switching decomposer
     fn lwe_decomposer(&self) -> &Self::D;
+    /// RLWE x RGSW decoposer
     fn rlwe_rgsw_decomposer(&self) -> &(Self::D, Self::D);
+    /// RLWE auto decomposer
     fn auto_decomposer(&self) -> &Self::D;
 
-    /// Modulus operators
+    /// LWE modulus operator
     fn modop_lweq(&self) -> &Self::LweModOp;
+    /// RLWE modulus operator
     fn modop_rlweq(&self) -> &Self::RlweModOp;
 
     /// Ntt operators
     fn nttop_rlweq(&self) -> &Self::NttOp;
 
-    /// Maps a \in Z^*_{q} to discrete log k, with generator g (i.e. g^k =
-    /// a). Returned vector is of size q that stores dlog of a at `vec[a]`.
-    /// For any a, if k is s.t. a = g^{k}, then k is expressed as k. If k is s.t
-    /// a = -g^{k}, then k is expressed as k=k+q/4
+    /// Maps a \in Z^*_{br_q} to discrete log k, with generator g (i.e. g^k =
+    /// a). Returned vector is of size q that stores dlog of `a` at `vec[a]`.
+    ///
+    /// For any `a`, if k is s.t. `a = g^{k} % br_q`, then `k` is expressed as
+    /// k. If `k` is s.t `a = -g^{k} % br_q`, then `k` is expressed as
+    /// k=k+q/4
     fn g_k_dlog_map(&self) -> &[usize];
-    /// Returns auto map and index vector for g^k. For -g use k == 0.
+    /// Returns auto map and index vector for auto element g^k. For auto element
+    /// -g set k = 0.
     fn rlwe_auto_map(&self, k: usize) -> &(Vec<usize>, Vec<bool>);
 }
 
@@ -123,7 +143,7 @@ pub(crate) fn pbs<
     );
     // println!("Time: {:?}", now.elapsed());
 
-    // odd mowdown Q_ks -> q
+    // odd moddown Q_ks -> q
     let g_k_dlog_map = pbs_info.g_k_dlog_map();
     let mut g_k_si = vec![vec![]; br_q >> 1];
     scratch_lwe_vec
@@ -214,7 +234,9 @@ pub(crate) fn pbs<
 
 /// LMKCY+ Blind rotation
 ///
-/// gk_to_si: [g^0, ..., g^{q/2-1}, -g^0, -g^1, .., -g^{q/2-1}]
+/// - gk_to_si: Contains LWE secret index `i` in array of secret indices at k^th
+///   index if a_i = g^k if k < q/4 or a_i = -g^k if k > q/4. [g^0, ...,
+///   g^{q/2-1}, -g^0, -g^1, .., -g^{q/2-1}]
 fn blind_rotation<
     Mmut: MatrixMut,
     RlweD: RlweDecomposer<Element = Mmut::MatElement>,

src/shortint/enc_dec.rs

@@ -27,6 +27,9 @@ impl FheUint8 {
     }
 }
 
+/// Stored a batch of Fhe Uint8 ciphertext as collection of RLWE ciphertexts
+///
+/// To extract Fhe Uint8 ciphertext at `index` call `self.extract(index)`
 pub struct BatchedFheUint8<C> {
     data: Vec<C>,
 }
@@ -35,6 +38,13 @@ impl SampleExtractor> for BatchedFheUint8
 where
     C: SampleExtractor<R>,
 {
+    /// Extract Fhe Uint8 ciphertext at `index`
+    ///
+    /// `Self` stores batch of Fhe uint8 ciphertext as vector of RLWE
+    /// ciphertexts. Since Fhe uint8 ciphertext is collection of 8 bool
+    /// ciphertexts, Fhe uint8 ciphertext at index `i` is stored in coefficients
+    /// `i*8...(i+1)*8`. To extract Fhe uint8 at index `i`, sample extract bool
+    /// ciphertext at indices `[i*8, ..., (i+1)*8)`
     fn extract(&self, index: usize) -> FheUint8<R> {
         BoolEvaluator::with_local(|e| {
             let ring_size = e.parameters().rlwe_n().0;
@@ -58,6 +68,8 @@ impl> From<&SeededBatchedFheUint8<
 where
     <M as Matrix>::R: RowMut,
 {
+    /// Unseeds collection of seeded RLWE ciphertext in SeededBatchedFheUint8
+    /// and returns as `Self`
     fn from(value: &SeededBatchedFheUint8<M::R, [u8; 32]>) -> Self {
         BoolEvaluator::with_local(|e| {
             let parameters = e.parameters();
@@ -104,6 +116,7 @@ impl Encryptor<[u8], SeededBatchedFheUint8> for K
 where
     K: Encryptor<[bool], (Vec<C>, S)>,
 {
+    /// Encrypt a slice of u8s of arbitray length as `SeededBatchedFheUint8`
     fn encrypt(&self, m: &[u8]) -> SeededBatchedFheUint8<C, S> {
         // convert vector of u8s to vector bools
         let m = m
@@ -138,6 +151,10 @@ impl KeySwitchWithId> for BatchedFheUint8
 where
     C: KeySwitchWithId<C>,
 {
+    /// Key switching collection of RLWE ciphertexts in `BatchedFheUint8` from
+    /// user j's RLWE secret u_j to ideal RLWE secret key `s` of the protocol.
+    ///
+    /// - user_id: user id of user j
     fn key_switch(&self, user_id: usize) -> BatchedFheUint8<C> {
         let data = self
             .data