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        // Keep a total float order and deterministic doc/ordinal tie breaks.
43        // The generic then_with closures inline; no callback allocation is needed.
44        other
45            .score
46            .total_cmp(&self.score)
47            .then_with(|| self.doc_id.cmp(&other.doc_id))
48            .then_with(|| self.ordinal.cmp(&other.ordinal))
49    }
50}
51
52impl PartialOrd for HeapEntry {
53    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
54        Some(self.cmp(other))
55    }
56}
57
58/// Efficient top-k collector using min-heap (internal, scoring-layer)
59///
60/// Maintains the k highest-scoring documents using a min-heap where the
61/// lowest score is at the top for O(1) threshold lookup and O(log k) eviction.
62/// No deduplication — caller must ensure each doc_id is inserted only once.
63///
64/// This is intentionally separate from `TopKCollector` in `collector.rs`:
65/// `ScoreCollector` is used inside `MaxScoreExecutor` where only `(doc_id,
66/// score, ordinal)` tuples exist — no `Scorer` trait, no position tracking,
67/// and the threshold must be inlined for tight block-max loops.
68/// `TopKCollector` wraps a `Scorer` and drives the full `DocSet`/`Scorer`
69/// protocol, collecting positions on demand.
70pub struct ScoreCollector {
71    /// Min-heap of top-k entries (lowest score at top for eviction)
72    heap: BinaryHeap<HeapEntry>,
73    pub k: usize,
74    /// Cached threshold: avoids repeated heap.peek() in hot loops.
75    /// Updated only when the heap changes (insert/pop).
76    cached_threshold: f32,
77    /// Score of the logical sentinel filling every unused top-k slot after
78    /// threshold seeding. Keeping one score here instead of `k - heap.len()`
79    /// entries makes filling those unused slots O(1) time and memory.
80    virtual_threshold: Option<f32>,
81}
82
83impl ScoreCollector {
84    /// Create a new collector for top-k results
85    pub fn new(k: usize) -> Self {
86        Self {
87            heap: BinaryHeap::with_capacity(k.min(MAX_INITIAL_SCORE_COLLECTOR_CAPACITY)),
88            k,
89            cached_threshold: 0.0,
90            virtual_threshold: None,
91        }
92    }
93
94    /// Current score threshold (minimum score to enter top-k)
95    #[inline]
96    pub fn threshold(&self) -> f32 {
97        self.cached_threshold
98    }
99
100    /// Recompute cached threshold from heap state
101    #[inline]
102    fn update_threshold(&mut self) {
103        self.cached_threshold = if let Some(threshold) = self.virtual_threshold {
104            threshold
105        } else if self.heap.len() >= self.k {
106            self.heap.peek().map(|e| e.score).unwrap_or(0.0)
107        } else {
108            0.0
109        };
110    }
111
112    /// Insert a document score. Returns true if inserted in top-k.
113    /// Caller must ensure each doc_id is inserted only once.
114    #[inline]
115    pub fn insert(&mut self, doc_id: DocId, score: f32) -> bool {
116        self.insert_with_ordinal(doc_id, score, 0)
117    }
118
119    /// Insert a document score with ordinal. Returns true if inserted in top-k.
120    /// Caller must ensure each doc_id is inserted only once.
121    #[inline]
122    pub fn insert_with_ordinal(&mut self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
123        if self.k == 0 {
124            return false;
125        }
126        let entry = HeapEntry {
127            doc_id,
128            score,
129            ordinal,
130        };
131        if self.heap.len() < self.k {
132            if let Some(threshold) = self.virtual_threshold {
133                let sentinel = HeapEntry {
134                    doc_id: u32::MAX,
135                    score: threshold,
136                    ordinal: 0,
137                };
138                if entry >= sentinel {
139                    return false;
140                }
141            }
142
143            self.heap.push(entry);
144            // The final real entry displaces the last virtual sentinel.
145            if self.heap.len() == self.k {
146                self.virtual_threshold = None;
147                self.update_threshold();
148            }
149            true
150        } else if self.heap.peek().is_some_and(|worst| entry < *worst) {
151            {
152                let mut worst = self.heap.peek_mut().expect("full heap has a root");
153                *worst = entry;
154            }
155            self.update_threshold();
156            true
157        } else {
158            false
159        }
160    }
161
162    /// Check if a score could potentially enter top-k
163    #[inline]
164    pub fn would_enter(&self, score: f32) -> bool {
165        self.len() < self.k || score > self.cached_threshold
166    }
167
168    /// Check whether this fully identified candidate ranks ahead of the current
169    /// worst retained entry, including deterministic tie breaks.
170    #[inline]
171    pub fn would_enter_candidate(&self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
172        if self.k == 0 {
173            return false;
174        }
175        let entry = HeapEntry {
176            doc_id,
177            score,
178            ordinal,
179        };
180        if let Some(threshold) = self.virtual_threshold {
181            let sentinel = HeapEntry {
182                doc_id: u32::MAX,
183                score: threshold,
184                ordinal: 0,
185            };
186            entry < sentinel
187        } else {
188            self.heap.len() < self.k || self.heap.peek().is_some_and(|worst| entry < *worst)
189        }
190    }
191
192    /// Get the conceptual heap length, including virtual threshold sentinels.
193    #[inline]
194    pub fn len(&self) -> usize {
195        if self.virtual_threshold.is_some() {
196            self.k
197        } else {
198            self.heap.len()
199        }
200    }
201
202    /// Number of real results retained, excluding threshold sentinels.
203    #[inline]
204    pub fn real_len(&self) -> usize {
205        self.heap.len()
206    }
207
208    /// Check if collector is empty
209    #[inline]
210    pub fn is_empty(&self) -> bool {
211        self.len() == 0
212    }
213
214    /// Seed the threshold from a cross-segment shared value.
215    ///
216    /// Logically fills unused slots and replaces retained entries below the new
217    /// floor with virtual dummy entries. This can be called repeatedly while
218    /// another segment raises the shared threshold; equal-scoring real
219    /// candidates win the deterministic doc-id tie break over sentinels.
220    pub fn seed_threshold(&mut self, initial_threshold: f32) {
221        if initial_threshold <= 0.0
222            || self.k == 0
223            || (self.len() >= self.k && initial_threshold <= self.cached_threshold)
224        {
225            return;
226        }
227
228        let sentinel = HeapEntry {
229            doc_id: u32::MAX,
230            score: initial_threshold,
231            ordinal: 0,
232        };
233
234        // When unused slots are already represented by a virtual sentinel, a
235        // new seed only changes the heap if it outranks the old floor. With an
236        // all-real full heap, it must similarly outrank the current root.
237        if let Some(current_threshold) = self.virtual_threshold {
238            let current = HeapEntry {
239                doc_id: u32::MAX,
240                score: current_threshold,
241                ordinal: 0,
242            };
243            if sentinel >= current {
244                return;
245            }
246        } else if self.heap.len() >= self.k
247            && !self.heap.peek().is_some_and(|worst| sentinel < *worst)
248        {
249            return;
250        }
251
252        self.virtual_threshold = Some(initial_threshold);
253        while self.heap.peek().is_some_and(|worst| sentinel < *worst) {
254            self.heap.pop();
255        }
256        self.update_threshold();
257    }
258
259    /// Convert to sorted top-k results (descending by score).
260    /// Filters out sentinel entries (doc_id == u32::MAX) from threshold seeding.
261    pub fn into_sorted_results(self) -> Vec<(DocId, f32, u16)> {
262        let mut results: Vec<(DocId, f32, u16)> = self
263            .heap
264            .into_vec()
265            .into_iter()
266            .filter(|e| e.doc_id != u32::MAX)
267            .map(|e| (e.doc_id, e.score, e.ordinal))
268            .collect();
269
270        // Sort by score descending, then doc_id ascending
271        results.sort_unstable_by(|a, b| {
272            b.1.total_cmp(&a.1)
273                .then_with(|| a.0.cmp(&b.0))
274                .then_with(|| a.2.cmp(&b.2))
275        });
276
277        results
278    }
279}
280
281/// Cross-segment top-k score floor, shared across the parallel/concurrent
282/// per-segment searches of a single query.
283///
284/// Stores an `f32` as raw bits in an atomic so it can be read and monotonically
285/// raised from many threads without a lock. Each segment reads the current
286/// floor as its initial pruning threshold (`ScorerOptions::initial_threshold`)
287/// and, once it has collected a *full* top-k of its own, raises the floor to
288/// its k-th score.
289///
290/// Safety of seeding: a segment only raises the floor after filling its own
291/// heap, so a floor value `v` is always backed by at least `k` real documents
292/// scoring `>= v`. The final merged k-th score is therefore `>= v`, and seeding
293/// any other segment with `v` can never drop a document that belongs in the
294/// final top-k. Completion order is arbitrary, so the floor is best-effort — it
295/// only changes how aggressively later segments prune, never correctness.
296///
297/// The floor carries the query's result-window depth `k` (`for_limit`).
298/// Publishing from a heap shallower than `k` is invalid — a segment with
299/// fewer documents than the window fills its clamped heap early, and its
300/// heap threshold says nothing about the query-global k-th score. Executors
301/// must check `SharedThreshold::covers` before raising the floor with a
302/// full-heap threshold.
303#[derive(Clone, Debug)]
304pub struct SharedThreshold {
305    floor: std::sync::Arc<std::sync::atomic::AtomicU32>,
306    /// Result-window depth the floor is valid for. `usize::MAX` means the
307    /// depth is unknown; reading stays safe, publishing is disabled.
308    k: usize,
309    /// Wall-clock budget of the whole query (anytime mode): executors that
310    /// honour it stop scoring once it passes and flag the result truncated.
311    deadline: Option<std::time::Instant>,
312    /// Set by any executor that stopped early because of `deadline`.
313    truncated: std::sync::Arc<std::sync::atomic::AtomicBool>,
314}
315
316impl Default for SharedThreshold {
317    fn default() -> Self {
318        Self::new()
319    }
320}
321
322impl SharedThreshold {
323    /// A fresh floor of 0.0 (no pruning seed) with an unknown window depth.
324    /// Executors can read and manually raise it, but never publish their own
325    /// full-heap thresholds into it.
326    pub fn new() -> Self {
327        Self::with_depth(usize::MAX)
328    }
329
330    /// A fresh floor valid for a query fetching `limit` results.
331    pub fn for_limit(limit: usize) -> Self {
332        Self::with_depth(limit)
333    }
334
335    fn with_depth(k: usize) -> Self {
336        Self {
337            // 0.0_f32.to_bits() == 0, matching AtomicU32::default().
338            floor: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
339            k,
340            deadline: None,
341            truncated: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
342        }
343    }
344
345    /// Attach a wall-clock budget (`None` = unbounded).
346    pub fn with_deadline(mut self, deadline: Option<std::time::Instant>) -> Self {
347        self.deadline = deadline;
348        self
349    }
350
351    /// The query's deadline, if any.
352    pub fn deadline(&self) -> Option<std::time::Instant> {
353        self.deadline
354    }
355
356    /// Keep cancellation/observability, but isolate a component's score space.
357    pub(crate) fn budget_only(&self) -> Self {
358        Self {
359            floor: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
360            k: usize::MAX,
361            deadline: self.deadline,
362            truncated: self.truncated.clone(),
363        }
364    }
365
366    #[inline]
367    pub(crate) fn stop_if_expired(&self) -> bool {
368        if self.expired() {
369            self.mark_truncated();
370            true
371        } else {
372            false
373        }
374    }
375
376    /// Whether the deadline has passed.
377    #[inline]
378    pub fn expired(&self) -> bool {
379        self.deadline
380            .is_some_and(|deadline| std::time::Instant::now() >= deadline)
381    }
382
383    /// Record that an executor stopped early because the deadline passed.
384    pub fn mark_truncated(&self) {
385        self.truncated
386            .store(true, std::sync::atomic::Ordering::Relaxed);
387    }
388
389    /// Whether any executor of this query stopped early.
390    pub fn truncated(&self) -> bool {
391        self.truncated.load(std::sync::atomic::Ordering::Relaxed)
392    }
393
394    /// True when a full heap of `heap_depth` distinct documents backs a valid
395    /// query-global floor for this threshold's result window.
396    #[inline]
397    pub(crate) fn covers(&self, heap_depth: usize) -> bool {
398        heap_depth >= self.k
399    }
400
401    /// Current floor.
402    #[inline]
403    pub fn get(&self) -> f32 {
404        f32::from_bits(self.floor.load(std::sync::atomic::Ordering::Relaxed))
405    }
406
407    /// Raise the floor to `score` if it is strictly higher. Monotonic; a lower
408    /// or non-positive `score` is ignored. Scores here are BM25/sparse and thus
409    /// non-negative, but the comparison is done on `f32` values (not raw bits)
410    /// so it stays correct regardless.
411    pub fn raise(&self, score: f32) {
412        // Ignore non-positive scores; a NaN falls through harmlessly (the CAS
413        // loop condition below is false for NaN, so nothing is stored).
414        if score <= 0.0 {
415            return;
416        }
417        use std::sync::atomic::Ordering::Relaxed;
418        let bits = score.to_bits();
419        let mut cur = self.floor.load(Relaxed);
420        while f32::from_bits(cur) < score {
421            match self
422                .floor
423                .compare_exchange_weak(cur, bits, Relaxed, Relaxed)
424            {
425                Ok(_) => break,
426                Err(actual) => cur = actual,
427            }
428        }
429    }
430}
431
432/// Search result from MaxScore execution
433#[derive(Debug, Clone, Copy)]
434pub struct ScoredDoc {
435    pub doc_id: DocId,
436    pub score: f32,
437    /// Ordinal for multi-valued fields (which vector in the field matched)
438    pub ordinal: u16,
439}
440
441/// Unified Block-Max MaxScore executor for top-k retrieval
442///
443/// Works with both full-text (BM25) and sparse vector (dot product) queries
444/// through the polymorphic `TermCursor`. Combines three optimizations:
445/// 1. **MaxScore partitioning** (Turtle & Flood 1995): terms split into essential
446///    (must check) and non-essential (only scored if candidate is promising)
447/// 2. **Block-max pruning** (Ding & Suel 2011): skip blocks where per-block
448///    upper bounds can't beat the current threshold
449/// 3. **Conjunction optimization** (Lucene/Grand 2023): progressively intersect
450///    essential terms as threshold rises, skipping docs that lack enough terms
451pub struct MaxScoreExecutor<'a> {
452    /// Metric labels (index, field) — set via `with_metric_labels`; empty
453    /// strings render as "unknown"/"?" is avoided by callers passing real
454    /// names from the schema.
455    metric_index: &'a str,
456    metric_field: &'a str,
457    cursors: Vec<TermCursor<'a>>,
458    prefix_sums: Vec<f32>,
459    collector: ScoreCollector,
460    inv_heap_factor: f32,
461    predicate: Option<super::DocPredicate<'a>>,
462    /// Query-global budget: checked every few thousand loop iterations;
463    /// an expired deadline ends traversal with the results so far.
464    budget: Option<SharedThreshold>,
465}
466
467/// Where a text cursor reads the length of a scoring unit: chunk lengths of a
468/// chunked field, or the persisted per-document field lengths (norms) of a
469/// plain field. Without either, `tf` stands in for the length.
470#[derive(Clone, Copy)]
471pub enum LengthSource<'a> {
472    Chunks(&'a crate::segment::chunk_map::ChunkMap),
473    Docs(&'a crate::segment::chunk_map::DocLengths),
474}
475
476impl LengthSource<'_> {
477    #[inline]
478    pub fn length(&self, id: u32) -> u32 {
479        match self {
480            LengthSource::Chunks(map) => map.bm25_length(id),
481            LengthSource::Docs(lengths) => lengths.length(id),
482        }
483    }
484}
485
486/// Unified term cursor for Block-Max MaxScore execution.
487///
488/// All per-position decode buffers (`doc_ids`, `scores`, `ordinals`) live in
489/// the struct directly and are filled by `ensure_block_loaded`.
490///
491/// Skip-list metadata is **not** materialized — it is read lazily from the
492/// underlying source (`BlockPostingList` for text, `SparseIndex` for sparse),
493/// both backed by zero-copy mmap'd `OwnedBytes`.
494pub(crate) struct TermCursor<'a> {
495    pub max_score: f32,
496    num_blocks: usize,
497    // ── Per-position state (filled by ensure_block_loaded) ──────────
498    block_idx: usize,
499    doc_ids: Vec<u32>,
500    scores: Vec<f32>,
501    ordinals: Vec<u16>,
502    pos: usize,
503    block_loaded: bool,
504    exhausted: bool,
505    // ── Lazy ordinal decode (sparse only) ───────────────────────────
506    /// When true, ordinal decode is deferred until ordinal_mut() is called.
507    /// Set to true for MaxScoreExecutor cursors (most blocks never need ordinals).
508    lazy_ordinals: bool,
509    /// Whether ordinals have been decoded for the current block.
510    ordinals_loaded: bool,
511    /// Stored sparse block for deferred ordinal decode (cheap Arc clone of mmap data).
512    current_sparse_block: Option<crate::structures::SparseBlock>,
513    // ── Block decode + skip access source ───────────────────────────
514    variant: CursorVariant<'a>,
515}
516
517// One cursor per query term; the text variant carries the decoded-block
518// state inline on purpose (no indirection on the scoring path).
519#[allow(clippy::large_enum_variant)]
520enum CursorVariant<'a> {
521    /// Full-text BM25 — in-memory BlockPostingList (skip list + block data)
522    Text {
523        list: crate::structures::BlockPostingList,
524        idf: f32,
525        /// Precomputed: idf * (BM25_K1 + 1.0) — numerator scale factor
526        idf_times_k1_plus_1: f32,
527        /// Precomputed: 1.0 + BM25_K1 * (BM25_B / avg_field_len) — denominator tf coefficient
528        denom_tf_coeff: f32,
529        /// Precomputed: BM25_K1 * (1.0 - BM25_B) — denominator constant
530        denom_const: f32,
531        /// Precomputed: BM25_K1 * BM25_B / avg_len — per-token length
532        /// coefficient, used when `lengths` supplies real chunk lengths.
533        denom_len_coeff: f32,
534        /// Real per-posting lengths (chunk lengths or document norms).
535        /// `None` keeps the historic `tf`-as-length approximation.
536        lengths: Option<LengthSource<'a>>,
537        /// Block bounds may use the block's minimum length: only when the
538        /// list stores one and scoring uses real lengths (a `tf`-as-length
539        /// score is not bounded by a real-length bound).
540        length_bounds: bool,
541        length_floor: u32,
542        block_bound: CachedScoreBound,
543        group_bound: CachedScoreBound,
544        /// Average length used by the bounds (matches the scoring average).
545        avg_len: f32,
546        /// Per-field k1/b, used by the block and group bounds.
547        params: super::Bm25Params,
548        tfs: Vec<u32>,
549        /// Deferred TF decode state: (block_offset, tf_start, count).
550        /// Set when doc_ids are decoded but TFs/scores are not yet computed.
551        deferred_tf: Option<(usize, usize, usize)>,
552    },
553    /// Sparse vector — mmap'd SparseIndex (skip entries + block data)
554    Sparse {
555        si: &'a crate::segment::SparseIndex,
556        query_weight: f32,
557        skip_start: usize,
558        block_data_offset: u64,
559    },
560}
561
562/// One query-local memoized bound, not a corpus-sized table. Atomic packing
563/// keeps shared read access Sync without a lock or separate key/value races.
564struct CachedScoreBound(std::sync::atomic::AtomicU64);
565
566impl CachedScoreBound {
567    fn new() -> Self {
568        Self(std::sync::atomic::AtomicU64::new(u64::MAX))
569    }
570
571    fn get_or_compute(&self, key: usize, compute: impl FnOnce() -> f32) -> f32 {
572        use std::sync::atomic::Ordering::Relaxed;
573        // Posting block counts are bounded by the u32 posting-id space.
574        let key = key as u64;
575        let entry = self.0.load(Relaxed);
576        if entry >> 32 == key {
577            return f32::from_bits(entry as u32);
578        }
579        let score = compute();
580        self.0
581            .store((key << 32) | u64::from(score.to_bits()), Relaxed);
582        score
583    }
584}
585
586// ── TermCursor async/sync macros ──────────────────────────────────────────
587//
588// Parameterised on:
589//   $load_block_fn – load_block_direct | load_block_direct_sync  (sparse I/O)
590//   $ensure_fn     – ensure_block_loaded | ensure_block_loaded_sync
591//   $($aw)*        – .await  (present for async, absent for sync)
592
593macro_rules! cursor_ensure_block {
594    ($self:ident, $load_block_fn:ident, $($aw:tt)*) => {{
595        if $self.exhausted || $self.block_loaded {
596            return Ok(!$self.exhausted);
597        }
598        match &mut $self.variant {
599            CursorVariant::Text {
600                list,
601                deferred_tf,
602                ..
603            } => {
604                if let Some(state) = list.decode_block_doc_ids_only($self.block_idx, &mut $self.doc_ids) {
605                    *deferred_tf = Some(state);
606                    $self.scores.clear();
607                    $self.pos = 0;
608                    $self.block_loaded = true;
609                    Ok(true)
610                } else {
611                    $self.exhausted = true;
612                    Ok(false)
613                }
614            }
615            CursorVariant::Sparse {
616                si,
617                query_weight,
618                skip_start,
619                block_data_offset,
620                ..
621            } => {
622                let block = si
623                    .$load_block_fn(*skip_start, *block_data_offset, $self.block_idx)
624                    $($aw)* ?;
625                match block {
626                    Some(b) => {
627                        b.decode_doc_ids_into(&mut $self.doc_ids);
628                        b.decode_scored_weights_into(*query_weight, &mut $self.scores);
629                        if $self.lazy_ordinals {
630                            // Defer ordinal decode until ordinal_mut() is called.
631                            // Stores cheap Arc-backed mmap slice, no copy.
632                            $self.current_sparse_block = Some(b);
633                            $self.ordinals_loaded = false;
634                        } else {
635                            b.decode_ordinals_into(&mut $self.ordinals);
636                            $self.ordinals_loaded = true;
637                            $self.current_sparse_block = None;
638                        }
639                        $self.pos = 0;
640                        $self.block_loaded = true;
641                        Ok(true)
642                    }
643                    None => {
644                        $self.exhausted = true;
645                        Ok(false)
646                    }
647                }
648            }
649        }
650    }};
651}
652
653macro_rules! cursor_advance {
654    ($self:ident, $ensure_fn:ident, $($aw:tt)*) => {{
655        if $self.exhausted {
656            return Ok(u32::MAX);
657        }
658        $self.$ensure_fn() $($aw)* ?;
659        if $self.exhausted {
660            return Ok(u32::MAX);
661        }
662        Ok($self.advance_pos())
663    }};
664}
665
666macro_rules! cursor_seek {
667    ($self:ident, $ensure_fn:ident, $target:expr, $($aw:tt)*) => {{
668        if let Some(doc) = $self.seek_prepare($target) {
669            return Ok(doc);
670        }
671        $self.$ensure_fn() $($aw)* ?;
672        if $self.seek_finish($target) {
673            $self.$ensure_fn() $($aw)* ?;
674        }
675        Ok($self.doc())
676    }};
677}
678
679impl<'a> TermCursor<'a> {
680    /// Full-text BM25 cursor with explicit per-field parameters.
681    pub fn text_with_params(
682        posting_list: crate::structures::BlockPostingList,
683        idf: f32,
684        avg_field_len: f32,
685        lengths: Option<LengthSource<'a>>,
686        params: super::Bm25Params,
687    ) -> Self {
688        Self::text_with_lengths(posting_list, idf, avg_field_len, lengths, params)
689    }
690
691    fn text_with_lengths(
692        posting_list: crate::structures::BlockPostingList,
693        idf: f32,
694        avg_field_len: f32,
695        lengths: Option<LengthSource<'a>>,
696        params: super::Bm25Params,
697    ) -> Self {
698        let max_tf = posting_list.max_tf() as f32;
699        let safe_avg = avg_field_len.max(1.0);
700        let length_bounds = lengths.is_some() && posting_list.min_len().is_some();
701        let length_floor = match lengths {
702            Some(LengthSource::Chunks(map)) => map.length_floor(),
703            _ => 0,
704        };
705        let max_score = match posting_list.min_len() {
706            Some(min_len) if length_bounds => params.upper_bound_with_len(
707                max_tf.max(1.0),
708                idf,
709                min_len.max(length_floor) as f32,
710                safe_avg,
711            ),
712            _ => params.upper_bound(max_tf.max(1.0), idf),
713        };
714        let num_blocks = posting_list.num_blocks();
715        Self {
716            max_score,
717            num_blocks,
718            block_idx: 0,
719            doc_ids: Vec::with_capacity(128),
720            scores: Vec::with_capacity(128),
721            ordinals: Vec::new(),
722            pos: 0,
723            block_loaded: false,
724            exhausted: num_blocks == 0,
725            lazy_ordinals: false,
726            ordinals_loaded: true, // text cursors never have ordinals
727            current_sparse_block: None,
728            variant: CursorVariant::Text {
729                list: posting_list,
730                idf,
731                idf_times_k1_plus_1: idf * (params.k1 + 1.0),
732                denom_tf_coeff: 1.0 + params.k1 * (params.b / safe_avg),
733                denom_const: params.k1 * (1.0 - params.b),
734                denom_len_coeff: params.k1 * params.b / safe_avg,
735                lengths,
736                length_bounds,
737                length_floor,
738                block_bound: CachedScoreBound::new(),
739                group_bound: CachedScoreBound::new(),
740                avg_len: safe_avg,
741                params,
742                tfs: Vec::with_capacity(128),
743                deferred_tf: None,
744            },
745        }
746    }
747
748    /// Create a sparse vector cursor with lazy block loading.
749    /// Skip entries are **not** copied — they are read from `SparseIndex` mmap on demand.
750    pub fn sparse(
751        si: &'a crate::segment::SparseIndex,
752        query_weight: f32,
753        skip_start: usize,
754        skip_count: usize,
755        global_max_weight: f32,
756        block_data_offset: u64,
757    ) -> Self {
758        Self {
759            max_score: query_weight.abs() * global_max_weight,
760            num_blocks: skip_count,
761            block_idx: 0,
762            doc_ids: Vec::with_capacity(256),
763            scores: Vec::with_capacity(256),
764            ordinals: Vec::with_capacity(256),
765            pos: 0,
766            block_loaded: false,
767            exhausted: skip_count == 0,
768            lazy_ordinals: false,
769            ordinals_loaded: true,
770            current_sparse_block: None,
771            variant: CursorVariant::Sparse {
772                si,
773                query_weight,
774                skip_start,
775                block_data_offset,
776            },
777        }
778    }
779
780    // ── Skip-entry access (lazy, zero-copy for sparse) ──────────────────
781
782    #[inline]
783    fn block_first_doc(&self, idx: usize) -> DocId {
784        match &self.variant {
785            CursorVariant::Text { list, .. } => list.block_first_doc(idx).unwrap_or(u32::MAX),
786            CursorVariant::Sparse { si, skip_start, .. } => {
787                si.read_skip_entry(*skip_start + idx).first_doc
788            }
789        }
790    }
791
792    #[inline]
793    fn block_last_doc(&self, idx: usize) -> DocId {
794        match &self.variant {
795            CursorVariant::Text { list, .. } => list.block_last_doc(idx).unwrap_or(0),
796            CursorVariant::Sparse { si, skip_start, .. } => {
797                si.read_skip_entry(*skip_start + idx).last_doc
798            }
799        }
800    }
801
802    // ── Read-only accessors ─────────────────────────────────────────────
803
804    #[inline]
805    pub fn doc(&self) -> DocId {
806        if self.exhausted {
807            return u32::MAX;
808        }
809        if self.block_loaded {
810            debug_assert!(self.pos < self.doc_ids.len());
811            // SAFETY: pos < doc_ids.len() is maintained by advance_pos/ensure_block_loaded.
812            unsafe { *self.doc_ids.get_unchecked(self.pos) }
813        } else {
814            self.block_first_doc(self.block_idx)
815        }
816    }
817
818    #[inline]
819    pub fn ordinal(&self) -> u16 {
820        if !self.block_loaded || self.ordinals.is_empty() {
821            return 0;
822        }
823        debug_assert!(self.pos < self.ordinals.len());
824        // SAFETY: pos < ordinals.len() is maintained by advance_pos/ensure_block_loaded.
825        unsafe { *self.ordinals.get_unchecked(self.pos) }
826    }
827
828    /// Lazily-decoded ordinal accessor for MaxScore executor.
829    ///
830    /// When `lazy_ordinals=true`, ordinals are not decoded during block loading.
831    /// This method triggers the deferred decode on first access, amortized over
832    /// the block. Subsequent calls within the same block are free.
833    #[inline]
834    pub fn ordinal_mut(&mut self) -> u16 {
835        if !self.block_loaded {
836            return 0;
837        }
838        if !self.ordinals_loaded {
839            if let Some(ref block) = self.current_sparse_block {
840                block.decode_ordinals_into(&mut self.ordinals);
841            }
842            self.ordinals_loaded = true;
843        }
844        if self.ordinals.is_empty() {
845            return 0;
846        }
847        debug_assert!(self.pos < self.ordinals.len());
848        unsafe { *self.ordinals.get_unchecked(self.pos) }
849    }
850
851    #[inline]
852    pub fn score(&self) -> f32 {
853        if !self.block_loaded {
854            return 0.0;
855        }
856        debug_assert!(self.pos < self.scores.len());
857        // SAFETY: pos < scores.len() is maintained by advance_pos/ensure_block_loaded.
858        unsafe { *self.scores.get_unchecked(self.pos) }
859    }
860
861    /// Ensure BM25 scores are computed for the current block (lazy TF decode).
862    ///
863    /// For text cursors, TF unpacking and BM25 scoring are deferred from block
864    /// loading until this method is called, saving work for blocks skipped by
865    /// block-max or conjunction pruning. No-op for sparse cursors.
866    #[inline]
867    pub fn ensure_scores(&mut self) {
868        if self.block_loaded && self.scores.is_empty() {
869            self.compute_deferred_scores();
870        }
871    }
872
873    #[inline]
874    pub fn current_block_max_score(&self) -> f32 {
875        if self.exhausted {
876            return 0.0;
877        }
878        match &self.variant {
879            CursorVariant::Text { .. } => self.text_block_bound(self.block_idx),
880            CursorVariant::Sparse {
881                si,
882                query_weight,
883                skip_start,
884                ..
885            } => query_weight.abs() * si.read_skip_entry(*skip_start + self.block_idx).max_weight,
886        }
887    }
888
889    /// Upper bound over the L1 group (eight blocks) containing the current
890    /// block, for text lists that store superblock bounds; `None` when the
891    /// cursor cannot bound a whole group (sparse, legacy lists).
892    #[inline]
893    pub fn current_group_max_score(&self) -> Option<f32> {
894        if self.exhausted {
895            return Some(0.0);
896        }
897        match &self.variant {
898            CursorVariant::Text { .. } => self.text_group_bound(self.block_idx),
899            CursorVariant::Sparse { .. } => None,
900        }
901    }
902
903    /// Whether this cursor reads an in-memory text posting list (all of its
904    /// I/O is synchronous, so the windowed executor can drive it).
905    #[inline]
906    pub(crate) fn is_text(&self) -> bool {
907        matches!(self.variant, CursorVariant::Text { .. })
908    }
909
910    /// Upper bound of text block `idx` from its `(max_tf, min_len)` word.
911    fn text_block_bound(&self, idx: usize) -> f32 {
912        match &self.variant {
913            CursorVariant::Text {
914                list,
915                idf,
916                length_bounds,
917                length_floor,
918                block_bound,
919                avg_len,
920                params,
921                ..
922            } => block_bound.get_or_compute(idx, || {
923                let (max_tf, min_len) = list.block_bounds(idx).unwrap_or((0, None));
924                match min_len {
925                    Some(min_len) if *length_bounds => params.upper_bound_with_len(
926                        (max_tf as f32).max(1.0),
927                        *idf,
928                        min_len.max(*length_floor) as f32,
929                        *avg_len,
930                    ),
931                    _ => params.upper_bound((max_tf as f32).max(1.0), *idf),
932                }
933            }),
934            CursorVariant::Sparse { .. } => self.max_score,
935        }
936    }
937
938    /// Upper bound of the L1 group containing text block `idx`.
939    fn text_group_bound(&self, idx: usize) -> Option<f32> {
940        match &self.variant {
941            CursorVariant::Text {
942                list,
943                idf,
944                length_bounds,
945                length_floor,
946                group_bound,
947                avg_len,
948                params,
949                ..
950            } => {
951                let (max_tf, min_len) = list.group_bounds(idx)?;
952                Some(group_bound.get_or_compute(list.next_group_block(idx), || {
953                    if *length_bounds {
954                        params.upper_bound_with_len(
955                            (max_tf as f32).max(1.0),
956                            *idf,
957                            min_len.max(*length_floor) as f32,
958                            *avg_len,
959                        )
960                    } else {
961                        params.upper_bound((max_tf as f32).max(1.0), *idf)
962                    }
963                }))
964            }
965            CursorVariant::Sparse { .. } => None,
966        }
967    }
968
969    /// Upper bound of this cursor's contribution to any id in `[from, to]`:
970    /// the largest block bound over the blocks intersecting the range, with
971    /// one L1 word standing in for a group that lies inside it. Reads skip
972    /// entries only; no block is decoded (Lucene `advanceShallow` +
973    /// `getMaxScore(upTo)`).
974    pub(crate) fn window_upper_bound(&self, from: DocId, to: DocId) -> f32 {
975        if self.exhausted {
976            return 0.0;
977        }
978        let CursorVariant::Text { list, .. } = &self.variant else {
979            return self.max_score;
980        };
981        // Postings the cursor has already passed cannot score again: the
982        // bound starts at its current id, not at the window start.
983        let start = from.max(self.doc());
984        if start > to {
985            return 0.0;
986        }
987        let Some(mut idx) = list.seek_block(start, self.block_idx) else {
988            return 0.0;
989        };
990        let mut bound = 0.0f32;
991        while idx < self.num_blocks {
992            if list.block_first_doc(idx).unwrap_or(u32::MAX) > to {
993                break;
994            }
995            if list.is_group_start(idx)
996                && list.group_last_doc(idx).is_some_and(|last| last <= to)
997                && let Some(group_bound) = self.text_group_bound(idx)
998            {
999                bound = bound.max(group_bound);
1000                idx = list.next_group_block(idx);
1001                continue;
1002            }
1003            bound = bound.max(self.text_block_bound(idx));
1004            idx += 1;
1005        }
1006        bound
1007    }
1008
1009    /// Add this cursor's scores for every id in `[from, to]` to the window
1010    /// buffers (`scores[id - from]`, bit `id - from` of `mask`) and leave the
1011    /// cursor on its first id after `to`. Whole runs of a block are
1012    /// processed in one pass over its decoded arrays. Text cursors only.
1013    pub(crate) fn score_window_sync(
1014        &mut self,
1015        from: DocId,
1016        to: DocId,
1017        scores: &mut [f32],
1018        mask: &mut [u64],
1019    ) -> crate::Result<u32> {
1020        let mut matched = 0u32;
1021        loop {
1022            if self.exhausted {
1023                return Ok(matched);
1024            }
1025            if !self.block_loaded {
1026                if self.block_first_doc(self.block_idx) > to {
1027                    return Ok(matched);
1028                }
1029                self.ensure_block_loaded_sync()?;
1030                if self.exhausted {
1031                    return Ok(matched);
1032                }
1033            }
1034            if self.doc_ids[self.pos] > to {
1035                return Ok(matched);
1036            }
1037            self.ensure_scores();
1038            let remaining = &self.doc_ids[self.pos..];
1039            let end = if to == u32::MAX {
1040                remaining.len()
1041            } else {
1042                crate::structures::simd::find_first_ge_u32(remaining, to + 1)
1043            };
1044            let block_scores = &self.scores[self.pos..self.pos + end];
1045            for (doc, score) in remaining[..end].iter().zip(block_scores) {
1046                let slot = (doc - from) as usize;
1047                scores[slot] += score;
1048                mask[slot >> 6] |= 1u64 << (slot & 63);
1049            }
1050            matched += end as u32;
1051            self.pos += end;
1052            if self.pos >= self.doc_ids.len() {
1053                self.block_idx += 1;
1054                self.block_loaded = false;
1055                if self.block_idx >= self.num_blocks {
1056                    self.exhausted = true;
1057                    return Ok(matched);
1058                }
1059            } else {
1060                return Ok(matched);
1061            }
1062        }
1063    }
1064
1065    /// Move past every id `<= to`, skipping whole blocks that end before it
1066    /// without decoding them.
1067    pub(crate) fn skip_past_sync(&mut self, to: DocId) -> crate::Result<()> {
1068        if to == u32::MAX {
1069            self.exhausted = true;
1070            return Ok(());
1071        }
1072        while !self.exhausted && self.block_last_doc(self.block_idx) <= to {
1073            self.skip_to_next_block();
1074        }
1075        if !self.exhausted && self.doc() <= to {
1076            self.seek_sync(to + 1)?;
1077        }
1078        Ok(())
1079    }
1080
1081    /// Last doc of the L1 group containing the current block (text only).
1082    #[inline]
1083    pub fn current_group_last_doc(&self) -> DocId {
1084        match &self.variant {
1085            CursorVariant::Text { list, .. } => list.group_last_doc(self.block_idx).unwrap_or(0),
1086            CursorVariant::Sparse { .. } => self.block_last_doc(self.block_idx),
1087        }
1088    }
1089
1090    /// Jump past the current L1 group (text) or block (sparse).
1091    pub fn skip_to_next_group(&mut self) -> DocId {
1092        if self.exhausted {
1093            return u32::MAX;
1094        }
1095        let next = match &self.variant {
1096            CursorVariant::Text { list, .. } => list.next_group_block(self.block_idx),
1097            CursorVariant::Sparse { .. } => self.block_idx + 1,
1098        };
1099        self.block_idx = next;
1100        self.block_loaded = false;
1101        if self.block_idx >= self.num_blocks {
1102            self.exhausted = true;
1103            return u32::MAX;
1104        }
1105        self.block_first_doc(self.block_idx)
1106    }
1107
1108    // ── Block navigation ────────────────────────────────────────────────
1109
1110    pub fn skip_to_next_block(&mut self) -> DocId {
1111        if self.exhausted {
1112            return u32::MAX;
1113        }
1114        self.block_idx += 1;
1115        self.block_loaded = false;
1116        if self.block_idx >= self.num_blocks {
1117            self.exhausted = true;
1118            return u32::MAX;
1119        }
1120        self.block_first_doc(self.block_idx)
1121    }
1122
1123    #[inline]
1124    fn advance_pos(&mut self) -> DocId {
1125        self.pos += 1;
1126        if self.pos >= self.doc_ids.len() {
1127            self.block_idx += 1;
1128            self.block_loaded = false;
1129            if self.block_idx >= self.num_blocks {
1130                self.exhausted = true;
1131                return u32::MAX;
1132            }
1133        }
1134        self.doc()
1135    }
1136
1137    /// Compute BM25 scores from deferred TF data (lazy decode for text cursors).
1138    #[inline(never)]
1139    fn compute_deferred_scores(&mut self) {
1140        if let CursorVariant::Text {
1141            list,
1142            idf_times_k1_plus_1,
1143            denom_tf_coeff,
1144            denom_const,
1145            denom_len_coeff,
1146            lengths,
1147            tfs,
1148            deferred_tf,
1149            ..
1150        } = &mut self.variant
1151            && let Some((block_offset, tf_start, count)) = deferred_tf.take()
1152        {
1153            list.decode_block_tfs_deferred(block_offset, tf_start, count, tfs);
1154            let num_scale = *idf_times_k1_plus_1;
1155            let d_tf = *denom_tf_coeff;
1156            let d_const = *denom_const;
1157            let d_len = *denom_len_coeff;
1158            self.scores.clear();
1159            self.scores.resize(count, 0.0);
1160            match lengths {
1161                // Real BM25 length normalisation per chunk or document.
1162                Some(source) => {
1163                    for i in 0..count {
1164                        let tf = unsafe { *tfs.get_unchecked(i) } as f32;
1165                        let vid = unsafe { *self.doc_ids.get_unchecked(i) };
1166                        let len = source.length(vid) as f32;
1167                        let score = (num_scale * tf) / (tf + d_const + d_len * len);
1168                        unsafe {
1169                            *self.scores.get_unchecked_mut(i) = score;
1170                        }
1171                    }
1172                }
1173                None => {
1174                    for i in 0..count {
1175                        let tf = unsafe { *tfs.get_unchecked(i) } as f32;
1176                        let score = (num_scale * tf) / (d_tf * tf + d_const);
1177                        unsafe {
1178                            *self.scores.get_unchecked_mut(i) = score;
1179                        }
1180                    }
1181                }
1182            }
1183        }
1184    }
1185
1186    // ── Block loading / advance / seek ─────────────────────────────────
1187    //
1188    // Macros parameterised on sparse I/O method + optional .await to
1189    // stamp out both async and sync variants without duplication.
1190
1191    pub async fn ensure_block_loaded(&mut self) -> crate::Result<bool> {
1192        cursor_ensure_block!(self, load_block_direct, .await)
1193    }
1194
1195    pub fn ensure_block_loaded_sync(&mut self) -> crate::Result<bool> {
1196        cursor_ensure_block!(self, load_block_direct_sync,)
1197    }
1198
1199    pub async fn advance(&mut self) -> crate::Result<DocId> {
1200        cursor_advance!(self, ensure_block_loaded, .await)
1201    }
1202
1203    pub fn advance_sync(&mut self) -> crate::Result<DocId> {
1204        cursor_advance!(self, ensure_block_loaded_sync,)
1205    }
1206
1207    pub async fn seek(&mut self, target: DocId) -> crate::Result<DocId> {
1208        cursor_seek!(self, ensure_block_loaded, target, .await)
1209    }
1210
1211    pub fn seek_sync(&mut self, target: DocId) -> crate::Result<DocId> {
1212        cursor_seek!(self, ensure_block_loaded_sync, target,)
1213    }
1214
1215    fn seek_prepare(&mut self, target: DocId) -> Option<DocId> {
1216        if self.exhausted {
1217            return Some(u32::MAX);
1218        }
1219
1220        // Fast path: target is within the currently loaded block
1221        if self.block_loaded
1222            && let Some(&last) = self.doc_ids.last()
1223        {
1224            if last >= target && self.doc_ids[self.pos] < target {
1225                let remaining = &self.doc_ids[self.pos..];
1226                self.pos += crate::structures::simd::find_first_ge_u32(remaining, target);
1227                if self.pos >= self.doc_ids.len() {
1228                    self.block_idx += 1;
1229                    self.block_loaded = false;
1230                    if self.block_idx >= self.num_blocks {
1231                        self.exhausted = true;
1232                        return Some(u32::MAX);
1233                    }
1234                }
1235                return Some(self.doc());
1236            }
1237            if self.doc_ids[self.pos] >= target {
1238                return Some(self.doc());
1239            }
1240        }
1241
1242        // Seek to the block containing target
1243        let lo = match &self.variant {
1244            // Text: SIMD-accelerated 2-level seek (L1 + L0)
1245            CursorVariant::Text { list, .. } => match list.seek_block(target, self.block_idx) {
1246                Some(idx) => idx,
1247                None => {
1248                    self.exhausted = true;
1249                    return Some(u32::MAX);
1250                }
1251            },
1252            // Sparse: binary search on skip entries (lazy mmap reads)
1253            CursorVariant::Sparse { .. } => {
1254                let mut lo = self.block_idx;
1255                let mut hi = self.num_blocks;
1256                while lo < hi {
1257                    let mid = lo + (hi - lo) / 2;
1258                    if self.block_last_doc(mid) < target {
1259                        lo = mid + 1;
1260                    } else {
1261                        hi = mid;
1262                    }
1263                }
1264                lo
1265            }
1266        };
1267        if lo >= self.num_blocks {
1268            self.exhausted = true;
1269            return Some(u32::MAX);
1270        }
1271        if lo != self.block_idx || !self.block_loaded {
1272            self.block_idx = lo;
1273            self.block_loaded = false;
1274        }
1275        None
1276    }
1277
1278    #[inline]
1279    fn seek_finish(&mut self, target: DocId) -> bool {
1280        if self.exhausted {
1281            return false;
1282        }
1283        self.pos = crate::structures::simd::find_first_ge_u32(&self.doc_ids, target);
1284        if self.pos >= self.doc_ids.len() {
1285            self.block_idx += 1;
1286            self.block_loaded = false;
1287            if self.block_idx >= self.num_blocks {
1288                self.exhausted = true;
1289                return false;
1290            }
1291            return true;
1292        }
1293        false
1294    }
1295}
1296
1297/// Macro to stamp out the Block-Max MaxScore loop for both async and sync paths.
1298///
1299/// `$ensure`, `$advance`, `$seek` are cursor method idents (async or _sync variants).
1300/// `$($aw:tt)*` captures `.await` for async or nothing for sync.
1301macro_rules! bms_execute_loop {
1302    ($self:ident, $ensure:ident, $advance:ident, $seek:ident, $($aw:tt)*) => {{
1303        let n = $self.cursors.len();
1304
1305        // Load first block for each cursor (ensures doc() returns real values)
1306        for cursor in &mut $self.cursors {
1307            cursor.$ensure() $($aw)* ?;
1308        }
1309
1310        let mut docs_scored = 0u64;
1311        let mut docs_skipped = 0u64;
1312        let mut blocks_skipped = 0u64;
1313        let mut groups_skipped = 0u64;
1314        let mut conjunction_skipped = 0u64;
1315        let mut ordinal_scores: Vec<(u16, f32)> = Vec::with_capacity(n * 2);
1316        let _bms_start = std::time::Instant::now();
1317
1318        let inv_heap_factor = $self.inv_heap_factor;
1319        let mut adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1320        let mut iterations: u64 = 0;
1321
1322        loop {
1323            // Anytime budget: a coarse deadline check (one clock read per
1324            // 4096 iterations); the results collected so far are returned
1325            // and the query is flagged truncated.
1326            iterations += 1;
1327            if iterations & 0xFFF == 0
1328                && let Some(budget) = &$self.budget
1329                && budget.expired()
1330            {
1331                budget.mark_truncated();
1332                log::debug!(
1333                    "MaxScoreExecutor: deadline reached after {} iterations, {} scored",
1334                    iterations,
1335                    docs_scored
1336                );
1337                break;
1338            }
1339            let partition = $self.find_partition();
1340            if partition >= n {
1341                break;
1342            }
1343
1344            // Find minimum doc_id across essential cursors and collect
1345            // which cursors are at min_doc (avoids redundant re-checks in
1346            // conjunction, block-max, predicate, and scoring passes).
1347            let mut min_doc = u32::MAX;
1348            // Smallest essential doc after min_doc: the first doc where a
1349            // cursor not at min_doc can contribute, hence the farthest a
1350            // block skip may safely go.
1351            let mut next_other = u32::MAX;
1352            let mut at_min_mask = 0u64; // bitset of cursor indices at min_doc
1353            for i in partition..n {
1354                let doc = $self.cursors[i].doc();
1355                match doc.cmp(&min_doc) {
1356                    std::cmp::Ordering::Less => {
1357                        next_other = min_doc;
1358                        min_doc = doc;
1359                        at_min_mask = 1u64 << (i as u32);
1360                    }
1361                    std::cmp::Ordering::Equal => {
1362                        at_min_mask |= 1u64 << (i as u32);
1363                    }
1364                    std::cmp::Ordering::Greater => {
1365                        if doc < next_other {
1366                            next_other = doc;
1367                        }
1368                    }
1369                }
1370            }
1371            if min_doc == u32::MAX {
1372                break;
1373            }
1374
1375            let non_essential_upper = if partition > 0 {
1376                $self.prefix_sums[partition - 1]
1377            } else {
1378                0.0
1379            };
1380
1381            // --- Conjunction optimization ---
1382            if $self.collector.len() >= $self.collector.k {
1383                let mut present_upper: f32 = 0.0;
1384                let mut mask = at_min_mask;
1385                while mask != 0 {
1386                    let i = mask.trailing_zeros() as usize;
1387                    present_upper += $self.cursors[i].max_score;
1388                    mask &= mask - 1;
1389                }
1390
1391                if present_upper + non_essential_upper < adjusted_threshold {
1392                    let mut mask = at_min_mask;
1393                    while mask != 0 {
1394                        let i = mask.trailing_zeros() as usize;
1395                        $self.cursors[i].$ensure() $($aw)* ?;
1396                        $self.cursors[i].$advance() $($aw)* ?;
1397                        mask &= mask - 1;
1398                    }
1399                    conjunction_skipped += 1;
1400                    continue;
1401                }
1402            }
1403
1404            // --- Block-max pruning ---
1405            if $self.collector.len() >= $self.collector.k {
1406                let mut block_max_sum: f32 = 0.0;
1407                let mut mask = at_min_mask;
1408                while mask != 0 {
1409                    let i = mask.trailing_zeros() as usize;
1410                    block_max_sum += $self.cursors[i].current_block_max_score();
1411                    mask &= mask - 1;
1412                }
1413
1414                if block_max_sum + non_essential_upper < adjusted_threshold {
1415                    // Block-Max MaxScore skip: every document before
1416                    // `next_other` is covered only by the cursors at min_doc
1417                    // (plus non-essential ones), whose block bounds cannot
1418                    // reach the threshold. A document at or after
1419                    // `next_other` may also receive another essential
1420                    // cursor's score, so no cursor jumps past it: skip the
1421                    // block when it ends before `next_other`, otherwise seek
1422                    // to `next_other` inside the block.
1423                    //
1424                    // Superblocks: when the cursors' L1 group bounds cannot
1425                    // reach the threshold either, the same argument covers
1426                    // the whole group of eight blocks, so a cursor may jump
1427                    // to its next group instead (bounded by `next_other` in
1428                    // the same way). A cursor without group bounds counts
1429                    // with its block bound and still skips one block.
1430                    let mut group_sum: f32 = 0.0;
1431                    let mut mask = at_min_mask;
1432                    while mask != 0 {
1433                        let i = mask.trailing_zeros() as usize;
1434                        group_sum += $self.cursors[i]
1435                            .current_group_max_score()
1436                            .unwrap_or_else(|| $self.cursors[i].current_block_max_score());
1437                        mask &= mask - 1;
1438                    }
1439                    let group_prunable = group_sum + non_essential_upper < adjusted_threshold;
1440                    let mut mask = at_min_mask;
1441                    while mask != 0 {
1442                        let i = mask.trailing_zeros() as usize;
1443                        let by_group =
1444                            group_prunable && $self.cursors[i].current_group_max_score().is_some();
1445                        let boundary = if by_group {
1446                            $self.cursors[i].current_group_last_doc()
1447                        } else {
1448                            $self.cursors[i].block_last_doc($self.cursors[i].block_idx)
1449                        };
1450                        if next_other > boundary {
1451                            if by_group {
1452                                $self.cursors[i].skip_to_next_group();
1453                                groups_skipped += 1;
1454                            } else {
1455                                $self.cursors[i].skip_to_next_block();
1456                            }
1457                            $self.cursors[i].$ensure() $($aw)* ?;
1458                        } else {
1459                            $self.cursors[i].$seek(next_other) $($aw)* ?;
1460                        }
1461                        mask &= mask - 1;
1462                    }
1463                    blocks_skipped += 1;
1464                    continue;
1465                }
1466            }
1467
1468            // --- Predicate filter (after block-max, before scoring) ---
1469            if let Some(ref pred) = $self.predicate {
1470                if !pred(min_doc) {
1471                    let mut mask = at_min_mask;
1472                    while mask != 0 {
1473                        let i = mask.trailing_zeros() as usize;
1474                        $self.cursors[i].$ensure() $($aw)* ?;
1475                        $self.cursors[i].$advance() $($aw)* ?;
1476                        mask &= mask - 1;
1477                    }
1478                    continue;
1479                }
1480            }
1481
1482            // --- Score essential cursors ---
1483            ordinal_scores.clear();
1484            {
1485                let mut mask = at_min_mask;
1486                while mask != 0 {
1487                    let i = mask.trailing_zeros() as usize;
1488                    $self.cursors[i].$ensure() $($aw)* ?;
1489                    $self.cursors[i].ensure_scores();
1490                    while $self.cursors[i].doc() == min_doc {
1491                        let ord = $self.cursors[i].ordinal_mut();
1492                        let sc = $self.cursors[i].score();
1493                        ordinal_scores.push((ord, sc));
1494                        $self.cursors[i].$advance() $($aw)* ?;
1495                    }
1496                    mask &= mask - 1;
1497                }
1498            }
1499
1500            let essential_total: f32 = ordinal_scores.iter().map(|(_, s)| *s).sum();
1501            if $self.collector.len() >= $self.collector.k
1502                && essential_total + non_essential_upper < adjusted_threshold
1503            {
1504                docs_skipped += 1;
1505                continue;
1506            }
1507
1508            // --- Score non-essential cursors (highest max_score first for early exit) ---
1509            let mut running_total = essential_total;
1510            for i in (0..partition).rev() {
1511                if $self.collector.len() >= $self.collector.k
1512                    && running_total + $self.prefix_sums[i] < adjusted_threshold
1513                {
1514                    break;
1515                }
1516
1517                let doc = $self.cursors[i].$seek(min_doc) $($aw)* ?;
1518                if doc == min_doc {
1519                    $self.cursors[i].ensure_scores();
1520                    while $self.cursors[i].doc() == min_doc {
1521                        let s = $self.cursors[i].score();
1522                        running_total += s;
1523                        let ord = $self.cursors[i].ordinal_mut();
1524                        ordinal_scores.push((ord, s));
1525                        $self.cursors[i].$advance() $($aw)* ?;
1526                    }
1527                }
1528            }
1529
1530            // --- Group by ordinal and insert ---
1531            // Fast path: single entry (common for single-valued fields) — skip sort + grouping
1532            if ordinal_scores.len() == 1 {
1533                let (ord, score) = ordinal_scores[0];
1534                if $self.collector.insert_with_ordinal(min_doc, score, ord) {
1535                    docs_scored += 1;
1536                    adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1537                } else {
1538                    docs_skipped += 1;
1539                }
1540            } else if !ordinal_scores.is_empty() {
1541                if ordinal_scores.len() > 2 {
1542                    ordinal_scores.sort_unstable_by_key(|(ord, _)| *ord);
1543                } else if ordinal_scores.len() == 2 && ordinal_scores[0].0 > ordinal_scores[1].0 {
1544                    ordinal_scores.swap(0, 1);
1545                }
1546                let mut j = 0;
1547                while j < ordinal_scores.len() {
1548                    let current_ord = ordinal_scores[j].0;
1549                    let mut score = 0.0f32;
1550                    while j < ordinal_scores.len() && ordinal_scores[j].0 == current_ord {
1551                        score += ordinal_scores[j].1;
1552                        j += 1;
1553                    }
1554                    if $self
1555                        .collector
1556                        .insert_with_ordinal(min_doc, score, current_ord)
1557                    {
1558                        docs_scored += 1;
1559                        adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1560                    } else {
1561                        docs_skipped += 1;
1562                    }
1563                }
1564            }
1565        }
1566
1567        let results: Vec<ScoredDoc> = $self
1568            .collector
1569            .into_sorted_results()
1570            .into_iter()
1571            .map(|(doc_id, score, ordinal)| ScoredDoc {
1572                doc_id,
1573                score,
1574                ordinal,
1575            })
1576            .collect();
1577
1578        let _bms_elapsed_ms = _bms_start.elapsed().as_millis() as u64;
1579        if _bms_elapsed_ms > 500 {
1580            warn!(
1581                "slow MaxScore: {}ms, cursors={}, scored={}, skipped={}, blocks_skipped={}, groups_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1582                _bms_elapsed_ms,
1583                n,
1584                docs_scored,
1585                docs_skipped,
1586                blocks_skipped,
1587                groups_skipped,
1588                conjunction_skipped,
1589                results.len(),
1590                results.first().map(|r| r.score).unwrap_or(0.0)
1591            );
1592        } else {
1593            debug!(
1594                "MaxScoreExecutor: {}ms, scored={}, skipped={}, blocks_skipped={}, groups_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1595                _bms_elapsed_ms,
1596                docs_scored,
1597                docs_skipped,
1598                blocks_skipped,
1599                groups_skipped,
1600                conjunction_skipped,
1601                results.len(),
1602                results.first().map(|r| r.score).unwrap_or(0.0)
1603            );
1604        }
1605
1606        Ok(results)
1607    }};
1608}
1609
1610impl<'a> MaxScoreExecutor<'a> {
1611    /// Create a new executor from pre-built cursors.
1612    ///
1613    /// Cursors are sorted by max_score ascending (non-essential first) and
1614    /// prefix sums are computed for the MaxScore partitioning.
1615    pub(crate) fn new(mut cursors: Vec<TermCursor<'a>>, k: usize, heap_factor: f32) -> Self {
1616        // The execution loop tracks cursors at the current document in a u64.
1617        // Query construction normally enforces this bound, but keep this
1618        // boundary defensive for direct/internal executor users as well.
1619        if cursors.len() > super::MAX_QUERY_TERMS {
1620            cursors.sort_unstable_by(|a, b| b.max_score.total_cmp(&a.max_score));
1621            cursors.truncate(super::MAX_QUERY_TERMS);
1622            log::warn!(
1623                "MaxScore cursor count exceeded {}; retaining the strongest cursors",
1624                super::MAX_QUERY_TERMS
1625            );
1626        }
1627
1628        // Enable lazy ordinal decode — ordinals are only decoded when a doc
1629        // actually reaches the scoring phase (saves ~100ns per skipped block).
1630        for c in &mut cursors {
1631            c.lazy_ordinals = true;
1632        }
1633
1634        // Sort by max_score ascending (non-essential first)
1635        cursors.sort_by(|a, b| {
1636            a.max_score
1637                .partial_cmp(&b.max_score)
1638                .unwrap_or(Ordering::Equal)
1639        });
1640
1641        let mut prefix_sums = Vec::with_capacity(cursors.len());
1642        let mut cumsum = 0.0f32;
1643        for c in &cursors {
1644            cumsum += c.max_score;
1645            prefix_sums.push(cumsum);
1646        }
1647
1648        let clamped_heap_factor = heap_factor.clamp(0.01, 1.0);
1649
1650        debug!(
1651            "Creating MaxScoreExecutor: num_cursors={}, k={}, total_upper={:.4}, heap_factor={:.2}",
1652            cursors.len(),
1653            k,
1654            cumsum,
1655            clamped_heap_factor
1656        );
1657
1658        Self {
1659            cursors,
1660            prefix_sums,
1661            collector: ScoreCollector::new(k),
1662            inv_heap_factor: 1.0 / clamped_heap_factor,
1663            predicate: None,
1664            budget: None,
1665            metric_index: "unknown",
1666            metric_field: "unknown",
1667        }
1668    }
1669
1670    /// Attach the query's wall-clock budget (anytime mode).
1671    pub fn with_budget(mut self, budget: Option<SharedThreshold>) -> Self {
1672        self.budget = budget.filter(|b| b.deadline().is_some());
1673        self
1674    }
1675
1676    /// Attach (index, field) labels for the metrics this executor emits.
1677    pub fn with_metric_labels(mut self, index: &'a str, field: &'a str) -> Self {
1678        self.metric_index = index;
1679        self.metric_field = field;
1680        self
1681    }
1682
1683    /// Create an executor for sparse vector queries.
1684    ///
1685    /// Builds `TermCursor::Sparse` for each matched dimension.
1686    pub fn sparse(
1687        sparse_index: &'a crate::segment::SparseIndex,
1688        query_terms: Vec<(u32, f32)>,
1689        k: usize,
1690        heap_factor: f32,
1691    ) -> Self {
1692        let cursors: Vec<TermCursor<'a>> = query_terms
1693            .iter()
1694            .filter_map(|&(dim_id, qw)| {
1695                let (skip_start, skip_count, global_max, block_data_offset) =
1696                    sparse_index.get_skip_range_full(dim_id)?;
1697                Some(TermCursor::sparse(
1698                    sparse_index,
1699                    qw,
1700                    skip_start,
1701                    skip_count,
1702                    global_max,
1703                    block_data_offset,
1704                ))
1705            })
1706            .collect();
1707        Self::new(cursors, k, heap_factor)
1708    }
1709
1710    /// Create an executor for full-text BM25 queries.
1711    ///
1712    /// Builds `TermCursor::Text` for each posting list.
1713    pub fn text(
1714        posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1715        avg_field_len: f32,
1716        k: usize,
1717        lengths: Option<&'a crate::segment::chunk_map::DocLengths>,
1718        params: super::Bm25Params,
1719        heap_factor: f32,
1720    ) -> Self {
1721        let cursors: Vec<TermCursor<'a>> = posting_lists
1722            .into_iter()
1723            .map(|(pl, idf)| {
1724                TermCursor::text_with_params(
1725                    pl,
1726                    idf,
1727                    avg_field_len,
1728                    lengths.map(LengthSource::Docs),
1729                    params,
1730                )
1731            })
1732            .collect();
1733        Self::new(cursors, k, heap_factor)
1734    }
1735
1736    /// Executor for BM25 over a chunked text field: posting ids are virtual
1737    /// chunk ids, scored with each chunk's real length. Results carry the
1738    /// virtual id in `doc_id`; the caller resolves it through `lengths`.
1739    pub fn text_chunked(
1740        posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1741        avg_chunk_len: f32,
1742        k: usize,
1743        lengths: &'a crate::segment::chunk_map::ChunkMap,
1744        params: super::Bm25Params,
1745        heap_factor: f32,
1746    ) -> Self {
1747        let cursors: Vec<TermCursor<'a>> = posting_lists
1748            .into_iter()
1749            .map(|(pl, idf)| {
1750                TermCursor::text_with_params(
1751                    pl,
1752                    idf,
1753                    avg_chunk_len,
1754                    Some(LengthSource::Chunks(lengths)),
1755                    params,
1756                )
1757            })
1758            .collect();
1759        Self::new(cursors, k, heap_factor)
1760    }
1761
1762    #[inline]
1763    fn find_partition(&self) -> usize {
1764        // Alpha < 1.0 raises the effective threshold → more terms become
1765        // non-essential → more aggressive pruning (approximate retrieval).
1766        // Use multiplication by reciprocal (cheaper than division).
1767        let threshold = self.collector.threshold() * self.inv_heap_factor;
1768        // Keep an equal-score candidate essential: it can still displace the
1769        // current worst hit through the deterministic doc/ordinal tie-break.
1770        self.prefix_sums.partition_point(|&sum| sum < threshold)
1771    }
1772
1773    /// Attach a per-doc predicate filter to this executor.
1774    ///
1775    /// Docs failing the predicate are skipped after block-max pruning but
1776    /// before scoring. The predicate does not affect thresholds or block-max
1777    /// comparisons — the heap stores pure sparse/text scores.
1778    pub fn with_predicate(mut self, predicate: super::DocPredicate<'a>) -> Self {
1779        self.predicate = Some(predicate);
1780        self
1781    }
1782
1783    /// Seed the collector with an initial threshold for tighter early pruning.
1784    pub fn seed_threshold(&mut self, initial_threshold: f32) {
1785        self.collector.seed_threshold(initial_threshold);
1786    }
1787
1788    /// Execute Block-Max MaxScore and return top-k results (async).
1789    ///
1790    /// Text cursors (in-memory posting lists) run the windowed executor;
1791    /// sparse cursors, whose blocks may need asynchronous I/O, run the
1792    /// document-at-a-time loop.
1793    pub async fn execute(mut self) -> crate::Result<Vec<ScoredDoc>> {
1794        if self.cursors.is_empty() {
1795            return Ok(Vec::new());
1796        }
1797        let t = crate::observe::Timer::start();
1798        let results = if self.all_text() {
1799            self.execute_windowed()
1800        } else {
1801            bms_execute_loop!(self, ensure_block_loaded, advance, seek, .await)
1802        };
1803        if let Ok(r) = &results {
1804            crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1805        }
1806        results
1807    }
1808
1809    /// Synchronous execution — works when all cursors are text or mmap-backed sparse.
1810    pub fn execute_sync(mut self) -> crate::Result<Vec<ScoredDoc>> {
1811        if self.cursors.is_empty() {
1812            return Ok(Vec::new());
1813        }
1814        let t = crate::observe::Timer::start();
1815        let results = if self.all_text() {
1816            self.execute_windowed()
1817        } else {
1818            bms_execute_loop!(self, ensure_block_loaded_sync, advance_sync, seek_sync,)
1819        };
1820        if let Ok(r) = &results {
1821            crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1822        }
1823        results
1824    }
1825
1826    /// The document-at-a-time loop on any cursors (the reference the
1827    /// windowed executor is checked against in tests).
1828    #[cfg(test)]
1829    pub(crate) fn execute_doc_at_a_time_sync(mut self) -> crate::Result<Vec<ScoredDoc>> {
1830        if self.cursors.is_empty() {
1831            return Ok(Vec::new());
1832        }
1833        bms_execute_loop!(self, ensure_block_loaded_sync, advance_sync, seek_sync,)
1834    }
1835
1836    fn all_text(&self) -> bool {
1837        self.cursors.iter().all(TermCursor::is_text)
1838    }
1839
1840    /// Window-at-a-time Block-Max MaxScore for text cursors.
1841    ///
1842    /// The id space is walked in windows of at most [`WINDOW_IDS`] ids that
1843    /// start at the first id a globally essential cursor can still reach and
1844    /// end at the smallest current block end among those cursors. Per window
1845    /// (Lucene `MaxScoreBulkScorer`; turbopuffer "batched iterator
1846    /// advancement"):
1847    ///
1848    /// 1. every cursor's bound over the window is read from its skip entries
1849    ///    (`window_upper_bound`), and the cursors are re-partitioned into
1850    ///    essential and non-essential by those bounds, so the partition is
1851    ///    the block-max one, not the list-max one;
1852    /// 2. a window whose summed bounds cannot reach the threshold is skipped
1853    ///    by every cursor without decoding a block;
1854    /// 3. each essential cursor scores all of its postings in the window in
1855    ///    one pass over its decoded block (dense `scores[id - from]` buffer
1856    ///    plus a match bitset), so the same iterator advances many times in
1857    ///    a row instead of alternating with the others;
1858    /// 4. the candidates are filtered branch-free against the threshold
1859    ///    minus what the remaining cursors could still add, and each
1860    ///    non-essential cursor is then sought to the survivors in id order
1861    ///    (again one iterator at a time), strongest bound first.
1862    ///
1863    /// Rank-safe: only documents whose window-essential score plus the
1864    /// non-essential bounds cannot reach the threshold are dropped. The
1865    /// approximate `heap_factor` mode scales the threshold as in the
1866    /// document-at-a-time loop.
1867    pub(crate) fn execute_windowed(&mut self) -> crate::Result<Vec<ScoredDoc>> {
1868        if self
1869            .budget
1870            .as_ref()
1871            .is_some_and(SharedThreshold::stop_if_expired)
1872        {
1873            return Ok(Vec::new());
1874        }
1875        let n = self.cursors.len();
1876        for cursor in &mut self.cursors {
1877            cursor.ensure_block_loaded_sync()?;
1878        }
1879        let inv_heap_factor = self.inv_heap_factor;
1880        let mut window_scores = vec![0.0f32; WINDOW_IDS];
1881        let mut window_mask = vec![0u64; WINDOW_IDS / 64];
1882        let mut cand_docs: Vec<u32> = Vec::with_capacity(WINDOW_IDS);
1883        let mut cand_scores: Vec<f32> = Vec::with_capacity(WINDOW_IDS);
1884        let mut wmax = vec![0.0f32; n];
1885        let mut order: Vec<usize> = (0..n).collect();
1886        let mut wprefix = vec![0.0f32; n];
1887        let mut windows = 0u64;
1888        let mut windows_skipped = 0u64;
1889        let mut candidates = 0u64;
1890        let mut docs_scored = 0u64;
1891        let started = std::time::Instant::now();
1892
1893        loop {
1894            windows += 1;
1895            if windows & 0x3F == 0
1896                && let Some(budget) = &self.budget
1897                && budget.expired()
1898            {
1899                budget.mark_truncated();
1900                log::debug!(
1901                    "MaxScoreExecutor(windowed): deadline reached after {} windows, {} scored",
1902                    windows,
1903                    docs_scored
1904                );
1905                break;
1906            }
1907            let partition = self.find_partition();
1908            if partition >= n {
1909                break;
1910            }
1911            // Window: from the first id a globally essential cursor can
1912            // still reach to the smallest current block end among them.
1913            let mut from = u32::MAX;
1914            let mut to = u32::MAX;
1915            for cursor in &self.cursors[partition..] {
1916                if cursor.exhausted {
1917                    continue;
1918                }
1919                from = from.min(cursor.doc());
1920                to = to.min(cursor.block_last_doc(cursor.block_idx));
1921            }
1922            if from == u32::MAX {
1923                break;
1924            }
1925            let to = to.max(from).min(from.saturating_add(WINDOW_IDS as u32 - 1));
1926            let width = (to - from) as usize + 1;
1927            let words = width.div_ceil(64);
1928
1929            // Block-max partition over the window.
1930            let heap_full = self.collector.len() >= self.collector.k;
1931            let threshold = if heap_full {
1932                self.collector.threshold() * inv_heap_factor - 1e-6
1933            } else {
1934                0.0
1935            };
1936            for (i, bound) in wmax.iter_mut().enumerate() {
1937                *bound = self.cursors[i].window_upper_bound(from, to);
1938            }
1939            order.sort_unstable_by(|&a, &b| wmax[a].total_cmp(&wmax[b]));
1940            let mut sum = 0.0f32;
1941            for (rank, &i) in order.iter().enumerate() {
1942                sum += wmax[i];
1943                wprefix[rank] = sum;
1944            }
1945            let wpartition = if heap_full {
1946                wprefix.partition_point(|&s| s < threshold)
1947            } else {
1948                0
1949            };
1950            if wpartition >= n {
1951                // Nothing in the window can compete: every cursor jumps past it.
1952                for cursor in &mut self.cursors {
1953                    if !cursor.exhausted && cursor.doc() <= to {
1954                        cursor.skip_past_sync(to)?;
1955                    }
1956                }
1957                windows_skipped += 1;
1958                continue;
1959            }
1960
1961            // Essential cursors: bulk-score into the window buffers.
1962            window_scores[..width].fill(0.0);
1963            window_mask[..words].fill(0);
1964            for &i in &order[wpartition..] {
1965                let cursor = &mut self.cursors[i];
1966                if cursor.exhausted {
1967                    continue;
1968                }
1969                if cursor.doc() < from {
1970                    cursor.seek_sync(from)?;
1971                }
1972                if cursor.exhausted || cursor.doc() > to {
1973                    continue;
1974                }
1975                cursor.score_window_sync(
1976                    from,
1977                    to,
1978                    &mut window_scores[..width],
1979                    &mut window_mask[..words],
1980                )?;
1981            }
1982
1983            // Candidates in id order.
1984            cand_docs.clear();
1985            cand_scores.clear();
1986            for (word_idx, word) in window_mask[..words].iter().enumerate() {
1987                let mut bits = *word;
1988                while bits != 0 {
1989                    let slot = (word_idx << 6) | bits.trailing_zeros() as usize;
1990                    bits &= bits - 1;
1991                    cand_docs.push(from + slot as u32);
1992                    cand_scores.push(window_scores[slot]);
1993                }
1994            }
1995            if let Some(pred) = &self.predicate {
1996                let mut kept = 0usize;
1997                for j in 0..cand_docs.len() {
1998                    let doc = cand_docs[j];
1999                    cand_docs[kept] = doc;
2000                    cand_scores[kept] = cand_scores[j];
2001                    kept += pred(doc) as usize;
2002                }
2003                cand_docs.truncate(kept);
2004                cand_scores.truncate(kept);
2005            }
2006
2007            // Non-essential cursors on the survivors, strongest bound first.
2008            let mut remaining = if wpartition > 0 {
2009                wprefix[wpartition - 1]
2010            } else {
2011                0.0
2012            };
2013            for rank in (0..wpartition).rev() {
2014                let i = order[rank];
2015                if heap_full {
2016                    filter_competitive(&mut cand_docs, &mut cand_scores, remaining, threshold);
2017                }
2018                if cand_docs.is_empty() {
2019                    break;
2020                }
2021                if wmax[i] > 0.0 {
2022                    let cursor = &mut self.cursors[i];
2023                    for (doc, score) in cand_docs.iter().zip(cand_scores.iter_mut()) {
2024                        if cursor.seek_sync(*doc)? == *doc {
2025                            cursor.ensure_scores();
2026                            *score += cursor.score();
2027                        }
2028                    }
2029                }
2030                remaining -= wmax[i];
2031            }
2032            if heap_full {
2033                filter_competitive(&mut cand_docs, &mut cand_scores, 0.0, threshold);
2034            }
2035            candidates += cand_docs.len() as u64;
2036            for (doc, score) in cand_docs.iter().zip(&cand_scores) {
2037                if self.collector.insert_with_ordinal(*doc, *score, 0) {
2038                    docs_scored += 1;
2039                }
2040            }
2041        }
2042
2043        let collector = std::mem::replace(&mut self.collector, ScoreCollector::new(0));
2044        let results: Vec<ScoredDoc> = collector
2045            .into_sorted_results()
2046            .into_iter()
2047            .map(|(doc_id, score, ordinal)| ScoredDoc {
2048                doc_id,
2049                score,
2050                ordinal,
2051            })
2052            .collect();
2053        let elapsed_ms = started.elapsed().as_millis() as u64;
2054        if elapsed_ms > 500 {
2055            warn!(
2056                "slow windowed MaxScore: {}ms, cursors={}, windows={}, windows_skipped={}, candidates={}, scored={}, returned={}, top_score={:.4}",
2057                elapsed_ms,
2058                n,
2059                windows,
2060                windows_skipped,
2061                candidates,
2062                docs_scored,
2063                results.len(),
2064                results.first().map(|r| r.score).unwrap_or(0.0)
2065            );
2066        } else {
2067            debug!(
2068                "MaxScoreExecutor(windowed): {}ms, cursors={}, windows={}, windows_skipped={}, candidates={}, scored={}, returned={}, top_score={:.4}",
2069                elapsed_ms,
2070                n,
2071                windows,
2072                windows_skipped,
2073                candidates,
2074                docs_scored,
2075                results.len(),
2076                results.first().map(|r| r.score).unwrap_or(0.0)
2077            );
2078        }
2079        Ok(results)
2080    }
2081}
2082
2083/// Ids per window of the windowed executor (Lucene's `INNER_WINDOW_SIZE`).
2084const WINDOW_IDS: usize = 4096;
2085
2086/// Keep the candidates that can still reach `threshold` once `remaining`
2087/// (the bounds of the cursors not yet applied) is added. Written without a
2088/// data-dependent branch, like Lucene's `VectorUtil.filterByScore`.
2089fn filter_competitive(docs: &mut Vec<u32>, scores: &mut Vec<f32>, remaining: f32, threshold: f32) {
2090    let mut kept = 0usize;
2091    for j in 0..docs.len() {
2092        let doc = docs[j];
2093        let score = scores[j];
2094        docs[kept] = doc;
2095        scores[kept] = score;
2096        kept += (score + remaining >= threshold) as usize;
2097    }
2098    docs.truncate(kept);
2099    scores.truncate(kept);
2100}
2101
2102#[cfg(test)]
2103mod tests {
2104    use super::*;
2105
2106    // ── Windowed executor parity ─────────────────────────────────────────
2107
2108    struct Corpus {
2109        /// Per term: sorted `(doc, tf)` postings.
2110        postings: Vec<Vec<(u32, u32)>>,
2111        lengths: Vec<u16>,
2112        n_docs: u32,
2113    }
2114
2115    fn xorshift(state: &mut u64) -> u64 {
2116        *state ^= *state << 13;
2117        *state ^= *state >> 7;
2118        *state ^= *state << 17;
2119        *state
2120    }
2121
2122    /// Terms with very different densities (from 0.5% to 60% of the
2123    /// documents), skewed term frequencies, and pseudo-random lengths.
2124    fn random_corpus(seed: u64, n_docs: u32, n_terms: usize) -> Corpus {
2125        let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
2126        let lengths: Vec<u16> = (0..n_docs)
2127            .map(|_| 1 + (xorshift(&mut state) % 400) as u16)
2128            .collect();
2129        let densities = [0.6, 0.25, 0.1, 0.03, 0.005];
2130        let postings = (0..n_terms)
2131            .map(|t| {
2132                let density = densities[t % densities.len()];
2133                let cutoff = (density * u32::MAX as f64) as u64;
2134                let mut postings = Vec::new();
2135                for doc in 0..n_docs {
2136                    if (xorshift(&mut state) & 0xFFFF_FFFF) >= cutoff {
2137                        continue;
2138                    }
2139                    let r = xorshift(&mut state) % 100;
2140                    let tf = if r < 70 {
2141                        1
2142                    } else if r < 90 {
2143                        2
2144                    } else {
2145                        3 + (r % 6) as u32
2146                    };
2147                    postings.push((doc, tf));
2148                }
2149                postings
2150            })
2151            .collect();
2152        Corpus {
2153            postings,
2154            lengths,
2155            n_docs,
2156        }
2157    }
2158
2159    fn build_lists(
2160        corpus: &Corpus,
2161        lengths: Option<&crate::segment::chunk_map::DocLengths>,
2162    ) -> Vec<(crate::structures::BlockPostingList, f32)> {
2163        corpus
2164            .postings
2165            .iter()
2166            .map(|postings| {
2167                let mut list = crate::structures::PostingList::new();
2168                for &(doc, tf) in postings {
2169                    list.push(doc, tf);
2170                }
2171                let length_of = lengths.map(|l| move |doc: DocId| l.length(doc));
2172                let block_list = crate::structures::BlockPostingList::from_posting_list_with(
2173                    &list,
2174                    false,
2175                    length_of.as_ref().map(|f| f as &dyn Fn(DocId) -> u32),
2176                )
2177                .unwrap();
2178                let idf = super::super::bm25_idf(postings.len() as f32, corpus.n_docs as f32);
2179                (block_list, idf)
2180            })
2181            .collect()
2182    }
2183
2184    /// Exhaustive per-document scores with the same formula the cursors use.
2185    fn exhaustive(
2186        corpus: &Corpus,
2187        lists: &[(crate::structures::BlockPostingList, f32)],
2188        real_lengths: bool,
2189        avg: f32,
2190        params: super::super::Bm25Params,
2191    ) -> std::collections::HashMap<u32, f32> {
2192        let mut scores: std::collections::HashMap<u32, f32> = std::collections::HashMap::new();
2193        for (postings, (_, idf)) in corpus.postings.iter().zip(lists) {
2194            for &(doc, tf) in postings {
2195                let len = if real_lengths {
2196                    corpus.lengths[doc as usize] as f32
2197                } else {
2198                    tf as f32
2199                };
2200                *scores.entry(doc).or_insert(0.0) += params.score(tf as f32, *idf, len, avg);
2201            }
2202        }
2203        scores
2204    }
2205
2206    fn check_top_k(
2207        label: &str,
2208        results: &[ScoredDoc],
2209        exhaustive: &std::collections::HashMap<u32, f32>,
2210        k: usize,
2211        predicate: Option<&dyn Fn(u32) -> bool>,
2212    ) {
2213        let mut expected: Vec<(u32, f32)> = exhaustive
2214            .iter()
2215            .filter(|(doc, _)| predicate.is_none_or(|p| p(**doc)))
2216            .map(|(doc, score)| (*doc, *score))
2217            .collect();
2218        expected.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
2219        let want = k.min(expected.len());
2220        assert_eq!(results.len(), want, "{label}: result count");
2221        for (rank, (got, exp)) in results.iter().zip(&expected).enumerate() {
2222            let tolerance = 1e-4 * exp.1.abs().max(1.0);
2223            assert!(
2224                (got.score - exp.1).abs() <= tolerance,
2225                "{label}: rank {rank} score {} vs exhaustive {} (doc {} vs {})",
2226                got.score,
2227                exp.1,
2228                got.doc_id,
2229                exp.0
2230            );
2231            let own = exhaustive[&got.doc_id];
2232            assert!(
2233                (got.score - own).abs() <= tolerance,
2234                "{label}: doc {} scored {} but exhaustive says {}",
2235                got.doc_id,
2236                got.score,
2237                own
2238            );
2239            if let Some(p) = predicate {
2240                assert!(
2241                    p(got.doc_id),
2242                    "{label}: doc {} fails the predicate",
2243                    got.doc_id
2244                );
2245            }
2246        }
2247        for pair in results.windows(2) {
2248            assert!(
2249                pair[0].score >= pair[1].score,
2250                "{label}: results not sorted"
2251            );
2252        }
2253    }
2254
2255    /// The windowed executor returns the exact top-k (against an exhaustive
2256    /// scorer and against the document-at-a-time loop) over corpora of
2257    /// different sizes and term mixes, with and without real lengths,
2258    /// predicates, and a seeded threshold.
2259    #[test]
2260    fn windowed_text_maxscore_matches_exhaustive_and_doc_at_a_time() {
2261        let params = super::super::Bm25Params::default();
2262        let predicate_fn = |doc: u32| !doc.is_multiple_of(3);
2263        let mut cases = 0usize;
2264        for seed in 1..=6u64 {
2265            for &n_docs in &[300u32, 2_500, 12_000] {
2266                for &n_terms in &[1usize, 2, 4, 9] {
2267                    let corpus = random_corpus(seed, n_docs, n_terms);
2268                    let doc_lengths =
2269                        crate::segment::chunk_map::DocLengths::from_lengths(&corpus.lengths);
2270                    for real_lengths in [true, false] {
2271                        let lengths = real_lengths.then_some(&doc_lengths);
2272                        let lists = build_lists(&corpus, lengths);
2273                        let avg = if real_lengths {
2274                            doc_lengths.avg_len()
2275                        } else {
2276                            1.0
2277                        };
2278                        let truth = exhaustive(&corpus, &lists, real_lengths, avg, params);
2279                        for &k in &[1usize, 10, 100] {
2280                            for with_predicate in [false, true] {
2281                                let label = format!(
2282                                    "seed={seed} docs={n_docs} terms={n_terms} lengths={real_lengths} k={k} pred={with_predicate}"
2283                                );
2284                                let pred: Option<&dyn Fn(u32) -> bool> =
2285                                    with_predicate.then_some(&predicate_fn);
2286                                let make = |seeded: f32| {
2287                                    let mut executor = MaxScoreExecutor::text(
2288                                        lists.clone(),
2289                                        avg,
2290                                        k,
2291                                        lengths,
2292                                        params,
2293                                        1.0,
2294                                    );
2295                                    if with_predicate {
2296                                        executor = executor.with_predicate(Box::new(predicate_fn));
2297                                    }
2298                                    if seeded > 0.0 {
2299                                        executor.seed_threshold(seeded);
2300                                    }
2301                                    executor
2302                                };
2303                                let windowed = make(0.0).execute_windowed().unwrap();
2304                                check_top_k(
2305                                    &format!("windowed {label}"),
2306                                    &windowed,
2307                                    &truth,
2308                                    k,
2309                                    pred,
2310                                );
2311                                let reference = make(0.0).execute_doc_at_a_time_sync().unwrap();
2312                                check_top_k(
2313                                    &format!("reference {label}"),
2314                                    &reference,
2315                                    &truth,
2316                                    k,
2317                                    pred,
2318                                );
2319                                // A floor below the k-th score keeps the exact top-k.
2320                                if let Some(kth) = windowed.last().map(|r| r.score)
2321                                    && windowed.len() == k
2322                                {
2323                                    let seeded = make(kth * 0.9).execute_windowed().unwrap();
2324                                    check_top_k(
2325                                        &format!("seeded {label}"),
2326                                        &seeded,
2327                                        &truth,
2328                                        k,
2329                                        pred,
2330                                    );
2331                                    // A floor above every score returns nothing.
2332                                    let above =
2333                                        make(windowed[0].score * 1.5).execute_windowed().unwrap();
2334                                    assert!(above.is_empty(), "{label}: floor above all scores");
2335                                }
2336                                cases += 1;
2337                            }
2338                        }
2339                    }
2340                }
2341            }
2342        }
2343        assert!(cases > 400);
2344    }
2345
2346    /// The approximate mode returns a subset of the exact top-k with exact
2347    /// scores.
2348    #[test]
2349    fn windowed_text_maxscore_heap_factor_is_a_subset_with_exact_scores() {
2350        let params = super::super::Bm25Params::default();
2351        let corpus = random_corpus(7, 20_000, 6);
2352        let doc_lengths = crate::segment::chunk_map::DocLengths::from_lengths(&corpus.lengths);
2353        let lists = build_lists(&corpus, Some(&doc_lengths));
2354        let avg = doc_lengths.avg_len();
2355        let truth = exhaustive(&corpus, &lists, true, avg, params);
2356        let exact = MaxScoreExecutor::text(lists.clone(), avg, 50, Some(&doc_lengths), params, 1.0)
2357            .execute_windowed()
2358            .unwrap();
2359        check_top_k("exact", &exact, &truth, 50, None);
2360        let approx = MaxScoreExecutor::text(lists, avg, 50, Some(&doc_lengths), params, 0.6)
2361            .execute_windowed()
2362            .unwrap();
2363        assert_eq!(approx.len(), 50);
2364        for hit in &approx {
2365            let own = truth[&hit.doc_id];
2366            assert!((hit.score - own).abs() <= 1e-4 * own.max(1.0));
2367        }
2368        // The usual heap-factor guarantee: every returned score is within the
2369        // factor of the exact k-th score, and the best document is exact.
2370        let exact_kth = exact.last().unwrap().score;
2371        assert!(approx.iter().all(|hit| hit.score >= exact_kth * 0.6 - 1e-4));
2372        assert_eq!(approx[0].doc_id, exact[0].doc_id);
2373        let overlap = approx
2374            .iter()
2375            .filter(|hit| exact.iter().any(|e| e.doc_id == hit.doc_id))
2376            .count();
2377        assert!(overlap >= 25, "overlap {overlap} of 50");
2378    }
2379
2380    #[test]
2381    fn text_heap_factor_below_one_actually_changes_pruning() {
2382        let params = super::super::Bm25Params::default();
2383        let corpus = random_corpus(7, 20_000, 6);
2384        let lengths = crate::segment::chunk_map::DocLengths::from_lengths(&corpus.lengths);
2385        let lists = build_lists(&corpus, Some(&lengths));
2386        let mut exact = MaxScoreExecutor::text(
2387            lists.clone(),
2388            lengths.avg_len(),
2389            50,
2390            Some(&lengths),
2391            params,
2392            1.0,
2393        );
2394        let mut approximate =
2395            MaxScoreExecutor::text(lists, lengths.avg_len(), 50, Some(&lengths), params, 0.01);
2396        let exact = exact.execute_windowed().unwrap();
2397        let approximate = approximate.execute_windowed().unwrap();
2398        assert_ne!(
2399            exact.iter().map(|hit| hit.doc_id).collect::<Vec<_>>(),
2400            approximate.iter().map(|hit| hit.doc_id).collect::<Vec<_>>()
2401        );
2402    }
2403
2404    #[test]
2405    fn text_and_sparse_factors_have_identical_conventions() {
2406        for factor in [0.0, 0.01, 0.5, 1.0] {
2407            let text = MaxScoreExecutor::text(
2408                Vec::new(),
2409                1.0,
2410                1,
2411                None,
2412                super::super::Bm25Params::default(),
2413                factor,
2414            );
2415            let sparse = MaxScoreExecutor::new(Vec::new(), 1, factor);
2416            assert_eq!(text.inv_heap_factor, sparse.inv_heap_factor);
2417            assert_eq!(text.inv_heap_factor, factor.clamp(0.01, 1.0).recip());
2418        }
2419    }
2420
2421    #[test]
2422    fn chunk_bounds_use_the_same_length_floor_as_scoring() {
2423        use crate::segment::chunk_map::{ChunkMapBuilder, read_chunk_maps, write_chunk_maps};
2424        use crate::structures::{BlockPostingList, PostingList};
2425        let mut chunks = ChunkMapBuilder::default();
2426        let mut postings = PostingList::new();
2427        for id in 0..2048 {
2428            chunks
2429                .push(id / 4, (id % 4) as u16, if id % 4 == 0 { 1 } else { 200 })
2430                .unwrap();
2431            postings.push(id, 1);
2432        }
2433        let mut bytes = Vec::new();
2434        write_chunk_maps(&mut bytes, &[(0, &chunks)], &[]).unwrap();
2435        let maps = read_chunk_maps(crate::directories::OwnedBytes::new(bytes)).unwrap();
2436        let map = &maps.chunk_maps[&0];
2437        let list =
2438            BlockPostingList::from_posting_list_with(&postings, false, Some(&|id| map.length(id)))
2439                .unwrap();
2440        assert_eq!(list.min_len(), Some(1));
2441        let params = super::super::Bm25Params::default();
2442        let cursor = TermCursor::text_with_params(
2443            list,
2444            2.0,
2445            map.avg_len(),
2446            Some(LengthSource::Chunks(map)),
2447            params,
2448        );
2449        let expected = params.upper_bound_with_len(1.0, 2.0, 200.0, map.avg_len());
2450        assert_eq!(cursor.max_score.to_bits(), expected.to_bits());
2451        for block in 0..cursor.num_blocks {
2452            assert_eq!(cursor.text_block_bound(block).to_bits(), expected.to_bits());
2453            assert_eq!(
2454                cursor.text_group_bound(block).unwrap().to_bits(),
2455                expected.to_bits()
2456            );
2457        }
2458    }
2459
2460    #[test]
2461    fn nested_scorers_keep_deadline_but_not_outer_score_floor() {
2462        let shared = SharedThreshold::for_limit(10).with_deadline(Some(std::time::Instant::now()));
2463        shared.raise(100.0);
2464        let options = super::super::ScorerOptions {
2465            shared_threshold: Some(shared.clone()),
2466            initial_threshold: 100.0,
2467            ..Default::default()
2468        }
2469        .without_threshold();
2470        assert_eq!(options.initial_threshold, 0.0);
2471        let nested = options.shared_threshold.unwrap();
2472        assert_eq!(nested.get(), 0.0);
2473        assert!(!nested.covers(10));
2474        assert!(nested.stop_if_expired());
2475        assert!(shared.truncated());
2476        nested.raise(200.0);
2477        assert_eq!(shared.get(), 100.0);
2478    }
2479
2480    #[test]
2481    fn score_bound_cache_reuses_only_its_last_key() {
2482        let cache = CachedScoreBound::new();
2483        assert_eq!(cache.get_or_compute(0, || 1.5), 1.5);
2484        assert_eq!(
2485            cache.get_or_compute(0, || panic!("recomputed active bound")),
2486            1.5
2487        );
2488        assert_eq!(cache.get_or_compute(1, || 2.5), 2.5);
2489        assert_eq!(cache.get_or_compute(0, || 3.5), 3.5);
2490    }
2491
2492    #[test]
2493    fn test_shared_threshold_monotonic_raise() {
2494        let shared = SharedThreshold::new();
2495        assert_eq!(shared.get(), 0.0);
2496
2497        shared.raise(2.5);
2498        assert_eq!(shared.get(), 2.5);
2499
2500        // Lower values never lower the floor.
2501        shared.raise(1.0);
2502        assert_eq!(shared.get(), 2.5);
2503
2504        // Higher values raise it.
2505        shared.raise(4.0);
2506        assert_eq!(shared.get(), 4.0);
2507
2508        // Non-positive and NaN are ignored.
2509        shared.raise(0.0);
2510        shared.raise(-3.0);
2511        shared.raise(f32::NAN);
2512        assert_eq!(shared.get(), 4.0);
2513
2514        // Clones share the same atomic cell.
2515        let clone = shared.clone();
2516        clone.raise(9.0);
2517        assert_eq!(shared.get(), 9.0);
2518    }
2519
2520    #[test]
2521    fn test_shared_threshold_seed_matches_manual() {
2522        // A collector seeded with a floor prunes anything at/below it, matching
2523        // the threshold a fully-populated heap would have produced.
2524        let mut seeded = ScoreCollector::new(2);
2525        seeded.seed_threshold(3.0);
2526        assert_eq!(seeded.threshold(), 3.0);
2527        // A score at/below the floor cannot enter.
2528        assert!(!seeded.would_enter(3.0));
2529        assert!(seeded.would_enter(3.5));
2530        // Real inserts above the floor evict the sentinels; results contain no
2531        // sentinel (doc_id == u32::MAX) entries.
2532        seeded.insert(1, 5.0);
2533        seeded.insert(2, 4.0);
2534        let results = seeded.into_sorted_results();
2535        assert_eq!(results.len(), 2);
2536        assert_eq!(results[0].0, 1);
2537        assert_eq!(results[1].0, 2);
2538    }
2539
2540    #[test]
2541    fn test_shared_threshold_can_raise_after_real_inserts() {
2542        let mut collector = ScoreCollector::new(3);
2543        collector.insert(1, 10.0);
2544        collector.insert(2, 4.0);
2545        assert_eq!(collector.real_len(), 2);
2546
2547        // Raising the floor after traversal has started removes retained work
2548        // that can no longer reach the global top-k.
2549        collector.seed_threshold(6.0);
2550        assert_eq!(collector.threshold(), 6.0);
2551        assert_eq!(collector.real_len(), 1);
2552
2553        // A real candidate tied with the floor displaces the sentinel because
2554        // its doc id wins the canonical tie break.
2555        assert!(collector.would_enter_candidate(3, 6.0, 0));
2556        assert!(collector.insert(3, 6.0));
2557        assert_eq!(collector.real_len(), 2);
2558        let results = collector.into_sorted_results();
2559        assert_eq!(results, vec![(1, 10.0, 0), (3, 6.0, 0)]);
2560    }
2561
2562    #[test]
2563    fn test_large_seed_uses_virtual_sentinels() {
2564        let k = 1_000_000_000;
2565        let mut collector = ScoreCollector::new(k);
2566        assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
2567
2568        collector.seed_threshold(42.0);
2569
2570        // Seeding a huge top-k is constant-time and does not materialize any
2571        // of its conceptual sentinel entries.
2572        assert_eq!(collector.heap.len(), 0);
2573        assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
2574        assert_eq!(collector.len(), k);
2575        assert_eq!(collector.real_len(), 0);
2576        assert_eq!(collector.threshold(), 42.0);
2577        assert!(!collector.is_empty());
2578
2579        // A real result tied with the floor beats the sentinel by doc-id, while
2580        // a lower score remains below the conceptual threshold.
2581        assert!(collector.insert_with_ordinal(9, 42.0, 7));
2582        assert!(!collector.insert(10, 41.0));
2583        assert!(collector.insert(11, 43.0));
2584        assert_eq!(collector.len(), k);
2585        assert_eq!(collector.real_len(), 2);
2586        assert_eq!(
2587            collector.into_sorted_results(),
2588            vec![(11, 43.0, 0), (9, 42.0, 7)]
2589        );
2590    }
2591
2592    #[test]
2593    fn test_virtual_sentinels_preserve_tie_order_when_filled() {
2594        let mut collector = ScoreCollector::new(3);
2595        collector.seed_threshold(5.0);
2596
2597        assert!(collector.insert_with_ordinal(3, 5.0, 2));
2598        assert!(collector.insert_with_ordinal(2, 5.0, 8));
2599        assert!(collector.insert_with_ordinal(1, 5.0, 4));
2600        assert_eq!(collector.real_len(), 3);
2601        assert!(collector.virtual_threshold.is_none());
2602
2603        // Once all virtual slots have been displaced, canonical doc/ordinal
2604        // ordering still controls root replacement at an equal score.
2605        assert!(collector.insert_with_ordinal(2, 5.0, 1));
2606        assert!(!collector.insert_with_ordinal(4, 5.0, 0));
2607        assert_eq!(
2608            collector.into_sorted_results(),
2609            vec![(1, 5.0, 4), (2, 5.0, 1), (2, 5.0, 8)]
2610        );
2611    }
2612
2613    #[test]
2614    fn test_score_collector_basic() {
2615        let mut collector = ScoreCollector::new(3);
2616
2617        collector.insert(1, 1.0);
2618        collector.insert(2, 2.0);
2619        collector.insert(3, 3.0);
2620        assert_eq!(collector.threshold(), 1.0);
2621
2622        collector.insert(4, 4.0);
2623        assert_eq!(collector.threshold(), 2.0);
2624
2625        let results = collector.into_sorted_results();
2626        assert_eq!(results.len(), 3);
2627        assert_eq!(results[0].0, 4); // Highest score
2628        assert_eq!(results[1].0, 3);
2629        assert_eq!(results[2].0, 2);
2630    }
2631
2632    #[test]
2633    fn test_score_collector_threshold() {
2634        let mut collector = ScoreCollector::new(2);
2635
2636        collector.insert(1, 5.0);
2637        collector.insert(2, 3.0);
2638        assert_eq!(collector.threshold(), 3.0);
2639
2640        // Should not enter (score too low)
2641        assert!(!collector.would_enter(2.0));
2642        assert!(!collector.insert(3, 2.0));
2643
2644        // Should enter (score high enough)
2645        assert!(collector.would_enter(4.0));
2646        assert!(collector.insert(4, 4.0));
2647        assert_eq!(collector.threshold(), 4.0);
2648    }
2649
2650    #[test]
2651    fn test_heap_entry_ordering() {
2652        let mut heap = BinaryHeap::new();
2653        heap.push(HeapEntry {
2654            doc_id: 1,
2655            score: 3.0,
2656            ordinal: 0,
2657        });
2658        heap.push(HeapEntry {
2659            doc_id: 2,
2660            score: 1.0,
2661            ordinal: 0,
2662        });
2663        heap.push(HeapEntry {
2664            doc_id: 3,
2665            score: 2.0,
2666            ordinal: 0,
2667        });
2668
2669        // Min-heap: lowest score should come out first
2670        assert_eq!(heap.pop().unwrap().score, 1.0);
2671        assert_eq!(heap.pop().unwrap().score, 2.0);
2672        assert_eq!(heap.pop().unwrap().score, 3.0);
2673    }
2674}