match compatible-with-abe

Cargo.toml

@@ -17,7 +17,7 @@ disk = ["libc", "memmap2", "tempfile"]
 cpu = []
 
 [dependencies]
-mxx = { git = "https://github.com/MachinaIO/mxx.git" }
+mxx = { git = "https://github.com/MachinaIO/mxx.git", branch = "feat/compatible-with-abe" }
 tokio = { version = "1", features = ["fs", "rt-multi-thread", "macros"] }
 libc = { version = "0.2", optional = true }
 num-bigint = { version = "0.4", features = ["serde"] }

dio/Cargo.toml

@@ -8,7 +8,7 @@ debug = ["diamond-io/debug"]
 disk = ["diamond-io/disk"]
 
 [dependencies]
-mxx = { git = "https://github.com/MachinaIO/mxx.git" }
+mxx = { git = "https://github.com/MachinaIO/mxx.git", branch = "feat/compatible-with-abe" }
 diamond-io = { path = "../", default-features = false }
 clap = { version = "4.5.36", features = ["derive"] }
 num-bigint = { version = "0.4", features = ["serde"] }

src/io/eval.rs

@@ -7,7 +7,7 @@ use itertools::Itertools;
 use mxx::parallel_iter;
 use mxx::{
     bgg::{digits_to_int::DigitsToInt, encoding::BggEncoding, sampler::BGGPublicKeySampler},
-    lookup::simple_eval::SimpleBggEncodingPltEvaluator,
+    lookup::lwe_eval::LweBggEncodingPltEvaluator,
     matrix::PolyMatrix,
     poly::{Poly, PolyParams},
     sampler::{PolyHashSampler, PolyTrapdoorSampler},
@@ -54,8 +54,8 @@ where
         hash_key.copy_from_slice(&bytes);
         hash_key
     };
-
-    let bgg_pubkey_sampler = BGGPublicKeySampler::<_, SH>::new(hash_key, d);
+    // explictly sample d+1 public keys
+    let bgg_pubkey_sampler = BGGPublicKeySampler::<_, SH>::new(hash_key, d + 1);
     let public_data = PublicSampledData::<SH>::sample(&obf_params, hash_key);
     log_mem("Sampled public data");
     let packed_input_size = public_data.packed_input_size;
@@ -341,7 +341,7 @@ where
         }
     }
     let bgg_encoding_plt_evaluator =
-        SimpleBggEncodingPltEvaluator::<M, SH>::new(hash_key, dir_path.clone(), p_l_plus_one);
+        LweBggEncodingPltEvaluator::<M, SH>::new(hash_key, dir_path.clone(), p_l_plus_one);
     let output_encodings = final_circuit.eval(
         params.as_ref(),
         &new_encodings[0],

src/io/obf.rs

@@ -11,7 +11,7 @@ use itertools::Itertools;
 use mxx::storage::store_and_drop_poly;
 use mxx::{
     bgg::{digits_to_int::DigitsToInt, public_key::BggPublicKey, sampler::BGGPublicKeySampler},
-    lookup::simple_eval::SimpleBggPubKeyEvaluator,
+    lookup::lwe_eval::LweBggPubKeyEvaluator,
     matrix::PolyMatrix,
     poly::{Poly, PolyParams},
     rlwe_enc::rlwe_encrypt,
@@ -62,7 +62,8 @@ where
     let d = obf_params.d;
     let sampler_uniform = SU::new();
     let sampler_trapdoor = ST::new(&params, obf_params.trapdoor_sigma);
-    let bgg_pubkey_sampler = BGGPublicKeySampler::<_, SH>::new(hash_key, d);
+    // sample d+1 public keys explictly
+    let bgg_pubkey_sampler = BGGPublicKeySampler::<_, SH>::new(hash_key, d + 1);
     let m_b = (1 + packed_input_size) * (d + 1) * (2 + log_base_q);
     let packed_output_size = public_data.packed_output_size;
 
@@ -80,9 +81,6 @@ where
     =============================================================================
     */
 
-    // Sample BGG+ encoding secret key
-    let s_bars = sampler_uniform.sample_uniform(&params, 1, d, DistType::BitDist).get_row(0);
-    log_mem("Sampled s_bars");
     // Sample FHE secret key t
     let t_bar = sampler_uniform.sample_uniform(&params, 1, 1, DistType::BitDist);
     log_mem("Sampled t_bar");
@@ -102,6 +100,9 @@ where
         reveals
     };
     let s_init = {
+        // Sample BGG+ encoding secret key
+        let s_bars = sampler_uniform.sample_uniform(&params, 1, d, DistType::BitDist).get_row(0);
+        log_mem("Sampled s_bars");
         let minus_one_poly = <SU::M as PolyMatrix>::P::const_minus_one(&params);
         let mut secrets = s_bars.to_vec();
         secrets.push(minus_one_poly);
@@ -401,7 +402,7 @@ where
         );
         log_mem("Computed final_circuit");
 
-        let bgg_plt_evaluator = SimpleBggPubKeyEvaluator::<M, SH, SU, ST>::new(
+        let bgg_plt_evaluator = LweBggPubKeyEvaluator::<M, SH, ST>::new(
             hash_key,
             sampler_trapdoor.clone(),
             b_l_plus_one.clone(),

tests/test_io_final_circuit.rs

@@ -3,7 +3,7 @@ use keccak_asm::Keccak256;
 use mxx::{
     bgg::{digits_to_int::DigitsToInt, public_key::BggPublicKey, sampler::BGGPublicKeySampler},
     circuit::PolyCircuit,
-    lookup::simple_eval::SimpleBggPubKeyEvaluator,
+    lookup::lwe_eval::LweBggPubKeyEvaluator,
     matrix::dcrt_poly::DCRTPolyMatrix,
     poly::{
         dcrt::{params::DCRTPolyParams, poly::DCRTPoly},
@@ -82,7 +82,7 @@ fn test_build_final_step_circuit() {
     let packed_input_size = input_size.div_ceil(dim) + 1;
 
     let bgg_pubkey_sampler =
-        BGGPublicKeySampler::<_, DCRTPolyHashSampler<Keccak256>>::new(hash_key, d);
+        BGGPublicKeySampler::<_, DCRTPolyHashSampler<Keccak256>>::new(hash_key, d + 1);
     // consider inserting t on plaintext
     let reveal_plaintexts = [vec![true; packed_input_size], vec![false; 1]].concat();
     let pubkeys = bgg_pubkey_sampler.sample(&params, b"BGG_PUBKEY_INPUT:", &reveal_plaintexts);
@@ -93,10 +93,9 @@ fn test_build_final_step_circuit() {
         &pubkeys[0],
         &pubkeys[1..],
         None::<
-            SimpleBggPubKeyEvaluator<
+            LweBggPubKeyEvaluator<
                 DCRTPolyMatrix,
                 DCRTPolyHashSampler<Keccak256>,
-                DCRTPolyUniformSampler,
                 DCRTPolyTrapdoorSampler,
             >,
         >,