[POP-4099] seed extended searches

iris-mpc-cpu/src/execution/hawk_main/search.rs

@@ -231,8 +231,9 @@
         );
         metrics::counter!("supermatcher_extended_searches").increment(1);
 
+        let seeded_nbhd = SortedNeighborhood::from_ascending_vec(edges.to_vec());
         let supermatch_neighbors = hnsw_supermatch
-            .search(aby3_store, graph_store, query, ef_supermatch)
+            .search_layer_0_seeded(aby3_store, graph_store, query, seeded_nbhd, ef_supermatch)
             .await?;
 
         let supermatch_classified = classify_edges(

iris-mpc-cpu/src/hnsw/searcher.rs

@@ -1382,6 +1382,44 @@
         Ok(W)
     }
 
+    /// Extended layer-0-only search, seeded with the neighborhood produced by a
+    /// previous search. Used to confirm/deny a potential supermatcher without
+    /// repeating the upper-layer greedy descent.
+    ///
+    /// The edges of  `seeded_nbhd` must be sorted ascending by distance.
+    #[instrument(level = "trace", target = "searcher::network", skip_all)]
+    pub async fn search_layer_0_seeded<V: VectorStore>(
+        &self,
+        store: &mut V,
+        graph: &GraphMem,
+        query: &V::QueryRef,
+        seeded_nbhd: SortedNeighborhood<V>,
+        ef: usize,
+    ) -> Result<SortedNeighborhood<V>> {
+        let start = Instant::now();
+
+        if seeded_nbhd.is_empty() {
+            let W = self.search(store, graph, query, ef).await?;
+            return Ok(W);
+        }
+
+        let mut W = seeded_nbhd;
+        Self::search_layer(
+            store,
+            graph,
+            query,
+            &mut W,
+            ef,
+            0,
+            self.min_layer_search_batch_size,
+        )
+        .await?;
+
+        W.edges.truncate(ef);
+        metrics::histogram!("search_layer_0_seeded_duration").record(start.elapsed().as_secs_f64());
+        Ok(W)
+    }
+
     /// Insert `query` into the HNSW index represented by `store` and `graph`.
     /// Return a `VectorId` representing the inserted vector.
     #[instrument(level = "trace", skip_all, target = "searcher::network")]
@@ -2042,4 +2080,87 @@
         assert!(ops.is_empty());
         Ok(())
     }
+
+    /// Build a plaintext HNSW graph with `n` random elements, returning the
+    /// store, graph, the inserted queries, and their assigned `VectorId`s (in
+    /// insertion order).
+    async fn build_seeded_test_graph(
+        db: &HnswSearcher,
+        seed: u64,
+        n: usize,
+    ) -> Result<(
+        PlaintextStore<FhdOps>,
+        GraphMem,
+        Vec<Arc<IrisCode>>,
+        Vec<VectorId>,
+    )> {
+        let mut store = PlaintextStore::<FhdOps>::new();
+        let mut graph = GraphMem::new();
+        let rng = &mut AesRng::seed_from_u64(seed);
+
+        let queries = IrisDB::new_random_rng(n, rng)
+            .db
+            .into_iter()
+            .map(Arc::new)
+            .collect::<Vec<_>>();
+
+        let mut ids = Vec::with_capacity(n);
+        for query in queries.iter() {
+            let insertion_layer = db.gen_layer_rng(rng)?;
+            let id = db
+                .insert(&mut store, &mut graph, query, insertion_layer)
+                .await?;
+            ids.push(id);
+        }
+
+        Ok((store, graph, queries, ids))
+    }
+
+    /// For a node that lives in layer 0, a layer-0-only search seeded with
+    /// the result of a regular top-down search should reproduce that same
+    /// result.
+    #[tokio::test]
+    async fn test_search_layer_0_seeded_matches_full_search() -> Result<()> {
+        let db = HnswSearcher::new_with_test_parameters();
+        let (mut store, graph, queries, ids) = build_seeded_test_graph(&db, 0, 256).await?;
+
+        let query_idx = 0;
+        assert!(graph.layers[0].links.contains_key(&ids[query_idx]));
+        let query = &queries[query_idx];
+        let ef = db.params.get_ef_search(1);
+
+        // Regular top-down search.
+        let full = db.search(&mut store, &graph, query, ef).await?;
+
+        // Seed a layer-0-only search with the full result; it must be a fixed point.
+        let seeded = db
+            .search_layer_0_seeded(&mut store, &graph, query, full.clone(), ef)
+            .await?;
+
+        assert_eq!(full.edge_ids(), seeded.edge_ids());
+        assert_eq!(seeded.edge_ids()[query_idx], ids[query_idx]);
+        Ok(())
+    }
+
+    /// An empty seeded neighborhood should fall back to a regular search and
+    /// therefore return exactly the same result as `search`.
+    #[tokio::test]
+    async fn test_search_layer_0_seeded_empty_falls_back_to_search() -> Result<()> {
+        let db = HnswSearcher::new_with_test_parameters();
+        let (mut store, graph, queries, ids) = build_seeded_test_graph(&db, 0, 256).await?;
+
+        let query_idx = 0;
+        assert!(graph.layers[0].links.contains_key(&ids[query_idx]));
+        let query = &queries[query_idx];
+        let ef = db.params.get_ef_search(1);
+
+        let full = db.search(&mut store, &graph, query, ef).await?;
+        let seeded = db
+            .search_layer_0_seeded(&mut store, &graph, query, SortedNeighborhood::new(), ef)
+            .await?;
+
+        assert_eq!(full.edge_ids(), seeded.edge_ids());
+        assert_eq!(seeded.edge_ids()[query_idx], ids[query_idx]);
+        Ok(())
+    }
 }