Port local copy of HawkSearcher for application-specific customizatio…

…ns (#898) * Port local copy of HawkSearcher for application-specific customizations * Modify Python integration test to store data files in ignored directory --------- Co-authored-by: Bryan Gillespie <[email protected]>
worldcoin · Jan 10, 2025 · 50c9bdd · 50c9bdd
1 parent 7d6de6f
commit 50c9bdd
Show file tree

Hide file tree

Showing 16 changed files with 590 additions and 58 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/iris-mpc-cpu/Cargo.toml b/iris-mpc-cpu/Cargo.toml
@@ -26,6 +26,7 @@ itertools.workspace = true
 num-traits.workspace = true
 prost = "0.13"
 rand.workspace = true
+rand_distr = "0.4.3"
 rstest = "0.23.0"
 serde.workspace = true
 serde_json.workspace = true
@@ -57,4 +58,4 @@ path = "bin/local_hnsw.rs"
 
 [[bin]]
 name = "generate_benchmark_data"
-path = "bin/generate_benchmark_data.rs"
+path = "bin/generate_benchmark_data.rs"
diff --git a/iris-mpc-cpu/benches/hnsw.rs b/iris-mpc-cpu/benches/hnsw.rs
@@ -1,11 +1,12 @@
 use aes_prng::AesRng;
 use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, SamplingMode};
-use hawk_pack::{graph_store::GraphMem, HawkSearcher};
+use hawk_pack::graph_store::GraphMem;
 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::{aby3_store::Aby3Store, plaintext_store::PlaintextStore},
+    hnsw::searcher::HnswSearcher,
     protocol::ops::{
         batch_signed_lift_vec, cross_compare, galois_ring_pairwise_distance, galois_ring_to_rep3,
     },
@@ -28,7 +29,7 @@ fn bench_plaintext_hnsw(c: &mut Criterion) {
             let mut rng = AesRng::seed_from_u64(0_u64);
             let mut vector = PlaintextStore::default();
             let mut graph = GraphMem::new();
-            let searcher = HawkSearcher::default();
+            let searcher = HnswSearcher::default();
 
             for _ in 0..database_size {
                 let raw_query = IrisCode::random_rng(&mut rng);
@@ -44,7 +45,7 @@ fn bench_plaintext_hnsw(c: &mut Criterion) {
             b.to_async(&rt).iter_batched(
                 || (vector.clone(), graph.clone()),
                 |(mut db_vectors, mut graph)| async move {
-                    let searcher = HawkSearcher::default();
+                    let searcher = HnswSearcher::default();
                     let mut rng = AesRng::seed_from_u64(0_u64);
                     let on_the_fly_query = IrisDB::new_random_rng(1, &mut rng).db[0].clone();
                     let query = db_vectors.prepare_query(on_the_fly_query);
@@ -201,7 +202,7 @@ fn bench_gr_ready_made_hnsw(c: &mut Criterion) {
                 b.to_async(&rt).iter_batched(
                     || secret_searcher.clone(),
                     |vectors_graphs| async move {
-                        let searcher = HawkSearcher::default();
+                        let searcher = HnswSearcher::default();
                         let mut rng = AesRng::seed_from_u64(0_u64);
                         let on_the_fly_query = IrisDB::new_random_rng(1, &mut rng).db[0].clone();
                         let raw_query = generate_galois_iris_shares(&mut rng, on_the_fly_query);
@@ -235,7 +236,7 @@ fn bench_gr_ready_made_hnsw(c: &mut Criterion) {
                 b.to_async(&rt).iter_batched(
                     || secret_searcher.clone(),
                     |vectors_graphs| async move {
-                        let searcher = HawkSearcher::default();
+                        let searcher = HnswSearcher::default();
                         let mut rng = AesRng::seed_from_u64(0_u64);
                         let on_the_fly_query = IrisDB::new_random_rng(1, &mut rng).db[0].clone();
                         let raw_query = generate_galois_iris_shares(&mut rng, on_the_fly_query);

diff --git a/iris-mpc-cpu/examples/hnsw-ex.rs b/iris-mpc-cpu/examples/hnsw-ex.rs
@@ -1,7 +1,7 @@
 use aes_prng::AesRng;
-use hawk_pack::{graph_store::GraphMem, HawkSearcher};
+use hawk_pack::graph_store::GraphMem;
 use iris_mpc_common::iris_db::iris::IrisCode;
-use iris_mpc_cpu::hawkers::plaintext_store::PlaintextStore;
+use iris_mpc_cpu::{hawkers::plaintext_store::PlaintextStore, hnsw::searcher::HnswSearcher};
 use rand::SeedableRng;
 
 const DATABASE_SIZE: usize = 1_000;
@@ -16,7 +16,7 @@ fn main() {
         let mut rng = AesRng::seed_from_u64(0_u64);
         let mut vector = PlaintextStore::default();
         let mut graph = GraphMem::new();
-        let searcher = HawkSearcher::default();
+        let searcher = HnswSearcher::default();
 
         for idx in 0..DATABASE_SIZE {
             let raw_query = IrisCode::random_rng(&mut rng);

diff --git a/iris-mpc-cpu/src/hawkers/aby3_store.rs b/iris-mpc-cpu/src/hawkers/aby3_store.rs
@@ -1,12 +1,12 @@
-use super::plaintext_store::PlaintextStore;
 use crate::{
     database_generators::{generate_galois_iris_shares, GaloisRingSharedIris},
     execution::{
         local::{generate_local_identities, LocalRuntime},
         player::Identity,
         session::{Session, SessionHandles},
     },
-    hawkers::plaintext_store::PointId,
+    hawkers::plaintext_store::{PlaintextStore, PointId},
+    hnsw::HnswSearcher,
     network::NetworkType,
     protocol::ops::{
         batch_signed_lift_vec, compare_threshold_and_open, cross_compare,
@@ -22,7 +22,7 @@ use aes_prng::AesRng;
 use hawk_pack::{
     data_structures::queue::FurthestQueue,
     graph_store::{graph_mem::Layer, GraphMem},
-    GraphStore, HawkSearcher, VectorStore,
+    GraphStore, VectorStore,
 };
 use iris_mpc_common::iris_db::db::IrisDB;
 use rand::{CryptoRng, RngCore, SeedableRng};
@@ -563,7 +563,7 @@ impl Aby3Store {
                 .collect::<Vec<_>>();
             jobs.spawn(async move {
                 let mut graph_store = GraphMem::new();
-                let searcher = HawkSearcher::default();
+                let searcher = HnswSearcher::default();
                 // insert queries
                 for query in queries.iter() {
                     searcher
@@ -603,9 +603,9 @@ impl Aby3Store {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::database_generators::generate_galois_iris_shares;
+    use crate::{database_generators::generate_galois_iris_shares, hnsw::HnswSearcher};
     use aes_prng::AesRng;
-    use hawk_pack::{graph_store::GraphMem, HawkSearcher};
+    use hawk_pack::graph_store::GraphMem;
     use itertools::Itertools;
     use rand::SeedableRng;
     use tracing_test::traced_test;
@@ -634,7 +634,7 @@ mod tests {
             let mut rng = rng.clone();
             jobs.spawn(async move {
                 let mut aby3_graph = GraphMem::new();
-                let db = HawkSearcher::default();
+                let db = HnswSearcher::default();
 
                 let mut inserted = vec![];
                 // insert queries
@@ -693,7 +693,7 @@ mod tests {
                 premade_v.storage.body.read().unwrap().points
             );
         }
-        let hawk_searcher = HawkSearcher::default();
+        let hawk_searcher = HnswSearcher::default();
 
         for i in 0..database_size {
             let cleartext_neighbors = hawk_searcher
@@ -840,7 +840,7 @@ mod tests {
     async fn test_gr_scratch_hnsw() {
         let mut rng = AesRng::seed_from_u64(0_u64);
         let database_size = 2;
-        let searcher = HawkSearcher::default();
+        let searcher = HnswSearcher::default();
         let mut vectors_and_graphs =
             Aby3Store::shared_random_setup(&mut rng, database_size, NetworkType::LocalChannel)
                 .await

diff --git a/iris-mpc-cpu/src/hawkers/plaintext_store.rs b/iris-mpc-cpu/src/hawkers/plaintext_store.rs
@@ -1,5 +1,6 @@
+use crate::hnsw::HnswSearcher;
 use aes_prng::AesRng;
-use hawk_pack::{graph_store::GraphMem, HawkSearcher, VectorStore};
+use hawk_pack::{graph_store::GraphMem, VectorStore};
 use iris_mpc_common::iris_db::{
     db::IrisDB,
     iris::{IrisCode, MATCH_THRESHOLD_RATIO},
@@ -148,7 +149,7 @@ impl PlaintextStore {
 
         let mut plaintext_vector_store = PlaintextStore::default();
         let mut plaintext_graph_store = GraphMem::new();
-        let searcher = HawkSearcher::default();
+        let searcher = HnswSearcher::default();
 
         for raw_query in cleartext_database.iter() {
             let query = plaintext_vector_store.prepare_query(raw_query.clone());
@@ -189,7 +190,7 @@ impl PlaintextStore {
         let mut rng_searcher1 = AesRng::from_rng(rng.clone())?;
 
         let mut plaintext_graph_store = GraphMem::new();
-        let searcher = HawkSearcher::default();
+        let searcher = HnswSearcher::default();
 
         for i in 0..graph_size {
             searcher
@@ -209,8 +210,8 @@ impl PlaintextStore {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::hnsw::HnswSearcher;
     use aes_prng::AesRng;
-    use hawk_pack::HawkSearcher;
     use iris_mpc_common::iris_db::db::IrisDB;
     use rand::SeedableRng;
     use tracing_test::traced_test;
@@ -292,7 +293,7 @@ mod tests {
     async fn test_plaintext_hnsw_matcher() {
         let mut rng = AesRng::seed_from_u64(0_u64);
         let database_size = 1;
-        let searcher = HawkSearcher::default();
+        let searcher = HnswSearcher::default();
         let (mut ptxt_vector, mut ptxt_graph) =
             PlaintextStore::create_random(&mut rng, database_size)
                 .await

diff --git a/iris-mpc-cpu/src/hnsw/mod.rs b/iris-mpc-cpu/src/hnsw/mod.rs
@@ -0,0 +1,3 @@
+pub mod searcher;
+
+pub use searcher::HnswSearcher;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		pub mod searcher;

		pub use searcher::HnswSearcher;