TraversalContext can be reused for IVF

rs/index/src/hnsw/index.rs

@@ -8,35 +8,12 @@ use quantization::typing::VectorOps;
 use rand::Rng;
 use utils::distance::l2::L2DistanceCalculatorImpl::StreamingSIMD;
 
-use super::utils::{GraphTraversal, TraversalContext};
+use super::utils::GraphTraversal;
 use crate::hnsw::writer::Header;
 use crate::index::Index;
-use crate::utils::{IdWithScore, SearchContext};
+use crate::utils::{IdWithScore, SearchContext, TraversalContext};
 use crate::vector::fixed_file::FixedFileVectorStorage;
 
-impl TraversalContext for SearchContext {
-    fn visited(&self, i: u32) -> bool {
-        self.visited.contains(i)
-    }
-
-    fn set_visited(&mut self, i: u32) {
-        self.visited.insert(i);
-    }
-
-    fn should_record_pages(&self) -> bool {
-        self.record_pages
-    }
-
-    fn record_pages(&mut self, page_id: String) {
-        match &mut self.visited_pages {
-            Some(visited_pages) => {
-                visited_pages.insert(page_id);
-            }
-            None => {}
-        }
-    }
-}
-
 pub struct Hnsw<Q: Quantizer> {
     // Need this for mmap
     #[allow(dead_code)]

rs/index/src/hnsw/utils.rs

@@ -4,6 +4,7 @@ use bit_vec::BitVec;
 use ordered_float::NotNan;
 use quantization::quantization::Quantizer;
 
+use crate::utils::TraversalContext;
 pub struct BuilderContext {
     visited: BitVec,
 }
@@ -39,13 +40,6 @@ pub struct PointAndDistance {
     pub distance: NotNan<f32>,
 }
 
-pub trait TraversalContext {
-    fn visited(&self, i: u32) -> bool;
-    fn set_visited(&mut self, i: u32);
-    fn should_record_pages(&self) -> bool;
-    fn record_pages(&mut self, page_id: String);
-}
-
 /// Move the traversal logic out, since it's used in both indexing and query path
 pub trait GraphTraversal<Q: Quantizer> {
     type ContextT: TraversalContext;

rs/index/src/utils.rs

@@ -35,6 +35,36 @@ impl SearchContext {
     }
 }
 
+pub trait TraversalContext {
+    fn visited(&self, i: u32) -> bool;
+    fn set_visited(&mut self, i: u32);
+    fn should_record_pages(&self) -> bool;
+    fn record_pages(&mut self, page_id: String);
+}
+
+impl TraversalContext for SearchContext {
+    fn visited(&self, i: u32) -> bool {
+        self.visited.contains(i)
+    }
+
+    fn set_visited(&mut self, i: u32) {
+        self.visited.insert(i);
+    }
+
+    fn should_record_pages(&self) -> bool {
+        self.record_pages
+    }
+
+    fn record_pages(&mut self, page_id: String) {
+        match &mut self.visited_pages {
+            Some(visited_pages) => {
+                visited_pages.insert(page_id);
+            }
+            None => {}
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct IdWithScore {
     pub id: u64,

rs/index/src/vector/fixed_file.rs

@@ -5,8 +5,7 @@ use memmap2::Mmap;
 use num_traits::ToBytes;
 use utils::mem::transmute_u8_to_slice;
 
-use crate::hnsw::utils::TraversalContext;
-use crate::utils::SearchContext;
+use crate::utils::{SearchContext, TraversalContext};
 
 pub struct FixedFileVectorStorage<T> {
     _marker: PhantomData<T>,