use std::collections::BinaryHeap; use std::iter::FromIterator; use common::fixed_length_priority_queue::FixedLengthPriorityQueue; use common::types::{ScoreType, ScoredPointOffset}; use num_traits::float::FloatCore; /// Structure that holds context of the search pub struct SearchContext { /// Overall nearest points found so far pub nearest: FixedLengthPriorityQueue, /// Current candidates to process pub candidates: BinaryHeap, } impl SearchContext { pub fn new(entry_point: ScoredPointOffset, ef: usize) -> Self { let mut nearest = FixedLengthPriorityQueue::new(ef); nearest.push(entry_point); SearchContext { nearest, candidates: BinaryHeap::from_iter([entry_point]), } } pub fn lower_bound(&self) -> ScoreType { match self.nearest.top() { None => ScoreType::min_value(), Some(worst_of_the_best) => worst_of_the_best.score, } } /// Updates search context with new scored point. /// If it is closer than existing - also add it to candidates for further search pub fn process_candidate(&mut self, score_point: ScoredPointOffset) { let was_added = match self.nearest.push(score_point) { None => true, Some(removed) => removed.idx != score_point.idx, }; if was_added { self.candidates.push(score_point); } } }