From 6fda656afb27791c78588e9da7f4ff257f6783c6 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Tue, 18 May 2021 13:42:28 +0200
Subject: [PATCH] Add HnswMap::iter() method

---
 instant-distance/src/lib.rs | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/instant-distance/src/lib.rs b/instant-distance/src/lib.rs
index fa82d43..3b75472 100644
--- a/instant-distance/src/lib.rs
+++ b/instant-distance/src/lib.rs
@@ -155,14 +155,26 @@ where
         &'a self,
         point: &P,
         search: &'a mut Search,
-    ) -> 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)
-            })
+    ) -> impl Iterator<Item = MapItem<'a, P, V>> + ExactSizeIterator + 'a {
+        self.hnsw.search(point, search).map(move |item| MapItem {
+            distance: item.distance,
+            pid: item.pid,
+            point: item.point,
+            value: &self.values[item.pid.0 as usize],
+        })
     }
+
+    /// Iterate over the keys and values in this index
+    pub fn iter(&self) -> impl Iterator<Item = (PointId, &P)> {
+        self.hnsw.iter()
+    }
+}
+
+pub struct MapItem<'a, P, V> {
+    pub distance: f32,
+    pub pid: PointId,
+    pub point: &'a P,
+    pub value: &'a V,
 }
 
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]