impl user facing Encryptor, Decrytor, KeySwitch, SampleExtract for FheBools

src/bool/mod.rs

@@ -229,7 +229,7 @@ mod impl_bool_frontend {
 mod common_mp_enc_dec {
     use itertools::Itertools;
 
-    use super::BoolEvaluator;
+    use super::{impl_bool_frontend::FheBool, BoolEvaluator};
     use crate::{
         pbs::{sample_extract, PbsInfo},
         utils::WithLocal,
@@ -238,6 +238,84 @@ mod common_mp_enc_dec {
 
     type Mat = Vec<Vec<u64>>;
 
+    /// `Self` stores batch of boolean ciphertexts as collection of `unseeded RLWE ciphertexts` encrypted under the ideal RLWE secret key of the protocol.
+    ///
+    /// Bool ciphertext at `index` can be extracted from the coefficient at `index %
+    /// N` of `index / N`th RLWE ciphertext.
+    pub(crate) struct BatchedFheBools<C> {
+        data: Vec<C>,
+        count: usize,
+    }
+
+    impl<C> BatchedFheBools<C> {
+        pub(crate) fn new(data: Vec<C>, count: usize) -> Self {
+            Self { data, count }
+        }
+
+        pub(crate) fn data(&self) -> &[C] {
+            &self.data
+        }
+
+        pub(crate) fn count(&self) -> usize {
+            self.count
+        }
+    }
+
+    impl<M: Matrix> SampleExtractor<FheBool<M::R>> for BatchedFheBools<M>
+    where
+        M: SampleExtractor<M::R>,
+    {
+        fn extract_all(&self) -> Vec<FheBool<M::R>> {
+            if self.data.len() > 0 {
+                let ring_size = self.data[0].dimension().0;
+
+                let index = 0;
+                let mut out = Vec::with_capacity(self.count);
+
+                while index < self.count {
+                    let row = index % ring_size;
+                    let col = index / ring_size;
+                    out.push(FheBool {
+                        data: SampleExtractor::extract_at(&self.data[col], row),
+                    });
+                }
+
+                out
+            } else {
+                vec![]
+            }
+        }
+        fn extract_at(&self, index: usize) -> FheBool<M::R> {
+            assert!(self.count > index);
+
+            let ring_size = self.data[0].dimension().0;
+            let row = index % ring_size;
+            let col = index / ring_size;
+
+            FheBool {
+                data: SampleExtractor::extract_at(&self.data[col], row),
+            }
+        }
+
+        fn extract_many(&self, how_many: usize) -> Vec<FheBool<M::R>> {
+            assert!(self.count >= how_many);
+            let ring_size = self.data[0].dimension().0;
+
+            let index = 0;
+            let mut out = Vec::with_capacity(self.count);
+
+            while index < how_many {
+                let row = index % ring_size;
+                let col = index / ring_size;
+                out.push(FheBool {
+                    data: SampleExtractor::extract_at(&self.data[col], row),
+                });
+            }
+
+            out
+        }
+    }
+
     impl SampleExtractor<<Mat as Matrix>::R> for Mat {
         /// Sample extract coefficient at `index` as a LWE ciphertext from RLWE
         /// ciphertext `Self`

src/bool/ni_mp_api.rs

@@ -179,38 +179,37 @@ impl Global for RuntimeServerKey {
     }
 }
 
+/// `Self::data` stores collection of seeded RLWE ciphertexts encrypted unser user j's RLWE secret `u_j`.
+pub(crate) struct NonInteractiveSeededFheBools<C, S> {
+    data: Vec<C>,
+    seed: S,
+    count: usize,
+}
+
 /// 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
+/// returns `BatchedFheBools` which stores vector of key switched RLWE
 /// ciphertext.
 pub 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 struct BatchedFheBools<C> {
-    pub(in super::super) data: Vec<C>,
+    count: usize,
 }
 
 /// Non interactive multi-party specfic encryptor decryptor routines
 mod impl_enc_dec {
     use crate::{
-        bool::{evaluator::BoolEncoding, keys::NonInteractiveMultiPartyClientKey},
+        bool::{
+            common_mp_enc_dec::BatchedFheBools, evaluator::BoolEncoding,
+            keys::NonInteractiveMultiPartyClientKey,
+        },
         multi_party::{
             multi_party_aggregate_decryption_shares_and_decrypt, multi_party_decryption_share,
         },
-        pbs::{sample_extract, PbsInfo, WithShoupRepr},
+        pbs::{PbsInfo, WithShoupRepr},
         random::{NewWithSeed, RandomFillUniformInModulus},
         rgsw::{rlwe_key_switch, seeded_secret_key_encrypt_rlwe},
         utils::TryConvertFrom1,
@@ -224,82 +223,23 @@ 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,
-    {
-        pub fn extract(&self, index: usize) -> C::R {
-            BoolEvaluator::with_local(|e| {
-                let ring_size = e.parameters().rlwe_n().0;
-                let ct_index = index / ring_size;
-                let coeff_index = index % ring_size;
-                let mut lwe_out = C::R::zeros(e.parameters().rlwe_n().0 + 1);
-                sample_extract(
-                    &mut lwe_out,
-                    &self.data[ct_index],
-                    e.pbs_info().modop_rlweq(),
-                    coeff_index,
-                );
-                lwe_out
-            })
-        }
-    }
-
-    impl<M: MatrixEntity + MatrixMut<MatElement = u64>> From<&(Vec<M::R>, [u8; 32])>
-        for NonInteractiveBatchedFheBools<M>
-    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();
-                let ring_size = parameters.rlwe_n().0;
-                let rlwe_q = parameters.rlwe_q();
-
-                let mut prng = DefaultSecureRng::new_seeded(value.1);
-                let rlwes = value
-                    .0
-                    .iter()
-                    .map(|partb| {
-                        let mut rlwe = M::zeros(2, ring_size);
-
-                        // sample A
-                        RandomFillUniformInModulus::random_fill(
-                            &mut prng,
-                            rlwe_q,
-                            rlwe.get_row_mut(0),
-                        );
-
-                        // Copy over B
-                        rlwe.get_row_mut(1).copy_from_slice(partb.as_ref());
-
-                        rlwe
-                    })
-                    .collect_vec();
-                Self { data: rlwes }
-            })
+    impl<C, S> NonInteractiveSeededFheBools<C, S> {
+        /// Unseed `Self`'s collection of RLWE ciphertexts into `NonInteractiveBatchedFheBools`
+        pub fn unseed<M>(&self) -> NonInteractiveBatchedFheBools<M>
+        where
+            NonInteractiveBatchedFheBools<M>: for<'a> From<&'a Self>,
+        {
+            NonInteractiveBatchedFheBools::from(self)
         }
     }
 
-    impl<K> Encryptor<[bool], NonInteractiveBatchedFheBools<Mat>> for K
-    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<C, S> From<NonInteractiveSeededFheBools<C, S>> for (Vec<C>, S) {
+        fn from(value: NonInteractiveSeededFheBools<C, S>) -> Self {
+            (value.data, value.seed)
         }
     }
 
-    impl<K> Encryptor<[bool], (Vec<<Mat as Matrix>::R>, [u8; 32])> for K
+    impl<K> Encryptor<[bool], NonInteractiveSeededFheBools<<Mat as Matrix>::R, [u8; 32]>> for K
     where
         K: NonInteractiveMultiPartyClientKey,
         <Mat as Matrix>::R:
@@ -307,7 +247,10 @@ mod impl_enc_dec {
     {
         /// 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]) {
+        fn encrypt(
+            &self,
+            m: &[bool],
+        ) -> NonInteractiveSeededFheBools<<Mat as Matrix>::R, [u8; 32]> {
             BoolEvaluator::with_local(|e| {
                 DefaultSecureRng::with_local_mut(|rng| {
                     let parameters = e.parameters();
@@ -356,46 +299,57 @@ mod impl_enc_dec {
                         })
                         .collect_vec();
 
-                    (rlwes, seed)
+                    NonInteractiveSeededFheBools {
+                        data: rlwes,
+                        seed,
+                        count: m.len(),
+                    }
                 })
             })
         }
     }
 
-    impl<K> MultiPartyDecryptor<bool, <Mat as Matrix>::R> for K
+    impl<M: MatrixEntity + MatrixMut<MatElement = u64>>
+        From<&NonInteractiveSeededFheBools<<Mat as Matrix>::R, [u8; 32]>>
+        for NonInteractiveBatchedFheBools<M>
     where
-        K: NonInteractiveMultiPartyClientKey,
-        <Mat as Matrix>::R:
-            TryConvertFrom1<[K::Element], CiphertextModulus<<Mat as Matrix>::MatElement>>,
+        <M as Matrix>::R: RowMut,
     {
-        type DecryptionShare = <Mat as Matrix>::MatElement;
-
-        fn gen_decryption_share(&self, c: &<Mat as Matrix>::R) -> Self::DecryptionShare {
+        /// 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: &NonInteractiveSeededFheBools<<Mat as Matrix>::R, [u8; 32]>) -> Self {
             BoolEvaluator::with_local(|e| {
-                DefaultSecureRng::with_local_mut(|rng| {
-                    multi_party_decryption_share(
-                        c,
-                        self.sk_rlwe().as_slice(),
-                        e.pbs_info().modop_rlweq(),
-                        rng,
-                    )
-                })
-            })
-        }
+                let parameters = e.parameters();
+                let ring_size = parameters.rlwe_n().0;
+                let rlwe_q = parameters.rlwe_q();
 
-        fn aggregate_decryption_shares(
-            &self,
-            c: &<Mat as Matrix>::R,
-            shares: &[Self::DecryptionShare],
-        ) -> bool {
-            BoolEvaluator::with_local(|e| {
-                let noisy_m = multi_party_aggregate_decryption_shares_and_decrypt(
-                    c,
-                    shares,
-                    e.pbs_info().modop_rlweq(),
-                );
+                let mut prng = DefaultSecureRng::new_seeded(value.seed);
+                let rlwes = value
+                    .data
+                    .iter()
+                    .map(|partb| {
+                        let mut rlwe = M::zeros(2, ring_size);
 
-                e.pbs_info().rlwe_q().decode(noisy_m)
+                        // sample A
+                        RandomFillUniformInModulus::random_fill(
+                            &mut prng,
+                            rlwe_q,
+                            rlwe.get_row_mut(0),
+                        );
+
+                        // Copy over B
+                        rlwe.get_row_mut(1).copy_from_slice(partb.as_ref());
+
+                        rlwe
+                    })
+                    .collect_vec();
+                Self {
+                    data: rlwes,
+                    count: value.count,
+                }
             })
         }
     }
@@ -443,7 +397,45 @@ mod impl_enc_dec {
                 .iter()
                 .map(|c| c.key_switch(user_id))
                 .collect_vec();
-            BatchedFheBools { data }
+            BatchedFheBools::new(data, self.count)
+        }
+    }
+
+    impl<K> MultiPartyDecryptor<bool, <Mat as Matrix>::R> for K
+    where
+        K: NonInteractiveMultiPartyClientKey,
+        <Mat as Matrix>::R:
+            TryConvertFrom1<[K::Element], CiphertextModulus<<Mat as Matrix>::MatElement>>,
+    {
+        type DecryptionShare = <Mat as Matrix>::MatElement;
+
+        fn gen_decryption_share(&self, c: &<Mat as Matrix>::R) -> Self::DecryptionShare {
+            BoolEvaluator::with_local(|e| {
+                DefaultSecureRng::with_local_mut(|rng| {
+                    multi_party_decryption_share(
+                        c,
+                        self.sk_rlwe().as_slice(),
+                        e.pbs_info().modop_rlweq(),
+                        rng,
+                    )
+                })
+            })
+        }
+
+        fn aggregate_decryption_shares(
+            &self,
+            c: &<Mat as Matrix>::R,
+            shares: &[Self::DecryptionShare],
+        ) -> bool {
+            BoolEvaluator::with_local(|e| {
+                let noisy_m = multi_party_aggregate_decryption_shares_and_decrypt(
+                    c,
+                    shares,
+                    e.pbs_info().modop_rlweq(),
+                );
+
+                e.pbs_info().rlwe_q().decode(noisy_m)
+            })
         }
     }
 }

src/bool/print_noise.rs

@@ -615,7 +615,6 @@ mod tests {
                     NonInteractiveServerKeyEvaluationDomain,
                 },
                 print_noise::collect_server_key_stats,
-                NonInteractiveBatchedFheBools,
             },
             gen_client_key, gen_server_key_share,
             parameters::CiphertextModulus,
@@ -624,7 +623,7 @@ mod tests {
             utils::{tests::Stats, Global, WithLocal},
             BoolEvaluator, BooleanGates, DefaultDecomposer, Encoder, Encryptor, KeySwitchWithId,
             ModInit, ModularOpsU64, MultiPartyDecryptor, NttBackendU64, ParameterSelector,
-            RuntimeServerKey,
+            RuntimeServerKey, SampleExtractor,
         };
 
         set_parameter_set(ParameterSelector::NonInteractiveLTE8Party);
@@ -654,14 +653,20 @@ mod tests {
         let mut m1 = true;
 
         let mut ct0 = {
-            let ct: NonInteractiveBatchedFheBools<_> = cks[0].encrypt(vec![m0].as_slice());
-            let ct = ct.key_switch(0);
-            ct.extract(0)
+            cks[0]
+                .encrypt(vec![m0].as_slice())
+                .unseed::<Vec<Vec<u64>>>()
+                .key_switch(0)
+                .extract_at(0)
+                .data
         };
         let mut ct1 = {
-            let ct: NonInteractiveBatchedFheBools<_> = cks[1].encrypt(vec![m1].as_slice());
-            let ct = ct.key_switch(1);
-            ct.extract(0)
+            cks[1]
+                .encrypt(vec![m1].as_slice())
+                .unseed::<Vec<Vec<u64>>>()
+                .key_switch(1)
+                .extract_at(0)
+                .data
         };
 
         let mut stats = Stats::new();
@@ -830,7 +835,7 @@ mod tests {
 
         use crate::{
             aggregate_server_key_shares,
-            bool::{keys::tests::ideal_sk_rlwe, ni_mp_api::NonInteractiveBatchedFheBools},
+            bool::keys::tests::ideal_sk_rlwe,
             gen_client_key, gen_server_key_share,
             rgsw::decrypt_rlwe,
             set_common_reference_seed, set_parameter_set,
@@ -862,8 +867,11 @@ mod tests {
         let m = (0..parameters.rlwe_n().0)
             .map(|_| thread_rng().gen_bool(0.5))
             .collect_vec();
-        let ct: NonInteractiveBatchedFheBools<_> = cks[0].encrypt(m.as_slice());
-        let ct = ct.key_switch(0);
+        let ct = cks[0]
+            .encrypt(m.as_slice())
+            .unseed::<Vec<Vec<u64>>>()
+            .key_switch(0);
+        assert!(ct.data().len() == 1);
 
         let ideal_rlwe_sk = ideal_sk_rlwe(&cks);
 
@@ -874,7 +882,7 @@ mod tests {
 
         let mut m_out = vec![0u64; parameters.rlwe_n().0];
         decrypt_rlwe(
-            &ct.data[0],
+            &ct.data()[0],
             &ideal_rlwe_sk,
             &mut m_out,
             &nttop,

src/lib.rs

@@ -28,6 +28,7 @@ pub trait Matrix: AsRef<[Self::R]> {
     type MatElement;
     type R: Row<Element = Self::MatElement>;
 
+    /// (Rows, Cols)
     fn dimension(&self) -> (usize, usize);
 
     fn get_row(&self, row_idx: usize) -> impl Iterator<Item = &Self::MatElement> {

src/shortint/enc_dec.rs

@@ -70,7 +70,7 @@ impl> From<&SeededBatchedFheUint8<
 where
     <M as Matrix>::R: RowMut,
 {
-    /// Unseeds collection of seeded RLWE ciphertext in SeededBatchedFheUint8
+    /// 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| {
@@ -237,9 +237,10 @@ pub struct SeededBatchedFheUint8 {
     count: usize,
 }
 
+#[cfg(feature = "non_interactive_mp")]
 impl<K, C, S> Encryptor<[u8], SeededBatchedFheUint8<C, S>> for K
 where
-    K: Encryptor<[bool], (Vec<C>, S)>,
+    K: Encryptor<[bool], crate::NonInteractiveSeededFheBools<C, S>>,
 {
     /// Encrypt a slice of u8s of arbitray length packed into collection of
     /// seeded RLWE ciphertexts and return `SeededBatchedFheUint8`
@@ -249,7 +250,7 @@ where
             .iter()
             .flat_map(|v| (0..8).into_iter().map(|i| (((*v) >> i) & 1) == 1))
             .collect_vec();
-        let (cts, seed) = K::encrypt(&self, &bool_m);
+        let (cts, seed): (Vec<C>, S) = K::encrypt(&self, &bool_m).into();
         SeededBatchedFheUint8 {
             data: cts,
             seed,