Refactor SMPC protocol implementation (#876)

* Rename SMPC implementation and detach local runtime * Keep point storage in RwLock
worldcoin · Jan 8, 2025 · 8d8f85b · 8d8f85b
1 parent 0a2e977
commit 8d8f85b
Show file tree

Hide file tree

Showing 7 changed files with 144 additions and 137 deletions.
diff --git a/iris-mpc-cpu/benches/hnsw.rs b/iris-mpc-cpu/benches/hnsw.rs
@@ -5,7 +5,7 @@ use iris_mpc_common::iris_db::{db::IrisDB, iris::IrisCode};
 use iris_mpc_cpu::{
     database_generators::{create_random_sharing, generate_galois_iris_shares},
     execution::local::LocalRuntime,
-    hawkers::{galois_store::LocalNetAby3NgStoreProtocol, plaintext_store::PlaintextStore},
+    hawkers::{aby3_store::Aby3Store, plaintext_store::PlaintextStore},
     protocol::ops::{
         batch_signed_lift_vec, cross_compare, galois_ring_pairwise_distance, galois_ring_to_rep3,
     },
@@ -75,7 +75,7 @@ fn bench_hnsw_primitives(c: &mut Criterion) {
             let runtime = LocalRuntime::mock_setup_with_grpc().await.unwrap();
 
             let mut jobs = JoinSet::new();
-            for (index, player) in runtime.identities.iter().enumerate() {
+            for (index, player) in runtime.get_identities().iter().enumerate() {
                 let d1i = d1[index].clone();
                 let d2i = d2[index].clone();
                 let t1i = t1[index].clone();
@@ -123,7 +123,7 @@ fn bench_gr_primitives(c: &mut Criterion) {
             let y2 = generate_galois_iris_shares(&mut rng, iris_db[3].clone());
 
             let mut jobs = JoinSet::new();
-            for (index, player) in runtime.identities.iter().enumerate() {
+            for (index, player) in runtime.get_identities().iter().enumerate() {
                 let x1 = x1[index].clone();
                 let mut y1 = y1[index].clone();
 
@@ -178,7 +178,7 @@ fn bench_gr_ready_made_hnsw(c: &mut Criterion) {
 
         let secret_searcher = rt.block_on(async move {
             let mut rng = AesRng::seed_from_u64(0_u64);
-            LocalNetAby3NgStoreProtocol::lazy_setup_from_files_with_grpc(
+            Aby3Store::lazy_setup_from_files_with_grpc(
                 "./data/store.ndjson",
                 &format!("./data/graph_{}.dat", database_size),
                 &mut rng,

diff --git a/iris-mpc-cpu/bin/local_hnsw.rs b/iris-mpc-cpu/bin/local_hnsw.rs
@@ -1,6 +1,6 @@
 use aes_prng::AesRng;
 use clap::Parser;
-use iris_mpc_cpu::hawkers::galois_store::LocalNetAby3NgStoreProtocol;
+use iris_mpc_cpu::hawkers::aby3_store::Aby3Store;
 use rand::SeedableRng;
 use std::error::Error;
 
@@ -18,7 +18,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     println!("Starting Local HNSW with {} vectors", database_size);
     let mut rng = AesRng::seed_from_u64(0_u64);
 
-    LocalNetAby3NgStoreProtocol::shared_random_setup_with_grpc(&mut rng, database_size).await?;
+    Aby3Store::shared_random_setup_with_grpc(&mut rng, database_size).await?;
 
     Ok(())
 }
diff --git a/iris-mpc-cpu/src/execution/local.rs b/iris-mpc-cpu/src/execution/local.rs
@@ -41,14 +41,22 @@ pub async fn get_free_local_addresses(num_ports: usize) -> eyre::Result<Vec<Stri
 
 #[derive(Debug, Clone)]
 pub struct LocalRuntime {
-    pub identities:       Vec<Identity>,
-    pub role_assignments: RoleAssignment,
-    pub seeds:            Vec<PrfSeed>,
     // only one session per player is created
-    pub sessions:         HashMap<Identity, Session>,
+    pub sessions: HashMap<Identity, Session>,
 }
 
 impl LocalRuntime {
+    pub fn get_session(&self, identity: &Identity) -> eyre::Result<Session> {
+        self.sessions
+            .get(identity)
+            .ok_or_else(|| eyre::eyre!(format!("Session not found for identity: {:?}", identity)))
+            .cloned()
+    }
+
+    pub fn get_identities(&self) -> Vec<Identity> {
+        self.sessions.keys().cloned().collect()
+    }
+
     pub async fn mock_setup(network_t: NetworkType) -> eyre::Result<Self> {
         let num_parties = 3;
         let identities = generate_local_identities();
@@ -138,12 +146,7 @@ impl LocalRuntime {
                 setup: prf,
             });
         }
-        Ok(LocalRuntime {
-            identities,
-            role_assignments,
-            seeds,
-            sessions,
-        })
+        Ok(LocalRuntime { sessions })
     }
 
     pub async fn new(identities: Vec<Identity>, seeds: Vec<PrfSeed>) -> eyre::Result<Self> {