Skip to main content

hermes_core/query/
scoring.rs

1//! Shared scoring abstractions for text and sparse vector search
2//!
3//! Provides common types and executors for efficient top-k retrieval:
4//! - `TermCursor`: Unified cursor for both BM25 text and sparse vector posting lists
5//! - `ScoreCollector`: Efficient min-heap for maintaining top-k results
6//! - `MaxScoreExecutor`: Unified Block-Max MaxScore with conjunction optimization
7//! - `ScoredDoc`: Result type with doc_id, score, and ordinal
8
9use std::cmp::Ordering;
10use std::collections::BinaryHeap;
11
12use log::{debug, warn};
13
14use crate::DocId;
15
16/// Avoid eagerly reserving an arbitrarily large top-k heap. Most searches
17/// return far fewer hits than a very large requested limit, so let the heap
18/// grow on demand beyond this point.
19const MAX_INITIAL_SCORE_COLLECTOR_CAPACITY: usize = 8 * 1024;
20
21/// Entry for top-k min-heap
22#[derive(Clone, Copy)]
23pub struct HeapEntry {
24    pub doc_id: DocId,
25    pub score: f32,
26    pub ordinal: u16,
27}
28
29impl PartialEq for HeapEntry {
30    fn eq(&self, other: &Self) -> bool {
31        self.score.to_bits() == other.score.to_bits()
32            && self.doc_id == other.doc_id
33            && self.ordinal == other.ordinal
34    }
35}
36
37impl Eq for HeapEntry {}
38
39impl Ord for HeapEntry {
40    fn cmp(&self, other: &Self) -> Ordering {
41        // Min-heap: lower scores come first (to be evicted).
42        // total_cmp is branchless (compiles to a single comparison instruction).
43        other
44            .score
45            .total_cmp(&self.score)
46            .then_with(|| self.doc_id.cmp(&other.doc_id))
47            .then_with(|| self.ordinal.cmp(&other.ordinal))
48    }
49}
50
51impl PartialOrd for HeapEntry {
52    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
53        Some(self.cmp(other))
54    }
55}
56
57/// Efficient top-k collector using min-heap (internal, scoring-layer)
58///
59/// Maintains the k highest-scoring documents using a min-heap where the
60/// lowest score is at the top for O(1) threshold lookup and O(log k) eviction.
61/// No deduplication — caller must ensure each doc_id is inserted only once.
62///
63/// This is intentionally separate from `TopKCollector` in `collector.rs`:
64/// `ScoreCollector` is used inside `MaxScoreExecutor` where only `(doc_id,
65/// score, ordinal)` tuples exist — no `Scorer` trait, no position tracking,
66/// and the threshold must be inlined for tight block-max loops.
67/// `TopKCollector` wraps a `Scorer` and drives the full `DocSet`/`Scorer`
68/// protocol, collecting positions on demand.
69pub struct ScoreCollector {
70    /// Min-heap of top-k entries (lowest score at top for eviction)
71    heap: BinaryHeap<HeapEntry>,
72    pub k: usize,
73    /// Cached threshold: avoids repeated heap.peek() in hot loops.
74    /// Updated only when the heap changes (insert/pop).
75    cached_threshold: f32,
76    /// Score of the logical sentinel filling every unused top-k slot after
77    /// threshold seeding. Keeping one score here instead of `k - heap.len()`
78    /// entries makes filling those unused slots O(1) time and memory.
79    virtual_threshold: Option<f32>,
80}
81
82impl ScoreCollector {
83    /// Create a new collector for top-k results
84    pub fn new(k: usize) -> Self {
85        Self {
86            heap: BinaryHeap::with_capacity(k.min(MAX_INITIAL_SCORE_COLLECTOR_CAPACITY)),
87            k,
88            cached_threshold: 0.0,
89            virtual_threshold: None,
90        }
91    }
92
93    /// Current score threshold (minimum score to enter top-k)
94    #[inline]
95    pub fn threshold(&self) -> f32 {
96        self.cached_threshold
97    }
98
99    /// Recompute cached threshold from heap state
100    #[inline]
101    fn update_threshold(&mut self) {
102        self.cached_threshold = if let Some(threshold) = self.virtual_threshold {
103            threshold
104        } else if self.heap.len() >= self.k {
105            self.heap.peek().map(|e| e.score).unwrap_or(0.0)
106        } else {
107            0.0
108        };
109    }
110
111    /// Insert a document score. Returns true if inserted in top-k.
112    /// Caller must ensure each doc_id is inserted only once.
113    #[inline]
114    pub fn insert(&mut self, doc_id: DocId, score: f32) -> bool {
115        self.insert_with_ordinal(doc_id, score, 0)
116    }
117
118    /// Insert a document score with ordinal. Returns true if inserted in top-k.
119    /// Caller must ensure each doc_id is inserted only once.
120    #[inline]
121    pub fn insert_with_ordinal(&mut self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
122        if self.k == 0 {
123            return false;
124        }
125        let entry = HeapEntry {
126            doc_id,
127            score,
128            ordinal,
129        };
130        if self.heap.len() < self.k {
131            if let Some(threshold) = self.virtual_threshold {
132                let sentinel = HeapEntry {
133                    doc_id: u32::MAX,
134                    score: threshold,
135                    ordinal: 0,
136                };
137                if entry >= sentinel {
138                    return false;
139                }
140            }
141
142            self.heap.push(entry);
143            // The final real entry displaces the last virtual sentinel.
144            if self.heap.len() == self.k {
145                self.virtual_threshold = None;
146                self.update_threshold();
147            }
148            true
149        } else if self.heap.peek().is_some_and(|worst| entry < *worst) {
150            {
151                let mut worst = self.heap.peek_mut().expect("full heap has a root");
152                *worst = entry;
153            }
154            self.update_threshold();
155            true
156        } else {
157            false
158        }
159    }
160
161    /// Check if a score could potentially enter top-k
162    #[inline]
163    pub fn would_enter(&self, score: f32) -> bool {
164        self.len() < self.k || score > self.cached_threshold
165    }
166
167    /// Check whether this fully identified candidate ranks ahead of the current
168    /// worst retained entry, including deterministic tie breaks.
169    #[inline]
170    pub fn would_enter_candidate(&self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
171        if self.k == 0 {
172            return false;
173        }
174        let entry = HeapEntry {
175            doc_id,
176            score,
177            ordinal,
178        };
179        if let Some(threshold) = self.virtual_threshold {
180            let sentinel = HeapEntry {
181                doc_id: u32::MAX,
182                score: threshold,
183                ordinal: 0,
184            };
185            entry < sentinel
186        } else {
187            self.heap.len() < self.k || self.heap.peek().is_some_and(|worst| entry < *worst)
188        }
189    }
190
191    /// Get the conceptual heap length, including virtual threshold sentinels.
192    #[inline]
193    pub fn len(&self) -> usize {
194        if self.virtual_threshold.is_some() {
195            self.k
196        } else {
197            self.heap.len()
198        }
199    }
200
201    /// Number of real results retained, excluding threshold sentinels.
202    #[inline]
203    pub fn real_len(&self) -> usize {
204        self.heap.len()
205    }
206
207    /// Check if collector is empty
208    #[inline]
209    pub fn is_empty(&self) -> bool {
210        self.len() == 0
211    }
212
213    /// Seed the threshold from a cross-segment shared value.
214    ///
215    /// Logically fills unused slots and replaces retained entries below the new
216    /// floor with virtual dummy entries. This can be called repeatedly while
217    /// another segment raises the shared threshold; equal-scoring real
218    /// candidates win the deterministic doc-id tie break over sentinels.
219    pub fn seed_threshold(&mut self, initial_threshold: f32) {
220        if initial_threshold <= 0.0
221            || self.k == 0
222            || (self.len() >= self.k && initial_threshold <= self.cached_threshold)
223        {
224            return;
225        }
226
227        let sentinel = HeapEntry {
228            doc_id: u32::MAX,
229            score: initial_threshold,
230            ordinal: 0,
231        };
232
233        // When unused slots are already represented by a virtual sentinel, a
234        // new seed only changes the heap if it outranks the old floor. With an
235        // all-real full heap, it must similarly outrank the current root.
236        if let Some(current_threshold) = self.virtual_threshold {
237            let current = HeapEntry {
238                doc_id: u32::MAX,
239                score: current_threshold,
240                ordinal: 0,
241            };
242            if sentinel >= current {
243                return;
244            }
245        } else if self.heap.len() >= self.k
246            && !self.heap.peek().is_some_and(|worst| sentinel < *worst)
247        {
248            return;
249        }
250
251        self.virtual_threshold = Some(initial_threshold);
252        while self.heap.peek().is_some_and(|worst| sentinel < *worst) {
253            self.heap.pop();
254        }
255        self.update_threshold();
256    }
257
258    /// Convert to sorted top-k results (descending by score).
259    /// Filters out sentinel entries (doc_id == u32::MAX) from threshold seeding.
260    pub fn into_sorted_results(self) -> Vec<(DocId, f32, u16)> {
261        let mut results: Vec<(DocId, f32, u16)> = self
262            .heap
263            .into_vec()
264            .into_iter()
265            .filter(|e| e.doc_id != u32::MAX)
266            .map(|e| (e.doc_id, e.score, e.ordinal))
267            .collect();
268
269        // Sort by score descending, then doc_id ascending
270        results.sort_unstable_by(|a, b| {
271            b.1.total_cmp(&a.1)
272                .then_with(|| a.0.cmp(&b.0))
273                .then_with(|| a.2.cmp(&b.2))
274        });
275
276        results
277    }
278}
279
280/// Cross-segment top-k score floor, shared across the parallel/concurrent
281/// per-segment searches of a single query.
282///
283/// Stores an `f32` as raw bits in an atomic so it can be read and monotonically
284/// raised from many threads without a lock. Each segment reads the current
285/// floor as its initial pruning threshold (`ScorerOptions::initial_threshold`)
286/// and, once it has collected a *full* top-k of its own, raises the floor to
287/// its k-th score.
288///
289/// Safety of seeding: a segment only raises the floor after filling its own
290/// heap, so a floor value `v` is always backed by at least `k` real documents
291/// scoring `>= v`. The final merged k-th score is therefore `>= v`, and seeding
292/// any other segment with `v` can never drop a document that belongs in the
293/// final top-k. Completion order is arbitrary, so the floor is best-effort — it
294/// only changes how aggressively later segments prune, never correctness.
295///
296/// The floor carries the query's result-window depth `k` (`for_limit`).
297/// Publishing from a heap shallower than `k` is invalid — a segment with
298/// fewer documents than the window fills its clamped heap early, and its
299/// heap threshold says nothing about the query-global k-th score. Executors
300/// must check `SharedThreshold::covers` before raising the floor with a
301/// full-heap threshold.
302#[derive(Clone, Debug)]
303pub struct SharedThreshold {
304    floor: std::sync::Arc<std::sync::atomic::AtomicU32>,
305    /// Result-window depth the floor is valid for. `usize::MAX` means the
306    /// depth is unknown; reading stays safe, publishing is disabled.
307    k: usize,
308}
309
310impl Default for SharedThreshold {
311    fn default() -> Self {
312        Self::new()
313    }
314}
315
316impl SharedThreshold {
317    /// A fresh floor of 0.0 (no pruning seed) with an unknown window depth.
318    /// Executors can read and manually raise it, but never publish their own
319    /// full-heap thresholds into it.
320    pub fn new() -> Self {
321        Self::with_depth(usize::MAX)
322    }
323
324    /// A fresh floor valid for a query fetching `limit` results.
325    pub fn for_limit(limit: usize) -> Self {
326        Self::with_depth(limit)
327    }
328
329    fn with_depth(k: usize) -> Self {
330        Self {
331            // 0.0_f32.to_bits() == 0, matching AtomicU32::default().
332            floor: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
333            k,
334        }
335    }
336
337    /// True when a full heap of `heap_depth` distinct documents backs a valid
338    /// query-global floor for this threshold's result window.
339    #[inline]
340    pub(crate) fn covers(&self, heap_depth: usize) -> bool {
341        heap_depth >= self.k
342    }
343
344    /// Current floor.
345    #[inline]
346    pub fn get(&self) -> f32 {
347        f32::from_bits(self.floor.load(std::sync::atomic::Ordering::Relaxed))
348    }
349
350    /// Raise the floor to `score` if it is strictly higher. Monotonic; a lower
351    /// or non-positive `score` is ignored. Scores here are BM25/sparse and thus
352    /// non-negative, but the comparison is done on `f32` values (not raw bits)
353    /// so it stays correct regardless.
354    pub fn raise(&self, score: f32) {
355        // Ignore non-positive scores; a NaN falls through harmlessly (the CAS
356        // loop condition below is false for NaN, so nothing is stored).
357        if score <= 0.0 {
358            return;
359        }
360        use std::sync::atomic::Ordering::Relaxed;
361        let bits = score.to_bits();
362        let mut cur = self.floor.load(Relaxed);
363        while f32::from_bits(cur) < score {
364            match self
365                .floor
366                .compare_exchange_weak(cur, bits, Relaxed, Relaxed)
367            {
368                Ok(_) => break,
369                Err(actual) => cur = actual,
370            }
371        }
372    }
373}
374
375/// Search result from MaxScore execution
376#[derive(Debug, Clone, Copy)]
377pub struct ScoredDoc {
378    pub doc_id: DocId,
379    pub score: f32,
380    /// Ordinal for multi-valued fields (which vector in the field matched)
381    pub ordinal: u16,
382}
383
384/// Unified Block-Max MaxScore executor for top-k retrieval
385///
386/// Works with both full-text (BM25) and sparse vector (dot product) queries
387/// through the polymorphic `TermCursor`. Combines three optimizations:
388/// 1. **MaxScore partitioning** (Turtle & Flood 1995): terms split into essential
389///    (must check) and non-essential (only scored if candidate is promising)
390/// 2. **Block-max pruning** (Ding & Suel 2011): skip blocks where per-block
391///    upper bounds can't beat the current threshold
392/// 3. **Conjunction optimization** (Lucene/Grand 2023): progressively intersect
393///    essential terms as threshold rises, skipping docs that lack enough terms
394pub struct MaxScoreExecutor<'a> {
395    /// Metric labels (index, field) — set via `with_metric_labels`; empty
396    /// strings render as "unknown"/"?" is avoided by callers passing real
397    /// names from the schema.
398    metric_index: &'a str,
399    metric_field: &'a str,
400    cursors: Vec<TermCursor<'a>>,
401    prefix_sums: Vec<f32>,
402    collector: ScoreCollector,
403    inv_heap_factor: f32,
404    predicate: Option<super::DocPredicate<'a>>,
405}
406
407/// Unified term cursor for Block-Max MaxScore execution.
408///
409/// All per-position decode buffers (`doc_ids`, `scores`, `ordinals`) live in
410/// the struct directly and are filled by `ensure_block_loaded`.
411///
412/// Skip-list metadata is **not** materialized — it is read lazily from the
413/// underlying source (`BlockPostingList` for text, `SparseIndex` for sparse),
414/// both backed by zero-copy mmap'd `OwnedBytes`.
415pub(crate) struct TermCursor<'a> {
416    pub max_score: f32,
417    num_blocks: usize,
418    // ── Per-position state (filled by ensure_block_loaded) ──────────
419    block_idx: usize,
420    doc_ids: Vec<u32>,
421    scores: Vec<f32>,
422    ordinals: Vec<u16>,
423    pos: usize,
424    block_loaded: bool,
425    exhausted: bool,
426    // ── Lazy ordinal decode (sparse only) ───────────────────────────
427    /// When true, ordinal decode is deferred until ordinal_mut() is called.
428    /// Set to true for MaxScoreExecutor cursors (most blocks never need ordinals).
429    lazy_ordinals: bool,
430    /// Whether ordinals have been decoded for the current block.
431    ordinals_loaded: bool,
432    /// Stored sparse block for deferred ordinal decode (cheap Arc clone of mmap data).
433    current_sparse_block: Option<crate::structures::SparseBlock>,
434    // ── Block decode + skip access source ───────────────────────────
435    variant: CursorVariant<'a>,
436}
437
438enum CursorVariant<'a> {
439    /// Full-text BM25 — in-memory BlockPostingList (skip list + block data)
440    Text {
441        list: crate::structures::BlockPostingList,
442        idf: f32,
443        /// Precomputed: idf * (BM25_K1 + 1.0) — numerator scale factor
444        idf_times_k1_plus_1: f32,
445        /// Precomputed: 1.0 + BM25_K1 * (BM25_B / avg_field_len) — denominator tf coefficient
446        denom_tf_coeff: f32,
447        /// Precomputed: BM25_K1 * (1.0 - BM25_B) — denominator constant
448        denom_const: f32,
449        tfs: Vec<u32>,
450        /// Deferred TF decode state: (block_offset, tf_start, count).
451        /// Set when doc_ids are decoded but TFs/scores are not yet computed.
452        deferred_tf: Option<(usize, usize, usize)>,
453    },
454    /// Sparse vector — mmap'd SparseIndex (skip entries + block data)
455    Sparse {
456        si: &'a crate::segment::SparseIndex,
457        query_weight: f32,
458        skip_start: usize,
459        block_data_offset: u64,
460    },
461}
462
463// ── TermCursor async/sync macros ──────────────────────────────────────────
464//
465// Parameterised on:
466//   $load_block_fn – load_block_direct | load_block_direct_sync  (sparse I/O)
467//   $ensure_fn     – ensure_block_loaded | ensure_block_loaded_sync
468//   $($aw)*        – .await  (present for async, absent for sync)
469
470macro_rules! cursor_ensure_block {
471    ($self:ident, $load_block_fn:ident, $($aw:tt)*) => {{
472        if $self.exhausted || $self.block_loaded {
473            return Ok(!$self.exhausted);
474        }
475        match &mut $self.variant {
476            CursorVariant::Text {
477                list,
478                deferred_tf,
479                ..
480            } => {
481                if let Some(state) = list.decode_block_doc_ids_only($self.block_idx, &mut $self.doc_ids) {
482                    *deferred_tf = Some(state);
483                    $self.scores.clear();
484                    $self.pos = 0;
485                    $self.block_loaded = true;
486                    Ok(true)
487                } else {
488                    $self.exhausted = true;
489                    Ok(false)
490                }
491            }
492            CursorVariant::Sparse {
493                si,
494                query_weight,
495                skip_start,
496                block_data_offset,
497                ..
498            } => {
499                let block = si
500                    .$load_block_fn(*skip_start, *block_data_offset, $self.block_idx)
501                    $($aw)* ?;
502                match block {
503                    Some(b) => {
504                        b.decode_doc_ids_into(&mut $self.doc_ids);
505                        b.decode_scored_weights_into(*query_weight, &mut $self.scores);
506                        if $self.lazy_ordinals {
507                            // Defer ordinal decode until ordinal_mut() is called.
508                            // Stores cheap Arc-backed mmap slice, no copy.
509                            $self.current_sparse_block = Some(b);
510                            $self.ordinals_loaded = false;
511                        } else {
512                            b.decode_ordinals_into(&mut $self.ordinals);
513                            $self.ordinals_loaded = true;
514                            $self.current_sparse_block = None;
515                        }
516                        $self.pos = 0;
517                        $self.block_loaded = true;
518                        Ok(true)
519                    }
520                    None => {
521                        $self.exhausted = true;
522                        Ok(false)
523                    }
524                }
525            }
526        }
527    }};
528}
529
530macro_rules! cursor_advance {
531    ($self:ident, $ensure_fn:ident, $($aw:tt)*) => {{
532        if $self.exhausted {
533            return Ok(u32::MAX);
534        }
535        $self.$ensure_fn() $($aw)* ?;
536        if $self.exhausted {
537            return Ok(u32::MAX);
538        }
539        Ok($self.advance_pos())
540    }};
541}
542
543macro_rules! cursor_seek {
544    ($self:ident, $ensure_fn:ident, $target:expr, $($aw:tt)*) => {{
545        if let Some(doc) = $self.seek_prepare($target) {
546            return Ok(doc);
547        }
548        $self.$ensure_fn() $($aw)* ?;
549        if $self.seek_finish($target) {
550            $self.$ensure_fn() $($aw)* ?;
551        }
552        Ok($self.doc())
553    }};
554}
555
556impl<'a> TermCursor<'a> {
557    /// Create a full-text BM25 cursor (lazy — no blocks decoded yet).
558    pub fn text(
559        posting_list: crate::structures::BlockPostingList,
560        idf: f32,
561        avg_field_len: f32,
562    ) -> Self {
563        let max_tf = posting_list.max_tf() as f32;
564        let max_score = super::bm25_upper_bound(max_tf.max(1.0), idf);
565        let num_blocks = posting_list.num_blocks();
566        let safe_avg = avg_field_len.max(1.0);
567        Self {
568            max_score,
569            num_blocks,
570            block_idx: 0,
571            doc_ids: Vec::with_capacity(128),
572            scores: Vec::with_capacity(128),
573            ordinals: Vec::new(),
574            pos: 0,
575            block_loaded: false,
576            exhausted: num_blocks == 0,
577            lazy_ordinals: false,
578            ordinals_loaded: true, // text cursors never have ordinals
579            current_sparse_block: None,
580            variant: CursorVariant::Text {
581                list: posting_list,
582                idf,
583                idf_times_k1_plus_1: idf * (super::BM25_K1 + 1.0),
584                denom_tf_coeff: 1.0 + super::BM25_K1 * (super::BM25_B / safe_avg),
585                denom_const: super::BM25_K1 * (1.0 - super::BM25_B),
586                tfs: Vec::with_capacity(128),
587                deferred_tf: None,
588            },
589        }
590    }
591
592    /// Create a sparse vector cursor with lazy block loading.
593    /// Skip entries are **not** copied — they are read from `SparseIndex` mmap on demand.
594    pub fn sparse(
595        si: &'a crate::segment::SparseIndex,
596        query_weight: f32,
597        skip_start: usize,
598        skip_count: usize,
599        global_max_weight: f32,
600        block_data_offset: u64,
601    ) -> Self {
602        Self {
603            max_score: query_weight.abs() * global_max_weight,
604            num_blocks: skip_count,
605            block_idx: 0,
606            doc_ids: Vec::with_capacity(256),
607            scores: Vec::with_capacity(256),
608            ordinals: Vec::with_capacity(256),
609            pos: 0,
610            block_loaded: false,
611            exhausted: skip_count == 0,
612            lazy_ordinals: false,
613            ordinals_loaded: true,
614            current_sparse_block: None,
615            variant: CursorVariant::Sparse {
616                si,
617                query_weight,
618                skip_start,
619                block_data_offset,
620            },
621        }
622    }
623
624    // ── Skip-entry access (lazy, zero-copy for sparse) ──────────────────
625
626    #[inline]
627    fn block_first_doc(&self, idx: usize) -> DocId {
628        match &self.variant {
629            CursorVariant::Text { list, .. } => list.block_first_doc(idx).unwrap_or(u32::MAX),
630            CursorVariant::Sparse { si, skip_start, .. } => {
631                si.read_skip_entry(*skip_start + idx).first_doc
632            }
633        }
634    }
635
636    #[inline]
637    fn block_last_doc(&self, idx: usize) -> DocId {
638        match &self.variant {
639            CursorVariant::Text { list, .. } => list.block_last_doc(idx).unwrap_or(0),
640            CursorVariant::Sparse { si, skip_start, .. } => {
641                si.read_skip_entry(*skip_start + idx).last_doc
642            }
643        }
644    }
645
646    // ── Read-only accessors ─────────────────────────────────────────────
647
648    #[inline]
649    pub fn doc(&self) -> DocId {
650        if self.exhausted {
651            return u32::MAX;
652        }
653        if self.block_loaded {
654            debug_assert!(self.pos < self.doc_ids.len());
655            // SAFETY: pos < doc_ids.len() is maintained by advance_pos/ensure_block_loaded.
656            unsafe { *self.doc_ids.get_unchecked(self.pos) }
657        } else {
658            self.block_first_doc(self.block_idx)
659        }
660    }
661
662    #[inline]
663    pub fn ordinal(&self) -> u16 {
664        if !self.block_loaded || self.ordinals.is_empty() {
665            return 0;
666        }
667        debug_assert!(self.pos < self.ordinals.len());
668        // SAFETY: pos < ordinals.len() is maintained by advance_pos/ensure_block_loaded.
669        unsafe { *self.ordinals.get_unchecked(self.pos) }
670    }
671
672    /// Lazily-decoded ordinal accessor for MaxScore executor.
673    ///
674    /// When `lazy_ordinals=true`, ordinals are not decoded during block loading.
675    /// This method triggers the deferred decode on first access, amortized over
676    /// the block. Subsequent calls within the same block are free.
677    #[inline]
678    pub fn ordinal_mut(&mut self) -> u16 {
679        if !self.block_loaded {
680            return 0;
681        }
682        if !self.ordinals_loaded {
683            if let Some(ref block) = self.current_sparse_block {
684                block.decode_ordinals_into(&mut self.ordinals);
685            }
686            self.ordinals_loaded = true;
687        }
688        if self.ordinals.is_empty() {
689            return 0;
690        }
691        debug_assert!(self.pos < self.ordinals.len());
692        unsafe { *self.ordinals.get_unchecked(self.pos) }
693    }
694
695    #[inline]
696    pub fn score(&self) -> f32 {
697        if !self.block_loaded {
698            return 0.0;
699        }
700        debug_assert!(self.pos < self.scores.len());
701        // SAFETY: pos < scores.len() is maintained by advance_pos/ensure_block_loaded.
702        unsafe { *self.scores.get_unchecked(self.pos) }
703    }
704
705    /// Ensure BM25 scores are computed for the current block (lazy TF decode).
706    ///
707    /// For text cursors, TF unpacking and BM25 scoring are deferred from block
708    /// loading until this method is called, saving work for blocks skipped by
709    /// block-max or conjunction pruning. No-op for sparse cursors.
710    #[inline]
711    pub fn ensure_scores(&mut self) {
712        if self.block_loaded && self.scores.is_empty() {
713            self.compute_deferred_scores();
714        }
715    }
716
717    #[inline]
718    pub fn current_block_max_score(&self) -> f32 {
719        if self.exhausted {
720            return 0.0;
721        }
722        match &self.variant {
723            CursorVariant::Text { list, idf, .. } => {
724                let block_max_tf = list.block_max_tf(self.block_idx).unwrap_or(0) as f32;
725                super::bm25_upper_bound(block_max_tf.max(1.0), *idf)
726            }
727            CursorVariant::Sparse {
728                si,
729                query_weight,
730                skip_start,
731                ..
732            } => query_weight.abs() * si.read_skip_entry(*skip_start + self.block_idx).max_weight,
733        }
734    }
735
736    // ── Block navigation ────────────────────────────────────────────────
737
738    pub fn skip_to_next_block(&mut self) -> DocId {
739        if self.exhausted {
740            return u32::MAX;
741        }
742        self.block_idx += 1;
743        self.block_loaded = false;
744        if self.block_idx >= self.num_blocks {
745            self.exhausted = true;
746            return u32::MAX;
747        }
748        self.block_first_doc(self.block_idx)
749    }
750
751    #[inline]
752    fn advance_pos(&mut self) -> DocId {
753        self.pos += 1;
754        if self.pos >= self.doc_ids.len() {
755            self.block_idx += 1;
756            self.block_loaded = false;
757            if self.block_idx >= self.num_blocks {
758                self.exhausted = true;
759                return u32::MAX;
760            }
761        }
762        self.doc()
763    }
764
765    /// Compute BM25 scores from deferred TF data (lazy decode for text cursors).
766    #[inline(never)]
767    fn compute_deferred_scores(&mut self) {
768        if let CursorVariant::Text {
769            list,
770            idf_times_k1_plus_1,
771            denom_tf_coeff,
772            denom_const,
773            tfs,
774            deferred_tf,
775            ..
776        } = &mut self.variant
777            && let Some((block_offset, tf_start, count)) = deferred_tf.take()
778        {
779            list.decode_block_tfs_deferred(block_offset, tf_start, count, tfs);
780            let num_scale = *idf_times_k1_plus_1;
781            let d_tf = *denom_tf_coeff;
782            let d_const = *denom_const;
783            self.scores.clear();
784            self.scores.resize(count, 0.0);
785            for i in 0..count {
786                let tf = unsafe { *tfs.get_unchecked(i) } as f32;
787                let score = (num_scale * tf) / (d_tf * tf + d_const);
788                unsafe {
789                    *self.scores.get_unchecked_mut(i) = score;
790                }
791            }
792        }
793    }
794
795    // ── Block loading / advance / seek ─────────────────────────────────
796    //
797    // Macros parameterised on sparse I/O method + optional .await to
798    // stamp out both async and sync variants without duplication.
799
800    pub async fn ensure_block_loaded(&mut self) -> crate::Result<bool> {
801        cursor_ensure_block!(self, load_block_direct, .await)
802    }
803
804    pub fn ensure_block_loaded_sync(&mut self) -> crate::Result<bool> {
805        cursor_ensure_block!(self, load_block_direct_sync,)
806    }
807
808    pub async fn advance(&mut self) -> crate::Result<DocId> {
809        cursor_advance!(self, ensure_block_loaded, .await)
810    }
811
812    pub fn advance_sync(&mut self) -> crate::Result<DocId> {
813        cursor_advance!(self, ensure_block_loaded_sync,)
814    }
815
816    pub async fn seek(&mut self, target: DocId) -> crate::Result<DocId> {
817        cursor_seek!(self, ensure_block_loaded, target, .await)
818    }
819
820    pub fn seek_sync(&mut self, target: DocId) -> crate::Result<DocId> {
821        cursor_seek!(self, ensure_block_loaded_sync, target,)
822    }
823
824    fn seek_prepare(&mut self, target: DocId) -> Option<DocId> {
825        if self.exhausted {
826            return Some(u32::MAX);
827        }
828
829        // Fast path: target is within the currently loaded block
830        if self.block_loaded
831            && let Some(&last) = self.doc_ids.last()
832        {
833            if last >= target && self.doc_ids[self.pos] < target {
834                let remaining = &self.doc_ids[self.pos..];
835                self.pos += crate::structures::simd::find_first_ge_u32(remaining, target);
836                if self.pos >= self.doc_ids.len() {
837                    self.block_idx += 1;
838                    self.block_loaded = false;
839                    if self.block_idx >= self.num_blocks {
840                        self.exhausted = true;
841                        return Some(u32::MAX);
842                    }
843                }
844                return Some(self.doc());
845            }
846            if self.doc_ids[self.pos] >= target {
847                return Some(self.doc());
848            }
849        }
850
851        // Seek to the block containing target
852        let lo = match &self.variant {
853            // Text: SIMD-accelerated 2-level seek (L1 + L0)
854            CursorVariant::Text { list, .. } => match list.seek_block(target, self.block_idx) {
855                Some(idx) => idx,
856                None => {
857                    self.exhausted = true;
858                    return Some(u32::MAX);
859                }
860            },
861            // Sparse: binary search on skip entries (lazy mmap reads)
862            CursorVariant::Sparse { .. } => {
863                let mut lo = self.block_idx;
864                let mut hi = self.num_blocks;
865                while lo < hi {
866                    let mid = lo + (hi - lo) / 2;
867                    if self.block_last_doc(mid) < target {
868                        lo = mid + 1;
869                    } else {
870                        hi = mid;
871                    }
872                }
873                lo
874            }
875        };
876        if lo >= self.num_blocks {
877            self.exhausted = true;
878            return Some(u32::MAX);
879        }
880        if lo != self.block_idx || !self.block_loaded {
881            self.block_idx = lo;
882            self.block_loaded = false;
883        }
884        None
885    }
886
887    #[inline]
888    fn seek_finish(&mut self, target: DocId) -> bool {
889        if self.exhausted {
890            return false;
891        }
892        self.pos = crate::structures::simd::find_first_ge_u32(&self.doc_ids, target);
893        if self.pos >= self.doc_ids.len() {
894            self.block_idx += 1;
895            self.block_loaded = false;
896            if self.block_idx >= self.num_blocks {
897                self.exhausted = true;
898                return false;
899            }
900            return true;
901        }
902        false
903    }
904}
905
906/// Macro to stamp out the Block-Max MaxScore loop for both async and sync paths.
907///
908/// `$ensure`, `$advance`, `$seek` are cursor method idents (async or _sync variants).
909/// `$($aw:tt)*` captures `.await` for async or nothing for sync.
910macro_rules! bms_execute_loop {
911    ($self:ident, $ensure:ident, $advance:ident, $seek:ident, $($aw:tt)*) => {{
912        let n = $self.cursors.len();
913
914        // Load first block for each cursor (ensures doc() returns real values)
915        for cursor in &mut $self.cursors {
916            cursor.$ensure() $($aw)* ?;
917        }
918
919        let mut docs_scored = 0u64;
920        let mut docs_skipped = 0u64;
921        let mut blocks_skipped = 0u64;
922        let mut conjunction_skipped = 0u64;
923        let mut ordinal_scores: Vec<(u16, f32)> = Vec::with_capacity(n * 2);
924        let _bms_start = std::time::Instant::now();
925
926        let inv_heap_factor = $self.inv_heap_factor;
927        let mut adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
928
929        loop {
930            let partition = $self.find_partition();
931            if partition >= n {
932                break;
933            }
934
935            // Find minimum doc_id across essential cursors and collect
936            // which cursors are at min_doc (avoids redundant re-checks in
937            // conjunction, block-max, predicate, and scoring passes).
938            let mut min_doc = u32::MAX;
939            let mut at_min_mask = 0u64; // bitset of cursor indices at min_doc
940            for i in partition..n {
941                let doc = $self.cursors[i].doc();
942                match doc.cmp(&min_doc) {
943                    std::cmp::Ordering::Less => {
944                        min_doc = doc;
945                        at_min_mask = 1u64 << (i as u32);
946                    }
947                    std::cmp::Ordering::Equal => {
948                        at_min_mask |= 1u64 << (i as u32);
949                    }
950                    _ => {}
951                }
952            }
953            if min_doc == u32::MAX {
954                break;
955            }
956
957            let non_essential_upper = if partition > 0 {
958                $self.prefix_sums[partition - 1]
959            } else {
960                0.0
961            };
962
963            // --- Conjunction optimization ---
964            if $self.collector.len() >= $self.collector.k {
965                let mut present_upper: f32 = 0.0;
966                let mut mask = at_min_mask;
967                while mask != 0 {
968                    let i = mask.trailing_zeros() as usize;
969                    present_upper += $self.cursors[i].max_score;
970                    mask &= mask - 1;
971                }
972
973                if present_upper + non_essential_upper < adjusted_threshold {
974                    let mut mask = at_min_mask;
975                    while mask != 0 {
976                        let i = mask.trailing_zeros() as usize;
977                        $self.cursors[i].$ensure() $($aw)* ?;
978                        $self.cursors[i].$advance() $($aw)* ?;
979                        mask &= mask - 1;
980                    }
981                    conjunction_skipped += 1;
982                    continue;
983                }
984            }
985
986            // --- Block-max pruning ---
987            if $self.collector.len() >= $self.collector.k {
988                let mut block_max_sum: f32 = 0.0;
989                let mut mask = at_min_mask;
990                while mask != 0 {
991                    let i = mask.trailing_zeros() as usize;
992                    block_max_sum += $self.cursors[i].current_block_max_score();
993                    mask &= mask - 1;
994                }
995
996                if block_max_sum + non_essential_upper < adjusted_threshold {
997                    let mut mask = at_min_mask;
998                    while mask != 0 {
999                        let i = mask.trailing_zeros() as usize;
1000                        $self.cursors[i].skip_to_next_block();
1001                        $self.cursors[i].$ensure() $($aw)* ?;
1002                        mask &= mask - 1;
1003                    }
1004                    blocks_skipped += 1;
1005                    continue;
1006                }
1007            }
1008
1009            // --- Predicate filter (after block-max, before scoring) ---
1010            if let Some(ref pred) = $self.predicate {
1011                if !pred(min_doc) {
1012                    let mut mask = at_min_mask;
1013                    while mask != 0 {
1014                        let i = mask.trailing_zeros() as usize;
1015                        $self.cursors[i].$ensure() $($aw)* ?;
1016                        $self.cursors[i].$advance() $($aw)* ?;
1017                        mask &= mask - 1;
1018                    }
1019                    continue;
1020                }
1021            }
1022
1023            // --- Score essential cursors ---
1024            ordinal_scores.clear();
1025            {
1026                let mut mask = at_min_mask;
1027                while mask != 0 {
1028                    let i = mask.trailing_zeros() as usize;
1029                    $self.cursors[i].$ensure() $($aw)* ?;
1030                    $self.cursors[i].ensure_scores();
1031                    while $self.cursors[i].doc() == min_doc {
1032                        let ord = $self.cursors[i].ordinal_mut();
1033                        let sc = $self.cursors[i].score();
1034                        ordinal_scores.push((ord, sc));
1035                        $self.cursors[i].$advance() $($aw)* ?;
1036                    }
1037                    mask &= mask - 1;
1038                }
1039            }
1040
1041            let essential_total: f32 = ordinal_scores.iter().map(|(_, s)| *s).sum();
1042            if $self.collector.len() >= $self.collector.k
1043                && essential_total + non_essential_upper < adjusted_threshold
1044            {
1045                docs_skipped += 1;
1046                continue;
1047            }
1048
1049            // --- Score non-essential cursors (highest max_score first for early exit) ---
1050            let mut running_total = essential_total;
1051            for i in (0..partition).rev() {
1052                if $self.collector.len() >= $self.collector.k
1053                    && running_total + $self.prefix_sums[i] < adjusted_threshold
1054                {
1055                    break;
1056                }
1057
1058                let doc = $self.cursors[i].$seek(min_doc) $($aw)* ?;
1059                if doc == min_doc {
1060                    $self.cursors[i].ensure_scores();
1061                    while $self.cursors[i].doc() == min_doc {
1062                        let s = $self.cursors[i].score();
1063                        running_total += s;
1064                        let ord = $self.cursors[i].ordinal_mut();
1065                        ordinal_scores.push((ord, s));
1066                        $self.cursors[i].$advance() $($aw)* ?;
1067                    }
1068                }
1069            }
1070
1071            // --- Group by ordinal and insert ---
1072            // Fast path: single entry (common for single-valued fields) — skip sort + grouping
1073            if ordinal_scores.len() == 1 {
1074                let (ord, score) = ordinal_scores[0];
1075                if $self.collector.insert_with_ordinal(min_doc, score, ord) {
1076                    docs_scored += 1;
1077                    adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1078                } else {
1079                    docs_skipped += 1;
1080                }
1081            } else if !ordinal_scores.is_empty() {
1082                if ordinal_scores.len() > 2 {
1083                    ordinal_scores.sort_unstable_by_key(|(ord, _)| *ord);
1084                } else if ordinal_scores.len() == 2 && ordinal_scores[0].0 > ordinal_scores[1].0 {
1085                    ordinal_scores.swap(0, 1);
1086                }
1087                let mut j = 0;
1088                while j < ordinal_scores.len() {
1089                    let current_ord = ordinal_scores[j].0;
1090                    let mut score = 0.0f32;
1091                    while j < ordinal_scores.len() && ordinal_scores[j].0 == current_ord {
1092                        score += ordinal_scores[j].1;
1093                        j += 1;
1094                    }
1095                    if $self
1096                        .collector
1097                        .insert_with_ordinal(min_doc, score, current_ord)
1098                    {
1099                        docs_scored += 1;
1100                        adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1101                    } else {
1102                        docs_skipped += 1;
1103                    }
1104                }
1105            }
1106        }
1107
1108        let results: Vec<ScoredDoc> = $self
1109            .collector
1110            .into_sorted_results()
1111            .into_iter()
1112            .map(|(doc_id, score, ordinal)| ScoredDoc {
1113                doc_id,
1114                score,
1115                ordinal,
1116            })
1117            .collect();
1118
1119        let _bms_elapsed_ms = _bms_start.elapsed().as_millis() as u64;
1120        if _bms_elapsed_ms > 500 {
1121            warn!(
1122                "slow MaxScore: {}ms, cursors={}, scored={}, skipped={}, blocks_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1123                _bms_elapsed_ms,
1124                n,
1125                docs_scored,
1126                docs_skipped,
1127                blocks_skipped,
1128                conjunction_skipped,
1129                results.len(),
1130                results.first().map(|r| r.score).unwrap_or(0.0)
1131            );
1132        } else {
1133            debug!(
1134                "MaxScoreExecutor: {}ms, scored={}, skipped={}, blocks_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1135                _bms_elapsed_ms,
1136                docs_scored,
1137                docs_skipped,
1138                blocks_skipped,
1139                conjunction_skipped,
1140                results.len(),
1141                results.first().map(|r| r.score).unwrap_or(0.0)
1142            );
1143        }
1144
1145        Ok(results)
1146    }};
1147}
1148
1149impl<'a> MaxScoreExecutor<'a> {
1150    /// Create a new executor from pre-built cursors.
1151    ///
1152    /// Cursors are sorted by max_score ascending (non-essential first) and
1153    /// prefix sums are computed for the MaxScore partitioning.
1154    pub(crate) fn new(mut cursors: Vec<TermCursor<'a>>, k: usize, heap_factor: f32) -> Self {
1155        // The execution loop tracks cursors at the current document in a u64.
1156        // Query construction normally enforces this bound, but keep this
1157        // boundary defensive for direct/internal executor users as well.
1158        if cursors.len() > super::MAX_QUERY_TERMS {
1159            cursors.sort_unstable_by(|a, b| b.max_score.total_cmp(&a.max_score));
1160            cursors.truncate(super::MAX_QUERY_TERMS);
1161            log::warn!(
1162                "MaxScore cursor count exceeded {}; retaining the strongest cursors",
1163                super::MAX_QUERY_TERMS
1164            );
1165        }
1166
1167        // Enable lazy ordinal decode — ordinals are only decoded when a doc
1168        // actually reaches the scoring phase (saves ~100ns per skipped block).
1169        for c in &mut cursors {
1170            c.lazy_ordinals = true;
1171        }
1172
1173        // Sort by max_score ascending (non-essential first)
1174        cursors.sort_by(|a, b| {
1175            a.max_score
1176                .partial_cmp(&b.max_score)
1177                .unwrap_or(Ordering::Equal)
1178        });
1179
1180        let mut prefix_sums = Vec::with_capacity(cursors.len());
1181        let mut cumsum = 0.0f32;
1182        for c in &cursors {
1183            cumsum += c.max_score;
1184            prefix_sums.push(cumsum);
1185        }
1186
1187        let clamped_heap_factor = heap_factor.clamp(0.01, 1.0);
1188
1189        debug!(
1190            "Creating MaxScoreExecutor: num_cursors={}, k={}, total_upper={:.4}, heap_factor={:.2}",
1191            cursors.len(),
1192            k,
1193            cumsum,
1194            clamped_heap_factor
1195        );
1196
1197        Self {
1198            cursors,
1199            prefix_sums,
1200            collector: ScoreCollector::new(k),
1201            inv_heap_factor: 1.0 / clamped_heap_factor,
1202            predicate: None,
1203            metric_index: "unknown",
1204            metric_field: "unknown",
1205        }
1206    }
1207
1208    /// Attach (index, field) labels for the metrics this executor emits.
1209    pub fn with_metric_labels(mut self, index: &'a str, field: &'a str) -> Self {
1210        self.metric_index = index;
1211        self.metric_field = field;
1212        self
1213    }
1214
1215    /// Create an executor for sparse vector queries.
1216    ///
1217    /// Builds `TermCursor::Sparse` for each matched dimension.
1218    pub fn sparse(
1219        sparse_index: &'a crate::segment::SparseIndex,
1220        query_terms: Vec<(u32, f32)>,
1221        k: usize,
1222        heap_factor: f32,
1223    ) -> Self {
1224        let cursors: Vec<TermCursor<'a>> = query_terms
1225            .iter()
1226            .filter_map(|&(dim_id, qw)| {
1227                let (skip_start, skip_count, global_max, block_data_offset) =
1228                    sparse_index.get_skip_range_full(dim_id)?;
1229                Some(TermCursor::sparse(
1230                    sparse_index,
1231                    qw,
1232                    skip_start,
1233                    skip_count,
1234                    global_max,
1235                    block_data_offset,
1236                ))
1237            })
1238            .collect();
1239        Self::new(cursors, k, heap_factor)
1240    }
1241
1242    /// Create an executor for full-text BM25 queries.
1243    ///
1244    /// Builds `TermCursor::Text` for each posting list.
1245    pub fn text(
1246        posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1247        avg_field_len: f32,
1248        k: usize,
1249    ) -> Self {
1250        let cursors: Vec<TermCursor<'a>> = posting_lists
1251            .into_iter()
1252            .map(|(pl, idf)| TermCursor::text(pl, idf, avg_field_len))
1253            .collect();
1254        Self::new(cursors, k, 1.0)
1255    }
1256
1257    #[inline]
1258    fn find_partition(&self) -> usize {
1259        // Alpha < 1.0 raises the effective threshold → more terms become
1260        // non-essential → more aggressive pruning (approximate retrieval).
1261        // Use multiplication by reciprocal (cheaper than division).
1262        let threshold = self.collector.threshold() * self.inv_heap_factor;
1263        // Keep an equal-score candidate essential: it can still displace the
1264        // current worst hit through the deterministic doc/ordinal tie-break.
1265        self.prefix_sums.partition_point(|&sum| sum < threshold)
1266    }
1267
1268    /// Attach a per-doc predicate filter to this executor.
1269    ///
1270    /// Docs failing the predicate are skipped after block-max pruning but
1271    /// before scoring. The predicate does not affect thresholds or block-max
1272    /// comparisons — the heap stores pure sparse/text scores.
1273    pub fn with_predicate(mut self, predicate: super::DocPredicate<'a>) -> Self {
1274        self.predicate = Some(predicate);
1275        self
1276    }
1277
1278    /// Seed the collector with an initial threshold for tighter early pruning.
1279    pub fn seed_threshold(&mut self, initial_threshold: f32) {
1280        self.collector.seed_threshold(initial_threshold);
1281    }
1282
1283    /// Execute Block-Max MaxScore and return top-k results (async).
1284    pub async fn execute(mut self) -> crate::Result<Vec<ScoredDoc>> {
1285        if self.cursors.is_empty() {
1286            return Ok(Vec::new());
1287        }
1288        let t = crate::observe::Timer::start();
1289        let results = bms_execute_loop!(self, ensure_block_loaded, advance, seek, .await);
1290        if let Ok(r) = &results {
1291            crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1292        }
1293        results
1294    }
1295
1296    /// Synchronous execution — works when all cursors are text or mmap-backed sparse.
1297    pub fn execute_sync(mut self) -> crate::Result<Vec<ScoredDoc>> {
1298        if self.cursors.is_empty() {
1299            return Ok(Vec::new());
1300        }
1301        let t = crate::observe::Timer::start();
1302        let results = bms_execute_loop!(self, ensure_block_loaded_sync, advance_sync, seek_sync,);
1303        if let Ok(r) = &results {
1304            crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1305        }
1306        results
1307    }
1308}
1309
1310#[cfg(test)]
1311mod tests {
1312    use super::*;
1313
1314    #[test]
1315    fn test_shared_threshold_monotonic_raise() {
1316        let shared = SharedThreshold::new();
1317        assert_eq!(shared.get(), 0.0);
1318
1319        shared.raise(2.5);
1320        assert_eq!(shared.get(), 2.5);
1321
1322        // Lower values never lower the floor.
1323        shared.raise(1.0);
1324        assert_eq!(shared.get(), 2.5);
1325
1326        // Higher values raise it.
1327        shared.raise(4.0);
1328        assert_eq!(shared.get(), 4.0);
1329
1330        // Non-positive and NaN are ignored.
1331        shared.raise(0.0);
1332        shared.raise(-3.0);
1333        shared.raise(f32::NAN);
1334        assert_eq!(shared.get(), 4.0);
1335
1336        // Clones share the same atomic cell.
1337        let clone = shared.clone();
1338        clone.raise(9.0);
1339        assert_eq!(shared.get(), 9.0);
1340    }
1341
1342    #[test]
1343    fn test_shared_threshold_seed_matches_manual() {
1344        // A collector seeded with a floor prunes anything at/below it, matching
1345        // the threshold a fully-populated heap would have produced.
1346        let mut seeded = ScoreCollector::new(2);
1347        seeded.seed_threshold(3.0);
1348        assert_eq!(seeded.threshold(), 3.0);
1349        // A score at/below the floor cannot enter.
1350        assert!(!seeded.would_enter(3.0));
1351        assert!(seeded.would_enter(3.5));
1352        // Real inserts above the floor evict the sentinels; results contain no
1353        // sentinel (doc_id == u32::MAX) entries.
1354        seeded.insert(1, 5.0);
1355        seeded.insert(2, 4.0);
1356        let results = seeded.into_sorted_results();
1357        assert_eq!(results.len(), 2);
1358        assert_eq!(results[0].0, 1);
1359        assert_eq!(results[1].0, 2);
1360    }
1361
1362    #[test]
1363    fn test_shared_threshold_can_raise_after_real_inserts() {
1364        let mut collector = ScoreCollector::new(3);
1365        collector.insert(1, 10.0);
1366        collector.insert(2, 4.0);
1367        assert_eq!(collector.real_len(), 2);
1368
1369        // Raising the floor after traversal has started removes retained work
1370        // that can no longer reach the global top-k.
1371        collector.seed_threshold(6.0);
1372        assert_eq!(collector.threshold(), 6.0);
1373        assert_eq!(collector.real_len(), 1);
1374
1375        // A real candidate tied with the floor displaces the sentinel because
1376        // its doc id wins the canonical tie break.
1377        assert!(collector.would_enter_candidate(3, 6.0, 0));
1378        assert!(collector.insert(3, 6.0));
1379        assert_eq!(collector.real_len(), 2);
1380        let results = collector.into_sorted_results();
1381        assert_eq!(results, vec![(1, 10.0, 0), (3, 6.0, 0)]);
1382    }
1383
1384    #[test]
1385    fn test_large_seed_uses_virtual_sentinels() {
1386        let k = 1_000_000_000;
1387        let mut collector = ScoreCollector::new(k);
1388        assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
1389
1390        collector.seed_threshold(42.0);
1391
1392        // Seeding a huge top-k is constant-time and does not materialize any
1393        // of its conceptual sentinel entries.
1394        assert_eq!(collector.heap.len(), 0);
1395        assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
1396        assert_eq!(collector.len(), k);
1397        assert_eq!(collector.real_len(), 0);
1398        assert_eq!(collector.threshold(), 42.0);
1399        assert!(!collector.is_empty());
1400
1401        // A real result tied with the floor beats the sentinel by doc-id, while
1402        // a lower score remains below the conceptual threshold.
1403        assert!(collector.insert_with_ordinal(9, 42.0, 7));
1404        assert!(!collector.insert(10, 41.0));
1405        assert!(collector.insert(11, 43.0));
1406        assert_eq!(collector.len(), k);
1407        assert_eq!(collector.real_len(), 2);
1408        assert_eq!(
1409            collector.into_sorted_results(),
1410            vec![(11, 43.0, 0), (9, 42.0, 7)]
1411        );
1412    }
1413
1414    #[test]
1415    fn test_virtual_sentinels_preserve_tie_order_when_filled() {
1416        let mut collector = ScoreCollector::new(3);
1417        collector.seed_threshold(5.0);
1418
1419        assert!(collector.insert_with_ordinal(3, 5.0, 2));
1420        assert!(collector.insert_with_ordinal(2, 5.0, 8));
1421        assert!(collector.insert_with_ordinal(1, 5.0, 4));
1422        assert_eq!(collector.real_len(), 3);
1423        assert!(collector.virtual_threshold.is_none());
1424
1425        // Once all virtual slots have been displaced, canonical doc/ordinal
1426        // ordering still controls root replacement at an equal score.
1427        assert!(collector.insert_with_ordinal(2, 5.0, 1));
1428        assert!(!collector.insert_with_ordinal(4, 5.0, 0));
1429        assert_eq!(
1430            collector.into_sorted_results(),
1431            vec![(1, 5.0, 4), (2, 5.0, 1), (2, 5.0, 8)]
1432        );
1433    }
1434
1435    #[test]
1436    fn test_score_collector_basic() {
1437        let mut collector = ScoreCollector::new(3);
1438
1439        collector.insert(1, 1.0);
1440        collector.insert(2, 2.0);
1441        collector.insert(3, 3.0);
1442        assert_eq!(collector.threshold(), 1.0);
1443
1444        collector.insert(4, 4.0);
1445        assert_eq!(collector.threshold(), 2.0);
1446
1447        let results = collector.into_sorted_results();
1448        assert_eq!(results.len(), 3);
1449        assert_eq!(results[0].0, 4); // Highest score
1450        assert_eq!(results[1].0, 3);
1451        assert_eq!(results[2].0, 2);
1452    }
1453
1454    #[test]
1455    fn test_score_collector_threshold() {
1456        let mut collector = ScoreCollector::new(2);
1457
1458        collector.insert(1, 5.0);
1459        collector.insert(2, 3.0);
1460        assert_eq!(collector.threshold(), 3.0);
1461
1462        // Should not enter (score too low)
1463        assert!(!collector.would_enter(2.0));
1464        assert!(!collector.insert(3, 2.0));
1465
1466        // Should enter (score high enough)
1467        assert!(collector.would_enter(4.0));
1468        assert!(collector.insert(4, 4.0));
1469        assert_eq!(collector.threshold(), 4.0);
1470    }
1471
1472    #[test]
1473    fn test_heap_entry_ordering() {
1474        let mut heap = BinaryHeap::new();
1475        heap.push(HeapEntry {
1476            doc_id: 1,
1477            score: 3.0,
1478            ordinal: 0,
1479        });
1480        heap.push(HeapEntry {
1481            doc_id: 2,
1482            score: 1.0,
1483            ordinal: 0,
1484        });
1485        heap.push(HeapEntry {
1486            doc_id: 3,
1487            score: 2.0,
1488            ordinal: 0,
1489        });
1490
1491        // Min-heap: lowest score should come out first
1492        assert_eq!(heap.pop().unwrap().score, 1.0);
1493        assert_eq!(heap.pop().unwrap().score, 2.0);
1494        assert_eq!(heap.pop().unwrap().score, 3.0);
1495    }
1496}