From 4b43a11e26036b4055a49a27a62e2067a77cea78 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Tue, 18 May 2021 13:38:20 +0200
Subject: [PATCH] Make search() return more information

---
 instant-distance/src/lib.rs   | 39 ++++++++++++++++++++++++++---------
 instant-distance/tests/all.rs |  2 +-
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/instant-distance/src/lib.rs b/instant-distance/src/lib.rs
index 9f5a252..fa82d43 100644
--- a/instant-distance/src/lib.rs
+++ b/instant-distance/src/lib.rs
@@ -155,11 +155,13 @@ where
         &'a self,
         point: &P,
         search: &'a mut Search,
-    ) -> impl Iterator<Item = (f32, &'a V)> + ExactSizeIterator + 'a {
-        self.hnsw.search(point, search).map(move |candidate| {
-            let value = &self.values[candidate.pid.0 as usize];
-            (candidate.distance.into(), value)
-        })
+    ) -> impl Iterator<Item = (f32, &'a P, &'a V)> + ExactSizeIterator + 'a {
+        self.hnsw
+            .search(point, search)
+            .map(move |(distance, pid, point)| {
+                let value = &self.values[pid.0 as usize];
+                (distance, point, value)
+            })
     }
 }
 
@@ -336,14 +338,15 @@ where
     /// The results are returned in the `out` parameter; the number of neighbors to search for
     /// is limited by the size of the `out` parameter, and the number of results found is returned
     /// in the return value.
-    pub fn search<'a>(
-        &self,
+    pub fn search<'a, 'b: 'a>(
+        &'b self,
         point: &P,
         search: &'a mut Search,
-    ) -> impl Iterator<Item = Candidate> + ExactSizeIterator + 'a {
+    ) -> impl Iterator<Item = Item<'b, P>> + ExactSizeIterator + 'a {
         search.reset();
+        let map = move |candidate| Item::new(candidate, self);
         if self.points.is_empty() {
-            return search.iter();
+            return search.iter().map(map);
         }
 
         search.visited.reserve_capacity(self.points.len());
@@ -365,7 +368,7 @@ where
             }
         }
 
-        search.iter()
+        search.iter().map(map)
     }
 
     /// Iterate over the keys and values in this index
@@ -377,6 +380,22 @@ where
     }
 }
 
+pub struct Item<'a, P> {
+    pub distance: f32,
+    pub pid: PointId,
+    pub point: &'a P,
+}
+
+impl<'a, P> Item<'a, P> {
+    fn new(candidate: Candidate, hnsw: &'a Hnsw<P>) -> Self {
+        Self {
+            distance: candidate.distance.into_inner(),
+            pid: candidate.pid,
+            point: &hnsw[candidate.pid],
+        }
+    }
+}
+
 struct Construction<'a, P: Point> {
     zero: &'a [RwLock<ZeroNode>],
     pool: SearchPool,
diff --git a/instant-distance/tests/all.rs b/instant-distance/tests/all.rs
index 92a9ead..3c25964 100644
--- a/instant-distance/tests/all.rs
+++ b/instant-distance/tests/all.rs
@@ -70,7 +70,7 @@ fn randomized(builder: Builder) -> (u64, usize) {
         .collect::<HashSet<_>>();
     let found = results
         .take(100)
-        .map(|candidate| candidate.pid)
+        .map(|item| item.pid)
         .collect::<HashSet<_>>();
     (seed, forced.intersection(&found).count())
 }