Strengthen arguments to robust_prune.

diskann/src/graph/index.rs

@@ -2601,26 +2601,19 @@ where
         //
         // Iteration preserves the pool's nearest-first order. Pruning stores results
         // by cache position and resolves each position to `local_id` afterward.
-        let sorted_cache: Vec<(f32, Option<_>)> = pool
-            .iter()
-            .map(|neighbor| {
-                // Filter out self loops.
-                let id = neighbor.id();
-                if exclude(*id) {
-                    (*neighbor.distance(), None)
-                } else {
-                    (*neighbor.distance(), map.get(*id))
-                }
-            })
-            .collect();
+        let mut storage = Vec::with_capacity(pool.len());
+        let cache = pool.map_in(
+            &mut storage,
+            |id| if exclude(*id) { None } else { map.get(*id) },
+        );
 
         let found = prune::robust_prune(
-            &sorted_cache,
+            cache,
             states,
             degree,
             alpha,
             self.config.prune_kind(),
-            |neighbor, result| {
+            |neighbor: &M::Element<'_>, result: &M::Element<'_>| {
                 computer.evaluate_similarity((*neighbor).reborrow(), result.reborrow())
             },
         );

diskann/src/graph/internal/prune.rs

@@ -5,11 +5,9 @@
 
 use thiserror::Error;
 
-use super::SortedNeighbors;
-
 use crate::{
     ANNError, error,
-    graph::{AdjacencyList, config::PruneKind},
+    graph::{AdjacencyList, config::PruneKind, internal::SortedNeighbors},
     neighbor::Neighbor,
     utils::{IntoUsize, VectorId},
 };
@@ -97,14 +95,14 @@ pub(crate) struct State {
 
 /// Select a degree-bounded neighbor set with Vamana RobustPrune.
 ///
-/// `sorted_cache` stores each source distance and candidate vector in ascending
+/// `candidates` stores each source distance and candidate vector in ascending
 /// source-distance order. `None` excludes that candidate without changing
 /// positional alignment. `states` has one entry for each candidate position.
 ///
 /// The function writes selected candidate indexes to `states[..result]` and
 /// returns `result`. The caller converts those indexes to graph IDs.
 pub(in crate::graph) fn robust_prune<V, D>(
-    sorted_cache: &[(f32, Option<V>)],
+    candidates: SortedNeighbors<'_, Option<V>>,
     states: &mut [State],
     degree: usize,
     alpha: f32,
@@ -114,11 +112,6 @@ pub(in crate::graph) fn robust_prune<V, D>(
 where
     D: FnMut(&V, &V) -> f32,
 {
-    debug_assert!(
-        sorted_cache.is_sorted_by_key(|(distance, _)| distance),
-        "candidate cache must be sorted by source distance"
-    );
-
     let mut current_alpha = 1.0f32;
     let increment_factor = alpha.min(1.2);
 
@@ -138,11 +131,11 @@ where
     //
     // On the implementation side, we use `states` in the following way:
     //
-    // * `states[n].neighbor` is the **index** in `sorted_cache` of the `n`th **neighbor**.
+    // * `states[n].neighbor` is the **index** in `candidates` of the `n`th **neighbor**.
     //   Note that a "neighbor" is a candidate that passes pruning.
     //
     //   Very important: to get the index `j` in the above description, we need to
-    //   check `sorted_cache[states[n].neighbor]`.
+    //   check `candidates[states[n].neighbor]`.
     //
     //   This indexing naturally skips candidates `j` that have not been promoted to
     //   neighbors.
@@ -152,14 +145,14 @@ where
     //   excludes it from future consideration.
     //
     // * `states[i].last_checked` is the highest value of `n` against which the
-    //   occlude factor for `j = sorted_cache[states[n].neighbor]` has been checked.
+    //   occlude factor for `j = candidates[states[n].neighbor]` has been checked.
     //
     //   The maximum value this should reach is `i`.
     //
     // Note that we use `states` for both "candidate" and "neighbor" tracking.
     let mut found = 0;
     while found < degree {
-        for (i, (neighbor_distance, neighbor)) in sorted_cache.iter().enumerate() {
+        for (i, neighbor) in candidates.iter().enumerate() {
             if found >= degree {
                 break;
             }
@@ -176,10 +169,11 @@ where
                 continue;
             }
 
-            // Retrieval from the cache might not be perfect.
+            // Retrieval from the candidates might not be perfect.
             //
-            // This neighbor did not end up in the cache, then just skip it.
-            let neighbor = match neighbor {
+            // This neighbor did not end up in the candidates, then just skip it.
+            let neighbor_distance = neighbor.distance();
+            let neighbor = match neighbor.id() {
                 Some(n) => n,
                 None => {
                     debug_assert!(states.get(i).is_some(), "index {i} is out of bounds");
@@ -198,7 +192,7 @@ where
                 let result_position = states[last_checked as usize].neighbor.into_usize();
                 last_checked += 1;
 
-                // If the position of this result in `sorted_cache` is greater than or equal
+                // If the position of this result in `candidates` is greater than or equal
                 // to the current working position, then skip this candidate.
                 if result_position >= i {
                     debug_assert!(states.get(i).is_some(), "index {i} is out of bounds");
@@ -209,9 +203,9 @@ where
 
                 // Otherwise, compute the distance between the result and this neighbor
                 // and update the occlude factor.
-                let distance = match &sorted_cache[result_position] {
-                    (_, Some(v)) => compute_distance(neighbor, v),
-                    (_, None) => f32::MAX,
+                let distance = match candidates[result_position].id() {
+                    Some(v) => compute_distance(neighbor, v),
+                    None => f32::MAX,
                 };
 
                 // Update occlude factor

diskann/src/graph/internal/sorted_neighbors.rs

@@ -11,19 +11,22 @@ use crate::neighbor::{self, Neighbor};
 
 /// A utility that asserts the contained neighbors are sorted by distance.
 #[derive(Debug)]
-pub struct SortedNeighbors<'a, I>(&'a [Neighbor<I>])
-where
-    I: Eq;
-
-impl<'a, I> SortedNeighbors<'a, I>
-where
-    I: Eq + std::fmt::Debug,
-{
+pub(crate) struct SortedNeighbors<'a, I>(&'a [Neighbor<I>]);
+
+impl<I> Clone for SortedNeighbors<'_, I> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<I> Copy for SortedNeighbors<'_, I> {}
+
+impl<'a, I> SortedNeighbors<'a, I> {
     /// Create a new `SortedNeighbors` around `neighbors` truncated to `max` length.
     ///
     /// As a by-product calling this method, `neighbors` will be resized to at most
     /// `max` and be sorted.
-    pub fn new(neighbors: &'a mut Vec<Neighbor<I>>, max: usize) -> Self {
+    pub(crate) fn new(neighbors: &'a mut Vec<Neighbor<I>>, max: usize) -> Self {
         // Here- we use `select_nth_unstable` to get the `position` index in the correct
         // location. We can then sort the prefix slice returned by that API.
         //
@@ -42,12 +45,29 @@ where
         neighbors.truncate(max);
         Self(&*neighbors)
     }
+
+    /// Apply the projection `f` to each element in `self` and store the result in `other`.
+    ///
+    /// The returned [`SortedNeighbors`] inherits the sorted property from `self`.
+    ///
+    /// # Side Effects
+    ///
+    /// This method removes all pre-existing elements from `storage`.
+    pub(crate) fn map_in<'b, F, J>(
+        self,
+        storage: &'b mut Vec<Neighbor<J>>,
+        mut f: F,
+    ) -> SortedNeighbors<'b, J>
+    where
+        F: FnMut(&I) -> J,
+    {
+        storage.clear();
+        storage.extend(self.iter().map(|n| Neighbor::new(f(n.id()), *n.distance())));
+        SortedNeighbors(storage)
+    }
 }
 
-impl<I> Deref for SortedNeighbors<'_, I>
-where
-    I: Eq,
-{
+impl<I> Deref for SortedNeighbors<'_, I> {
     type Target = [Neighbor<I>];
     fn deref(&self) -> &Self::Target {
         self.0
@@ -108,4 +128,46 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_map() {
+        let messages = ["a", "b", "c", "d", "e", "f"];
+
+        let mut storage = vec![Neighbor::new("foo", 1.0), Neighbor::new("bar", 0.0)];
+
+        {
+            let mut neighbors = vec![
+                Neighbor::new(0usize, 5.0f32),
+                Neighbor::new(1, 4.0),
+                Neighbor::new(2, 3.0),
+                Neighbor::new(3, 2.0),
+                Neighbor::new(4, 1.0),
+                Neighbor::new(5, 0.0),
+            ];
+
+            let sorted = SortedNeighbors::new(&mut neighbors, 6);
+            let cache = sorted.map_in(&mut storage, |id: &usize| messages[*id]);
+
+            assert_eq_verbose!(
+                *cache,
+                [
+                    Neighbor::new("f", 0.0f32),
+                    Neighbor::new("e", 1.0),
+                    Neighbor::new("d", 2.0),
+                    Neighbor::new("c", 3.0),
+                    Neighbor::new("b", 4.0),
+                    Neighbor::new("a", 5.0),
+                ]
+                .as_slice()
+            );
+        }
+
+        // Empty
+        {
+            let mut neighbors = Vec::<Neighbor<usize>>::new();
+            let sorted = SortedNeighbors::new(&mut neighbors, 10);
+            let cache = sorted.map_in(&mut storage, |id: &usize| messages[*id]);
+            assert!(cache.is_empty());
+        }
+    }
 }

diskann/src/test/cmp.rs

@@ -108,6 +108,15 @@ pub(crate) use assert_eq_verbose;
 // Implementation //
 ////////////////////
 
+impl<T> VerboseEq for &T
+where
+    T: VerboseEq + ?Sized,
+{
+    fn verbose_eq(&self, other: &Self) -> ANNResult<()> {
+        (*self).verbose_eq(*other)
+    }
+}
+
 /// Display implementation for recording a field mismatch.
 #[derive(Debug)]
 pub(crate) struct Field(pub(crate) &'static str);
@@ -130,7 +139,7 @@ macro_rules! impl_via_partial_eq {
         impl $crate::test::cmp::VerboseEq for $T {
             fn verbose_eq(&self, other: &Self) -> $crate::ANNResult<()> {
                 if self != other {
-                    Err($crate::ANNError::new(NotEq(self.clone(), other.clone())))
+                    Err($crate::ANNError::new(NotEq(self.to_owned(), other.to_owned())))
                 } else {
                     Ok(())
                 }
@@ -143,7 +152,7 @@ macro_rules! impl_via_partial_eq {
 }
 
 impl_via_partial_eq!(
-    u8, u16, u32, u64, i8, i16, i32, i64, usize, f32, f64, String, bool,
+    u8, u16, u32, u64, i8, i16, i32, i64, usize, f32, f64, String, str, bool,
 );
 
 macro_rules! impl_tuple {