Skip to main content

zeppelin_embed/fts/
search.rs

1//! The exhaustive-OR scorer: task 13's deliverable and task 14's oracle.
2//!
3//! This walks every posting of every query term and scores every document
4//! that matches at least one of them. It is deliberately the simplest
5//! correct thing. Task 14 adds block-max pruning that must return exactly
6//! what this returns — same ids, same scores, same tie-break — and this
7//! module stays forever as the reference that claim is checked against.
8//!
9//! # Tie-break, pinned
10//!
11//! Results are ordered by descending score, then by ascending global
12//! document id. The tie-break is part of the contract because task 14's
13//! equivalence property compares result vectors element-wise: an unpinned
14//! tie-break would make pruning look wrong when it was merely differently
15//! ordered, or hide a real reordering bug behind a sort.
16//!
17//! # Global document ids
18//!
19//! A [`GlobalDocId`](crate::fts::search::GlobalDocId) pairs a segment ordinal
20//! with a segment-local row. Row ids are dense and segment-local by task 07
21//! invariant, so a store-wide identity has to carry both. Ordering is by
22//! segment then row, which is insertion order, which is time order.
23
24use std::borrow::Cow;
25
26use super::bm25::{Bm25Params, Df, DocLen, TermScorer, Tf};
27use super::index::{FieldId, IndexError, LexicalIndex};
28use super::sealed::{SealedSegment, TermStream};
29
30/// A store-wide document identity.
31#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
32pub struct GlobalDocId {
33    /// Index of the segment in seal order.
34    pub segment: u32,
35    /// Dense segment-local row id.
36    pub row: u32,
37}
38
39/// One scored result.
40#[derive(Clone, Copy, Debug, PartialEq)]
41pub struct ScoredDoc {
42    /// The document.
43    pub doc: GlobalDocId,
44    /// Its BM25 score. Higher is better; never negative.
45    pub score: f64,
46}
47
48/// Per-field weight applied to term frequencies and lengths.
49///
50/// # One entry per field, ascending
51///
52/// The table is normalized at construction: sorted by field id, with a
53/// repeated field collapsed to its first weight. Without that invariant a
54/// caller who names a field twice makes the scorer count that field's length
55/// twice, which is a silently wrong score rather than a rejected query, and
56/// it makes iterating the table differ from looking each field up by id.
57#[derive(Clone, Debug, Eq, PartialEq)]
58pub struct FieldWeights {
59    weights: Vec<(FieldId, u32)>,
60}
61
62impl FieldWeights {
63    /// Equal weight over the supplied fields: the flat scorer.
64    ///
65    /// This is what reproduces the flat BEIR column.
66    #[must_use]
67    pub fn flat(fields: &[FieldId]) -> Self {
68        Self::normalized(fields.iter().map(|field| (*field, 1_000)).collect())
69    }
70
71    /// Explicit per-field weights, in thousandths.
72    ///
73    /// Integer thousandths rather than floats so the weight table is exactly
74    /// comparable and a configuration digest over it is stable.
75    #[must_use]
76    pub fn new(weights: &[(FieldId, u32)]) -> Self {
77        Self::normalized(weights.to_vec())
78    }
79
80    /// Sorts by field and collapses a repeated field to its first weight.
81    fn normalized(mut weights: Vec<(FieldId, u32)>) -> Self {
82        // A stable sort keeps "first weight wins" meaning first as the
83        // caller wrote it, which is what `weight` has always returned.
84        weights.sort_by_key(|(field, _)| *field);
85        weights.dedup_by_key(|(field, _)| *field);
86        Self { weights }
87    }
88
89    /// Returns the fields carrying weight.
90    ///
91    /// This allocates. Hot paths iterate [`Self::iter`] instead.
92    #[must_use]
93    pub fn fields(&self) -> Vec<FieldId> {
94        self.weights.iter().map(|(field, _)| *field).collect()
95    }
96
97    pub(crate) fn fields_controlled<E>(
98        &self,
99        work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
100    ) -> Result<Vec<FieldId>, E> {
101        work.check_now()?;
102        let mut fields = Vec::with_capacity(self.weights.len());
103        for (field, _) in &self.weights {
104            work.step()?;
105            fields.push(*field);
106        }
107        work.check_now()?;
108        Ok(fields)
109    }
110
111    pub(crate) fn clone_controlled<E>(
112        &self,
113        work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
114    ) -> Result<Self, E> {
115        work.check_now()?;
116        let mut weights = Vec::with_capacity(self.weights.len());
117        for entry in &self.weights {
118            work.step()?;
119            #[cfg(test)]
120            PREPARED_FIELD_COPIES.with(|count| count.set(count.get() + 1));
121            weights.push(*entry);
122        }
123        work.check_now()?;
124        // The source already has normalized order and first-weight semantics.
125        Ok(Self { weights })
126    }
127
128    /// Returns the `(field, weight)` pairs, ascending by field.
129    ///
130    /// The allocation-free counterpart of [`Self::fields`]. Because the
131    /// table holds one entry per field, `iter` and [`Self::weight`] cannot
132    /// disagree.
133    pub fn iter(&self) -> impl Iterator<Item = (FieldId, u32)> + '_ {
134        self.weights.iter().copied()
135    }
136
137    pub(crate) fn allocation_bytes(&self) -> Option<usize> {
138        self.weights
139            .len()
140            .checked_mul(std::mem::size_of::<(FieldId, u32)>())
141    }
142
143    /// Returns one field's weight in thousandths.
144    #[must_use]
145    pub fn weight(&self, field: FieldId) -> u32 {
146        self.weights
147            .iter()
148            .find(|(candidate, _)| *candidate == field)
149            .map_or(0, |(_, weight)| *weight)
150    }
151
152    /// Returns true when every named field carries unit weight.
153    ///
154    /// Unit weight is 1,000 thousandths, the identity: a field weighted
155    /// this way contributes its length and its term frequencies unchanged.
156    #[must_use]
157    pub fn is_unit(&self) -> bool {
158        !self.weights.is_empty() && self.weights.iter().all(|(_, weight)| *weight == 1_000)
159    }
160
161    /// Returns true when every weight is equal.
162    #[must_use]
163    pub fn is_flat(&self) -> bool {
164        let mut iterator = self.weights.iter().map(|(_, weight)| *weight);
165        let Some(first) = iterator.next() else {
166            return true;
167        };
168        iterator.all(|weight| weight == first)
169    }
170}
171
172/// A structured lexical query.
173///
174/// Terms are already analyzed. There is no string query path into the
175/// engine: parsing text into a structured query is a bindings-layer
176/// convenience with its own typed errors.
177#[derive(Clone, Debug, Eq, PartialEq)]
178pub struct TermQuery {
179    /// Analyzed query terms.
180    pub terms: Vec<Vec<u8>>,
181    /// Field weights to score over.
182    pub fields: FieldWeights,
183}
184
185impl TermQuery {
186    /// Builds a flat single-field query over already-analyzed terms.
187    #[must_use]
188    pub fn flat(terms: Vec<Vec<u8>>, fields: &[FieldId]) -> Self {
189        Self {
190            terms,
191            fields: FieldWeights::flat(fields),
192        }
193    }
194}
195
196/// Deterministic counters, extended by task 14's pruning contracts.
197#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
198pub struct SearchCounters {
199    /// Documents whose score was actually computed.
200    pub docs_evaluated: u64,
201    /// Posting entries decoded.
202    pub postings_decoded: u64,
203    /// Blocks whose contents were decoded.
204    pub blocks_decoded: u64,
205    /// Blocks skipped without decoding. Always zero for the exhaustive path.
206    pub blocks_skipped: u64,
207}
208
209/// A search outcome with its counters.
210#[derive(Clone, Debug, Default, PartialEq)]
211pub struct SearchResult {
212    /// Scored documents, best first.
213    pub hits: Vec<ScoredDoc>,
214    /// Counters for this search.
215    pub counters: SearchCounters,
216}
217
218/// The weighted analyzed length of one row.
219#[must_use]
220pub fn row_length(segment: &SealedSegment, row: u32, weights: &FieldWeights) -> u32 {
221    weighted_length(segment, row, weights)
222}
223
224/// The weighted analyzed length of every row in one segment.
225///
226/// # Why this borrows
227///
228/// The flat single-field scorer asks for `length * 1000 / 1000`, which is
229/// the field's own dense array unchanged. Returning it borrowed makes the
230/// per-query length preparation free rather than `O(row_count)`, and the
231/// identity is exact for every `u32` length, not approximate.
232///
233/// Any other weighting materializes the array once per segment per query,
234/// which is still one allocation instead of one lookup per scored posting.
235#[must_use]
236pub fn weighted_lengths<'segment>(
237    segment: &'segment SealedSegment,
238    weights: &FieldWeights,
239) -> Cow<'segment, [u32]> {
240    let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
241    match weighted_lengths_controlled(segment, weights, &mut work) {
242        Ok(lengths) => lengths,
243        Err(never) => match never {},
244    }
245}
246
247pub(crate) fn weighted_lengths_controlled<'segment, E>(
248    segment: &'segment SealedSegment,
249    weights: &FieldWeights,
250    work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
251) -> Result<Cow<'segment, [u32]>, E> {
252    work.check_now()?;
253    let row_count = usize::try_from(segment.row_count()).unwrap_or(0);
254    let mut entries = weights.iter();
255    if let (Some((field, 1_000)), None) = (entries.next(), entries.next())
256        && let Some(lengths) = segment.field_lengths_controlled(field, work)?
257        && lengths.len() == row_count
258    {
259        return Ok(Cow::Borrowed(lengths));
260    }
261    // The flat MULTI-field case, which is what every BEIR run uses: unit
262    // weight over a set of fields covering the segment. The weighted length
263    // is then just the document's total length, which does not depend on
264    // the query at all and is precomputed once at seal.
265    //
266    // Without this the scorer rebuilt a row-count-long array on EVERY
267    // query -- 171,332 entries on TREC-COVID, each costing a linear scan of
268    // the field table per field.
269    let mut all_unit = !weights.weights.is_empty();
270    for (_, weight) in weights.iter() {
271        work.step()?;
272        if weight != 1_000 {
273            all_unit = false;
274            break;
275        }
276    }
277    if all_unit {
278        for field in segment.fields() {
279            work.step()?;
280            let mut found = false;
281            for (candidate, _) in weights.iter() {
282                work.step()?;
283                if candidate == field {
284                    found = true;
285                    break;
286                }
287            }
288            if !found {
289                all_unit = false;
290                break;
291            }
292        }
293        if all_unit && segment.total_lengths().len() == row_count {
294            return Ok(Cow::Borrowed(segment.total_lengths()));
295        }
296    }
297    let mut lengths = Vec::with_capacity(row_count);
298    for row in 0..row_count {
299        work.step()?;
300        lengths.push(row_length_controlled(
301            segment,
302            u32::try_from(row).unwrap_or(u32::MAX),
303            weights,
304            work,
305        )?);
306    }
307    work.check_now()?;
308    Ok(Cow::Owned(lengths))
309}
310
311/// Scores every matching document exhaustively and returns the top `k`.
312///
313/// # Errors
314///
315/// Returns [`IndexError::Stats`] when the index holds no documents, because
316/// `avgdl` is undefined and every score would be a division artefact.
317pub fn search(
318    index: &LexicalIndex,
319    query: &TermQuery,
320    k: usize,
321    params: Bm25Params,
322) -> Result<SearchResult, IndexError> {
323    let stats = index.corpus_stats()?;
324    let fields = query.fields.fields();
325    let mut counters = SearchCounters::default();
326    let mut accumulator: Vec<(GlobalDocId, f64)> = Vec::new();
327
328    // Document frequency is store-wide, computed once per term.
329    let mut frequencies: Vec<(usize, u32)> = Vec::with_capacity(query.terms.len());
330    for (position, term) in query.terms.iter().enumerate() {
331        frequencies.push((position, index.prepared_document_frequency(term, &fields)?));
332    }
333
334    for (segment_ordinal, segment) in index.segments().iter().enumerate() {
335        let segment_index = u32::try_from(segment_ordinal).unwrap_or(u32::MAX);
336        // Row ids are dense, so a flat array indexed by row is both the
337        // fastest accumulator and the one that keeps per-document summation
338        // in query-term order — which task 14's pruning must reproduce to
339        // the last bit. `touched` keeps the sweep proportional to matches
340        // rather than to corpus size.
341        let row_count = usize::try_from(segment.row_count()).unwrap_or(0);
342        let mut row_scores: Vec<f64> = vec![0.0; row_count];
343        let mut touched: Vec<u32> = Vec::new();
344        // Prepared once per segment, borrowed outright in the flat
345        // single-field case, instead of once per scored posting.
346        let lengths = weighted_lengths(segment, &query.fields);
347
348        for (position, term) in query.terms.iter().enumerate() {
349            let df = frequencies
350                .iter()
351                .find(|(slot, _)| *slot == position)
352                .map_or(0, |(_, value)| *value);
353            if df == 0 {
354                continue;
355            }
356            let Some(mut stream) = TermStream::open(segment, term, &query.fields) else {
357                continue;
358            };
359            // `idf` costs a `ln()` and `avgdl` costs a division; both are
360            // constant across this term's postings, so they are computed
361            // once here rather than once per posting.
362            let scorer = TermScorer::new(Df(df), &stats, params);
363            // Blocks arrive from the sealed stream, decoded on demand into
364            // reused scratch buffers. Nothing is materialized.
365            while let Some(row) = stream.current_row() {
366                counters.postings_decoded = counters.postings_decoded.saturating_add(1);
367                let tf = stream.current_tf().unwrap_or(0);
368                let length = usize::try_from(row)
369                    .ok()
370                    .and_then(|slot| lengths.get(slot).copied())
371                    .unwrap_or(0);
372                let score = scorer.score(Tf(tf), DocLen(length));
373                if let Ok(slot) = usize::try_from(row)
374                    && let Some(entry) = row_scores.get_mut(slot)
375                {
376                    if *entry == 0.0 {
377                        touched.push(row);
378                    }
379                    *entry += score;
380                }
381                stream.advance();
382            }
383            counters.blocks_decoded = counters
384                .blocks_decoded
385                .saturating_add(stream.blocks_decoded());
386        }
387
388        touched.sort_unstable();
389        touched.dedup();
390        for row in touched {
391            let score = usize::try_from(row)
392                .ok()
393                .and_then(|slot| row_scores.get(slot).copied())
394                .unwrap_or(0.0);
395            counters.docs_evaluated = counters.docs_evaluated.saturating_add(1);
396            accumulator.push((
397                GlobalDocId {
398                    segment: segment_index,
399                    row,
400                },
401                score,
402            ));
403        }
404    }
405
406    // Descending score, then ascending document id. Pinned; see module docs.
407    accumulator.sort_by(|left, right| {
408        right
409            .1
410            .partial_cmp(&left.1)
411            .unwrap_or(std::cmp::Ordering::Equal)
412            .then(left.0.cmp(&right.0))
413    });
414    accumulator.truncate(k);
415
416    Ok(SearchResult {
417        hits: accumulator
418            .into_iter()
419            .map(|(doc, score)| ScoredDoc { doc, score })
420            .collect(),
421        counters,
422    })
423}
424
425/// Scores only allow-listed rows by driving iteration from those row ids.
426///
427/// The caller validates one bitmap per lexical segment and bounds every row
428/// before entering this crate-internal exact branch.
429pub(crate) fn search_allow_list_driven(
430    index: &LexicalIndex,
431    query: &TermQuery,
432    k: usize,
433    params: Bm25Params,
434    allow_lists: &[&crate::meta::DocBitmap],
435) -> Result<SearchResult, IndexError> {
436    match search_allow_list_driven_controlled(index, query, k, params, allow_lists, || {
437        Ok::<(), std::convert::Infallible>(())
438    }) {
439        Ok(result) => Ok(result),
440        Err(ControlledSearchError::Index(error)) => Err(error),
441        Err(ControlledSearchError::Control(error)) => match error {},
442    }
443}
444
445pub(crate) enum ControlledSearchError<Control> {
446    Index(IndexError),
447    Control(Control),
448}
449
450impl<Control> From<IndexError> for ControlledSearchError<Control> {
451    fn from(error: IndexError) -> Self {
452        Self::Index(error)
453    }
454}
455
456impl ControlledSearchError<std::convert::Infallible> {
457    pub(crate) fn into_index_error(self) -> IndexError {
458        match self {
459            Self::Index(error) => error,
460            Self::Control(never) => match never {},
461        }
462    }
463}
464
465#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
466pub(crate) struct CandidateScoringWork {
467    pub(crate) streams_opened: u64,
468    pub(crate) scorers_prepared: u64,
469    pub(crate) seeks: u64,
470    pub(crate) rows_visited: u64,
471}
472
473pub(crate) struct CandidateScores {
474    /// One explicit match/nonmatch in each original caller position.
475    pub(crate) scores: Vec<Option<f64>>,
476    pub(crate) counters: SearchCounters,
477}
478
479/// Owned analyzed terms and corpus-dependent constants for one pinned query.
480/// All three term arrays are constructed together and remain immutable.
481pub(crate) struct PreparedTermQuery {
482    query: TermQuery,
483    stats: super::bm25::CorpusStats,
484    frequencies: Vec<u32>,
485    scorers: Vec<TermScorer>,
486    params: Bm25Params,
487}
488
489#[cfg(test)]
490std::thread_local! {
491    static PREPARED_TERM_COPIES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
492    static PREPARED_FIELD_COPIES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
493    static PREPARED_SIZE_VISITS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
494}
495
496impl PreparedTermQuery {
497    pub(crate) fn new_controlled<E>(
498        index: &LexicalIndex,
499        query: &TermQuery,
500        params: Bm25Params,
501        mut check: impl FnMut() -> Result<(), E>,
502    ) -> Result<Self, ControlledSearchError<E>> {
503        let mut work =
504            super::control::WorkCheck::new(|| check().map_err(ControlledSearchError::Control));
505        work.check_now()?;
506        let mut terms = Vec::with_capacity(query.terms.len());
507        for term in &query.terms {
508            work.step()?;
509            #[cfg(test)]
510            PREPARED_TERM_COPIES.with(|count| count.set(count.get() + 1));
511            terms.push(term.clone());
512        }
513        let fields = query.fields.clone_controlled(&mut work)?;
514        Self::from_owned_controlled(index, TermQuery { terms, fields }, params, check)
515    }
516
517    pub(crate) fn from_owned_controlled<E>(
518        index: &LexicalIndex,
519        query: TermQuery,
520        params: Bm25Params,
521        mut check: impl FnMut() -> Result<(), E>,
522    ) -> Result<Self, ControlledSearchError<E>> {
523        check().map_err(ControlledSearchError::Control)?;
524        let stats = index.corpus_stats()?;
525        let fields = query
526            .fields
527            .fields_controlled(&mut super::control::WorkCheck::new(|| {
528                check().map_err(ControlledSearchError::Control)
529            }))?;
530        let frequencies = query
531            .terms
532            .iter()
533            .map(|term| index.prepared_document_frequency_controlled(term, &fields, &mut check))
534            .collect::<Result<Vec<_>, _>>()?;
535        let mut scorers = Vec::with_capacity(frequencies.len());
536        let mut work =
537            super::control::WorkCheck::new(|| check().map_err(ControlledSearchError::Control));
538        for df in &frequencies {
539            work.step()?;
540            scorers.push(TermScorer::new(Df(*df), &stats, params));
541        }
542        work.check_now()?;
543        Ok(Self {
544            query,
545            stats,
546            frequencies,
547            scorers,
548            params,
549        })
550    }
551
552    #[cfg(test)]
553    pub(crate) fn new(
554        index: &LexicalIndex,
555        query: &TermQuery,
556        params: Bm25Params,
557    ) -> Result<Self, IndexError> {
558        Self::from_owned(index, query.clone(), params)
559    }
560
561    #[cfg(test)]
562    pub(crate) fn from_owned(
563        index: &LexicalIndex,
564        query: TermQuery,
565        params: Bm25Params,
566    ) -> Result<Self, IndexError> {
567        let stats = index.corpus_stats()?;
568        let fields = query.fields.fields();
569        let frequencies: Vec<_> = query
570            .terms
571            .iter()
572            .map(|term| index.prepared_document_frequency(term, &fields))
573            .collect::<Result<Vec<_>, _>>()?;
574        let scorers = frequencies
575            .iter()
576            .map(|df| TermScorer::new(Df(*df), &stats, params))
577            .collect();
578        Ok(Self {
579            query,
580            stats,
581            frequencies,
582            scorers,
583            params,
584        })
585    }
586
587    pub(crate) fn query(&self) -> &TermQuery {
588        &self.query
589    }
590    pub(crate) fn stats(&self) -> super::bm25::CorpusStats {
591        self.stats
592    }
593    pub(crate) fn frequencies(&self) -> &[u32] {
594        &self.frequencies
595    }
596    pub(crate) fn scorer(&self, slot: usize) -> Option<TermScorer> {
597        self.scorers.get(slot).copied()
598    }
599    pub(crate) fn params(&self) -> Bm25Params {
600        self.params
601    }
602
603    #[cfg(test)]
604    pub(crate) fn allocation_bytes(query: &TermQuery) -> Option<usize> {
605        let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
606        match Self::allocation_bytes_controlled(query, &mut work) {
607            Ok(bytes) => bytes,
608            Err(never) => match never {},
609        }
610    }
611
612    pub(crate) fn allocation_bytes_controlled<E>(
613        query: &TermQuery,
614        work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
615    ) -> Result<Option<usize>, E> {
616        work.check_now()?;
617        let mut term_bytes = 0_usize;
618        for term in &query.terms {
619            work.step()?;
620            #[cfg(test)]
621            PREPARED_SIZE_VISITS.with(|count| count.set(count.get() + 1));
622            let Some(sum) = term_bytes.checked_add(term.len()) else {
623                return Ok(None);
624            };
625            term_bytes = sum;
626        }
627        work.check_now()?;
628        Ok(query
629            .terms
630            .len()
631            .checked_mul(
632                std::mem::size_of::<Vec<u8>>()
633                    + std::mem::size_of::<u32>()
634                    + std::mem::size_of::<TermScorer>(),
635            )
636            .and_then(|bytes| bytes.checked_add(term_bytes))
637            .and_then(|bytes| {
638                bytes.checked_add(
639                    query
640                        .fields
641                        .weights
642                        .len()
643                        .checked_mul(std::mem::size_of::<(FieldId, u32)>())?,
644                )
645            }))
646    }
647
648    pub(crate) fn single_allocation_bytes(
649        fields: &FieldWeights,
650        term_bytes: usize,
651    ) -> Option<usize> {
652        (std::mem::size_of::<Vec<u8>>()
653            + std::mem::size_of::<u32>()
654            + std::mem::size_of::<TermScorer>())
655        .checked_add(term_bytes)?
656        .checked_add(
657            fields
658                .weights
659                .len()
660                .checked_mul(std::mem::size_of::<(FieldId, u32)>())?,
661        )
662    }
663}
664
665/// Frozen query statistics shared by candidate batches across segments.
666/// Eligibility is a caller contract; row bounds are checked before scoring.
667pub(crate) struct CandidateScoring<'query> {
668    query: &'query TermQuery,
669    stats: super::bm25::CorpusStats,
670    frequencies: Cow<'query, [u32]>,
671    params: Bm25Params,
672    scorers: Option<&'query [TermScorer]>,
673}
674
675impl<'query> CandidateScoring<'query> {
676    pub(crate) fn new_controlled<E>(
677        index: &LexicalIndex,
678        query: &'query TermQuery,
679        params: Bm25Params,
680        mut check: impl FnMut() -> Result<(), E>,
681    ) -> Result<Self, ControlledSearchError<E>> {
682        check().map_err(ControlledSearchError::Control)?;
683        let stats = index.corpus_stats()?;
684        let fields = query
685            .fields
686            .fields_controlled(&mut super::control::WorkCheck::new(|| {
687                check().map_err(ControlledSearchError::Control)
688            }))?;
689        let frequencies = query
690            .terms
691            .iter()
692            .map(|term| index.prepared_document_frequency_controlled(term, &fields, &mut check))
693            .collect::<Result<Vec<_>, _>>()?;
694        check().map_err(ControlledSearchError::Control)?;
695        Ok(Self {
696            query,
697            stats,
698            frequencies: Cow::Owned(frequencies),
699            params,
700            scorers: None,
701        })
702    }
703
704    #[cfg(test)]
705    pub(crate) fn new(
706        index: &LexicalIndex,
707        query: &'query TermQuery,
708        params: Bm25Params,
709    ) -> Result<Self, IndexError> {
710        let stats = index.corpus_stats()?;
711        let fields = query.fields.fields();
712        let frequencies = query
713            .terms
714            .iter()
715            .map(|term| index.prepared_document_frequency(term, &fields))
716            .collect::<Result<Vec<_>, _>>()?;
717        Ok(Self {
718            query,
719            stats,
720            frequencies: Cow::Owned(frequencies),
721            params,
722            scorers: None,
723        })
724    }
725
726    pub(crate) fn prepared(prepared: &'query PreparedTermQuery) -> Self {
727        Self {
728            query: &prepared.query,
729            stats: prepared.stats,
730            frequencies: Cow::Borrowed(&prepared.frequencies),
731            params: prepared.params,
732            scorers: Some(&prepared.scorers),
733        }
734    }
735
736    fn term_scorer(&self, slot: usize, df: u32) -> TermScorer {
737        self.scorers
738            .and_then(|scorers| scorers.get(slot))
739            .copied()
740            .unwrap_or_else(|| TermScorer::new(Df(df), &self.stats, self.params))
741    }
742
743    pub(crate) fn score_rows<Control>(
744        &self,
745        ordinal: usize,
746        segment: &SealedSegment,
747        rows: &[u32],
748        work: &mut CandidateScoringWork,
749        checkpoint: impl FnMut(&CandidateScoringWork) -> Result<(), Control>,
750    ) -> Result<CandidateScores, ControlledSearchError<Control>> {
751        self.score_rows_with(ordinal, segment, rows, work, checkpoint, |_, score| score)
752    }
753
754    pub(crate) fn score_weighted_rows<Control>(
755        &self,
756        ordinal: usize,
757        segment: &SealedSegment,
758        rows: &[u32],
759        work: &mut CandidateScoringWork,
760        checkpoint: impl FnMut(&CandidateScoringWork) -> Result<(), Control>,
761        expansions: &[super::query::LexicalExpansion],
762    ) -> Result<CandidateScores, ControlledSearchError<Control>> {
763        self.score_rows_with(ordinal, segment, rows, work, checkpoint, |slot, score| {
764            score
765                * expansions
766                    .get(slot)
767                    .map_or(0.0, |entry| f64::from(entry.boost_thousandths) / 1_000.0)
768        })
769    }
770
771    #[allow(clippy::too_many_arguments)]
772    fn score_rows_with<Control>(
773        &self,
774        ordinal: usize,
775        segment: &SealedSegment,
776        rows: &[u32],
777        work: &mut CandidateScoringWork,
778        mut checkpoint: impl FnMut(&CandidateScoringWork) -> Result<(), Control>,
779        contribution: impl Fn(usize, f64) -> f64,
780    ) -> Result<CandidateScores, ControlledSearchError<Control>> {
781        checkpoint(work).map_err(ControlledSearchError::Control)?;
782        // Validate the whole input before any cursor or score is produced.
783        for (position, &row) in rows.iter().enumerate() {
784            if row >= segment.row_count() {
785                return Err(ControlledSearchError::Index(
786                    IndexError::LiveRowOutOfRange {
787                        segment: ordinal,
788                        row,
789                        row_count: segment.row_count(),
790                    },
791                ));
792            }
793            if position % 64 == 63 {
794                checkpoint(work).map_err(ControlledSearchError::Control)?;
795            }
796        }
797        let lengths = {
798            let mut nested = super::control::WorkCheck::new(|| {
799                checkpoint(work).map_err(ControlledSearchError::Control)
800            });
801            weighted_lengths_controlled(segment, &self.query.fields, &mut nested)?
802        };
803        let routing = {
804            let mut nested = super::control::WorkCheck::new(|| {
805                checkpoint(work).map_err(ControlledSearchError::Control)
806            });
807            let mut routing = Vec::with_capacity(rows.len());
808            for (position, &row) in rows.iter().enumerate() {
809                nested.step()?;
810                routing.push((row, position));
811            }
812            super::control::sort_by(
813                &mut routing,
814                |left, right| {
815                    #[cfg(test)]
816                    CANDIDATE_ROUTING_COMPARISONS.with(|value| value.set(value.get() + 1));
817                    left.0.cmp(&right.0)
818                },
819                &mut nested,
820            )?;
821            routing
822        };
823        let mut scores = vec![None; rows.len()];
824        let mut counters = SearchCounters::default();
825        if rows.is_empty() {
826            return Ok(CandidateScores { scores, counters });
827        }
828        let mut streams = Vec::with_capacity(self.query.terms.len());
829        // Preserve query-term occurrences and order, including repeated terms.
830        for (slot, (term, &df)) in self
831            .query
832            .terms
833            .iter()
834            .zip(self.frequencies.iter())
835            .enumerate()
836        {
837            if df != 0 {
838                let mut nested = super::control::WorkCheck::new(|| {
839                    checkpoint(work).map_err(ControlledSearchError::Control)
840                });
841                if let Some(stream) =
842                    TermStream::open_controlled(segment, term, &self.query.fields, &mut nested)?
843                {
844                    work.streams_opened = work.streams_opened.saturating_add(1);
845                    let scorer = self.term_scorer(slot, df);
846                    work.scorers_prepared = work
847                        .scorers_prepared
848                        .saturating_add(u64::from(self.scorers.is_none()));
849                    streams.push((stream, scorer, slot));
850                }
851            }
852            if slot % 64 == 63 {
853                checkpoint(work).map_err(ControlledSearchError::Control)?;
854            }
855        }
856        let mut group_start = 0;
857        while let Some(&(row, _)) = routing.get(group_start) {
858            let mut group_end = group_start + 1;
859            {
860                let mut nested = super::control::WorkCheck::new(|| {
861                    checkpoint(work).map_err(ControlledSearchError::Control)
862                });
863                while routing.get(group_end).is_some_and(|&(next, _)| next == row) {
864                    nested.step()?;
865                    group_end += 1;
866                }
867            }
868            let length = usize::try_from(row)
869                .ok()
870                .and_then(|slot| lengths.get(slot).copied())
871                .ok_or(ControlledSearchError::Index(
872                    IndexError::LiveLengthMissing {
873                        segment: ordinal,
874                        row,
875                    },
876                ))?;
877            let mut total = 0.0_f64;
878            let mut matched = false;
879            for (stream, scorer, slot) in &mut streams {
880                {
881                    let mut nested = super::control::WorkCheck::new(|| {
882                        checkpoint(work).map_err(ControlledSearchError::Control)
883                    });
884                    stream.seek_controlled(row, &mut nested)?;
885                }
886                work.seeks = work.seeks.saturating_add(1);
887                if work.seeks.is_multiple_of(64) {
888                    checkpoint(work).map_err(ControlledSearchError::Control)?;
889                }
890                if stream.current_row() == Some(row) {
891                    counters.postings_decoded = counters.postings_decoded.saturating_add(1);
892                    total += contribution(
893                        *slot,
894                        scorer.score(
895                            Tf({
896                                let mut nested = super::control::WorkCheck::new(|| {
897                                    checkpoint(work).map_err(ControlledSearchError::Control)
898                                });
899                                stream.current_tf_controlled(&mut nested)?.unwrap_or(0)
900                            }),
901                            DocLen(length),
902                        ),
903                    );
904                    matched = true;
905                }
906            }
907            work.rows_visited = work.rows_visited.saturating_add(1);
908            if work.rows_visited.is_multiple_of(64) {
909                checkpoint(work).map_err(ControlledSearchError::Control)?;
910            }
911            if matched {
912                counters.docs_evaluated = counters.docs_evaluated.saturating_add(1);
913            }
914            let mut nested = super::control::WorkCheck::new(|| {
915                checkpoint(work).map_err(ControlledSearchError::Control)
916            });
917            // These private boundaries come only from successful routing.get
918            // probes above, so both splits remain inside the owned slice.
919            let group = routing
920                .split_at(group_start)
921                .1
922                .split_at(group_end - group_start)
923                .0;
924            for &(_, position) in group {
925                nested.step()?;
926                #[cfg(test)]
927                CANDIDATE_DUPLICATE_COPIES.with(|value| value.set(value.get() + 1));
928                // Every position was generated from this exact output length.
929                if let Some(output) = scores.get_mut(position) {
930                    *output = matched.then_some(total);
931                }
932            }
933            group_start = group_end;
934        }
935        // Reused cursors own cumulative block receipts; count each once.
936        for (slot, (stream, _, _)) in streams.into_iter().enumerate() {
937            if slot % 64 == 0 {
938                checkpoint(work).map_err(ControlledSearchError::Control)?;
939            }
940            counters.blocks_decoded = counters
941                .blocks_decoded
942                .saturating_add(stream.blocks_decoded());
943            counters.blocks_skipped = counters
944                .blocks_skipped
945                .saturating_add(stream.blocks_skipped());
946        }
947        checkpoint(work).map_err(ControlledSearchError::Control)?;
948        Ok(CandidateScores { scores, counters })
949    }
950}
951
952#[derive(Clone, Copy, Debug, Eq, PartialEq)]
953enum AllowListScoringBranch {
954    RowDriven,
955    PostingDriven,
956}
957
958fn allow_list_scoring_branch(allowed: u64, corpus: u64) -> AllowListScoringBranch {
959    if allowed.saturating_mul(crate::planner::LEXICAL_ALLOW_LIST_DIVISOR) <= corpus {
960        AllowListScoringBranch::RowDriven
961    } else {
962        AllowListScoringBranch::PostingDriven
963    }
964}
965
966/// Allow-list scorer with deterministic cancellation checkpoints at entry,
967/// exit, segment boundaries, and every 64 posting/document work units.
968pub(crate) fn search_allow_list_driven_controlled<Control>(
969    index: &LexicalIndex,
970    query: &TermQuery,
971    k: usize,
972    params: Bm25Params,
973    allow_lists: &[&crate::meta::DocBitmap],
974    checkpoint: impl FnMut() -> Result<(), Control>,
975) -> Result<SearchResult, ControlledSearchError<Control>> {
976    let allowed = allow_lists
977        .iter()
978        .map(|allow_list| allow_list.cardinality())
979        .fold(0_u64, u64::saturating_add);
980    let branch = allow_list_scoring_branch(allowed, index.document_count());
981    search_allow_list_driven_controlled_with_branch(
982        index,
983        query,
984        k,
985        params,
986        allow_lists,
987        branch,
988        checkpoint,
989        None,
990    )
991}
992
993pub(crate) fn search_allow_list_prepared_controlled<Control>(
994    index: &LexicalIndex,
995    prepared: &PreparedTermQuery,
996    k: usize,
997    allow_lists: &[&crate::meta::DocBitmap],
998    checkpoint: impl FnMut() -> Result<(), Control>,
999) -> Result<SearchResult, ControlledSearchError<Control>> {
1000    let allowed = allow_lists
1001        .iter()
1002        .map(|list| list.cardinality())
1003        .fold(0_u64, u64::saturating_add);
1004    let branch = allow_list_scoring_branch(allowed, index.document_count());
1005    search_allow_list_driven_controlled_with_branch(
1006        index,
1007        prepared.query(),
1008        k,
1009        prepared.params,
1010        allow_lists,
1011        branch,
1012        checkpoint,
1013        Some(prepared),
1014    )
1015}
1016
1017#[allow(clippy::too_many_arguments)]
1018fn search_allow_list_driven_controlled_with_branch<Control>(
1019    index: &LexicalIndex,
1020    query: &TermQuery,
1021    k: usize,
1022    params: Bm25Params,
1023    allow_lists: &[&crate::meta::DocBitmap],
1024    branch: AllowListScoringBranch,
1025    mut checkpoint: impl FnMut() -> Result<(), Control>,
1026    prepared_query: Option<&PreparedTermQuery>,
1027) -> Result<SearchResult, ControlledSearchError<Control>> {
1028    checkpoint().map_err(ControlledSearchError::Control)?;
1029    let prepared = match prepared_query {
1030        Some(prepared) => CandidateScoring::prepared(prepared),
1031        None => CandidateScoring::new_controlled(index, query, params, &mut checkpoint)?,
1032    };
1033    let frequencies = &prepared.frequencies;
1034    let mut counters = SearchCounters::default();
1035    let mut accumulator = Vec::<(GlobalDocId, f64)>::new();
1036    let mut work_units = 0_u64;
1037
1038    match branch {
1039        AllowListScoringBranch::RowDriven => {
1040            for (ordinal, segment) in index.segments().iter().enumerate() {
1041                checkpoint().map_err(ControlledSearchError::Control)?;
1042                let Some(allow_list) = allow_lists.get(ordinal) else {
1043                    continue;
1044                };
1045                let segment_index = u32::try_from(ordinal).unwrap_or(u32::MAX);
1046                let rows = allow_list.iter().collect::<Vec<_>>();
1047                let batch = prepared.score_rows(
1048                    ordinal,
1049                    segment,
1050                    &rows,
1051                    &mut CandidateScoringWork::default(),
1052                    |_| checkpoint(),
1053                )?;
1054                counters.docs_evaluated = counters
1055                    .docs_evaluated
1056                    .saturating_add(batch.counters.docs_evaluated);
1057                counters.postings_decoded = counters
1058                    .postings_decoded
1059                    .saturating_add(batch.counters.postings_decoded);
1060                counters.blocks_decoded = counters
1061                    .blocks_decoded
1062                    .saturating_add(batch.counters.blocks_decoded);
1063                counters.blocks_skipped = counters
1064                    .blocks_skipped
1065                    .saturating_add(batch.counters.blocks_skipped);
1066                for (row, score) in rows.into_iter().zip(batch.scores) {
1067                    if let Some(score) = score {
1068                        accumulator.push((
1069                            GlobalDocId {
1070                                segment: segment_index,
1071                                row,
1072                            },
1073                            score,
1074                        ));
1075                    }
1076                }
1077            }
1078        }
1079        AllowListScoringBranch::PostingDriven => {
1080            let largest_segment = index
1081                .segments()
1082                .iter()
1083                .filter_map(|segment| usize::try_from(segment.row_count()).ok())
1084                .max()
1085                .unwrap_or(0);
1086            let mut row_scores = vec![0.0_f64; largest_segment];
1087            let mut touched = Vec::<u32>::new();
1088
1089            for (ordinal, segment) in index.segments().iter().enumerate() {
1090                checkpoint().map_err(ControlledSearchError::Control)?;
1091                let Some(allow_list) = allow_lists.get(ordinal) else {
1092                    continue;
1093                };
1094                let segment_index = u32::try_from(ordinal).unwrap_or(u32::MAX);
1095                let lengths = {
1096                    let mut nested = super::control::WorkCheck::new(|| {
1097                        checkpoint().map_err(ControlledSearchError::Control)
1098                    });
1099                    weighted_lengths_controlled(segment, &query.fields, &mut nested)?
1100                };
1101
1102                for (slot, term) in query.terms.iter().enumerate() {
1103                    let df = frequencies.get(slot).copied().unwrap_or(0);
1104                    if df == 0 {
1105                        continue;
1106                    }
1107                    let mut nested = super::control::WorkCheck::new(|| {
1108                        checkpoint().map_err(ControlledSearchError::Control)
1109                    });
1110                    let Some(mut stream) =
1111                        TermStream::open_controlled(segment, term, &query.fields, &mut nested)?
1112                    else {
1113                        continue;
1114                    };
1115                    let scorer = prepared.term_scorer(slot, df);
1116                    while let Some(row) = stream.current_row() {
1117                        work_units = work_units.saturating_add(1);
1118                        if work_units.is_multiple_of(64) {
1119                            nested.check_now()?;
1120                        }
1121                        counters.postings_decoded = counters.postings_decoded.saturating_add(1);
1122                        if allow_list.contains(row) {
1123                            let tf = stream.current_tf_controlled(&mut nested)?.unwrap_or(0);
1124                            let length = usize::try_from(row)
1125                                .ok()
1126                                .and_then(|position| lengths.get(position).copied())
1127                                .unwrap_or(0);
1128                            let score = scorer.score(Tf(tf), DocLen(length));
1129                            if let Ok(position) = usize::try_from(row)
1130                                && let Some(entry) = row_scores.get_mut(position)
1131                            {
1132                                if *entry == 0.0 {
1133                                    touched.push(row);
1134                                }
1135                                *entry += score;
1136                            }
1137                        }
1138                        stream.advance_controlled(&mut nested)?;
1139                    }
1140                    counters.blocks_decoded = counters
1141                        .blocks_decoded
1142                        .saturating_add(stream.blocks_decoded());
1143                    counters.blocks_skipped = counters
1144                        .blocks_skipped
1145                        .saturating_add(stream.blocks_skipped());
1146                }
1147
1148                touched.sort_unstable();
1149                touched.dedup();
1150                for row in touched.drain(..) {
1151                    let score = usize::try_from(row)
1152                        .ok()
1153                        .and_then(|position| row_scores.get_mut(position))
1154                        .map(|entry| {
1155                            let score = *entry;
1156                            *entry = 0.0;
1157                            score
1158                        })
1159                        .unwrap_or(0.0);
1160                    counters.docs_evaluated = counters.docs_evaluated.saturating_add(1);
1161                    accumulator.push((
1162                        GlobalDocId {
1163                            segment: segment_index,
1164                            row,
1165                        },
1166                        score,
1167                    ));
1168                }
1169            }
1170        }
1171    }
1172    accumulator.sort_by(|left, right| {
1173        right
1174            .1
1175            .partial_cmp(&left.1)
1176            .unwrap_or(std::cmp::Ordering::Equal)
1177            .then(left.0.cmp(&right.0))
1178    });
1179    accumulator.truncate(k);
1180    checkpoint().map_err(ControlledSearchError::Control)?;
1181    Ok(SearchResult {
1182        hits: accumulator
1183            .into_iter()
1184            .map(|(doc, score)| ScoredDoc { doc, score })
1185            .collect(),
1186        counters,
1187    })
1188}
1189
1190#[cfg(test)]
1191thread_local! {
1192    static WEIGHTED_LENGTH_ROWS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1193    static CANDIDATE_ROUTING_COMPARISONS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1194    static CANDIDATE_DUPLICATE_COPIES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1195}
1196
1197fn weighted_length(segment: &SealedSegment, row: u32, weights: &FieldWeights) -> u32 {
1198    let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
1199    match row_length_controlled(segment, row, weights, &mut work) {
1200        Ok(length) => length,
1201        Err(never) => match never {},
1202    }
1203}
1204
1205pub(crate) fn row_length_controlled<E>(
1206    segment: &SealedSegment,
1207    row: u32,
1208    weights: &FieldWeights,
1209    work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1210) -> Result<u32, E> {
1211    #[cfg(test)]
1212    WEIGHTED_LENGTH_ROWS.with(|count| count.set(count.get() + 1));
1213    let mut total = 0_u64;
1214    for (field, weight) in weights.iter() {
1215        work.step()?;
1216        let length = segment
1217            .field_lengths_controlled(field, work)?
1218            .and_then(|lengths| usize::try_from(row).ok().and_then(|slot| lengths.get(slot)))
1219            .copied()
1220            .unwrap_or(0);
1221        total = total.saturating_add(u64::from(length) * u64::from(weight));
1222    }
1223    Ok(u32::try_from(total / 1_000).unwrap_or(u32::MAX))
1224}
1225
1226#[cfg(test)]
1227#[allow(
1228    clippy::expect_used,
1229    clippy::indexing_slicing,
1230    clippy::panic,
1231    clippy::unwrap_used
1232)]
1233mod tests {
1234    use std::convert::Infallible;
1235
1236    use proptest::prelude::*;
1237    use proptest::test_runner::{Config, RngSeed, TestRunner};
1238
1239    use super::*;
1240
1241    fn prepared_copy_case(fields: bool) {
1242        let index = index_of(&["alpha"]);
1243        let terms = vec![b"alpha".to_vec(); if fields { 1 } else { 4_096 }];
1244        let query = TermQuery {
1245            terms,
1246            fields: FieldWeights::flat(
1247                &(0..if fields { 4_096 } else { 1 })
1248                    .map(FieldId)
1249                    .collect::<Vec<_>>(),
1250            ),
1251        };
1252        PREPARED_TERM_COPIES.with(|count| count.set(0));
1253        PREPARED_FIELD_COPIES.with(|count| count.set(0));
1254        let copied = || {
1255            if fields {
1256                PREPARED_FIELD_COPIES.with(std::cell::Cell::get)
1257            } else {
1258                PREPARED_TERM_COPIES.with(std::cell::Cell::get)
1259            }
1260        };
1261        let result =
1262            PreparedTermQuery::new_controlled(&index, &query, Bm25Params::default(), || {
1263                if copied() >= 64 {
1264                    Err("cancelled")
1265                } else {
1266                    Ok(())
1267                }
1268            });
1269        println!("prepared copy fields={fields}, copied={}", copied());
1270        assert!(matches!(
1271            result,
1272            Err(ControlledSearchError::Control("cancelled"))
1273        ));
1274        assert!(
1275            (64..=128).contains(&copied()),
1276            "owned copying must check within the query"
1277        );
1278        let clean =
1279            PreparedTermQuery::new_controlled(&index, &query, Bm25Params::default(), || {
1280                Ok::<(), Infallible>(())
1281            })
1282            .unwrap_or_else(|_| panic!("fresh control failed"));
1283        assert_eq!(clean.query().terms, query.terms);
1284        assert_eq!(clean.query().fields, query.fields);
1285    }
1286
1287    #[test]
1288    fn astra_18_prepared_term_copy_cancels_before_statistics() {
1289        prepared_copy_case(false);
1290    }
1291
1292    #[test]
1293    fn astra_18_prepared_field_copy_cancels_before_statistics() {
1294        prepared_copy_case(true);
1295    }
1296
1297    #[test]
1298    fn astra_18_prepared_term_size_cancels_inside_walk() {
1299        let query = TermQuery::flat(vec![b"alpha".to_vec(); 4_096], &[DEFAULT_FIELD]);
1300        PREPARED_SIZE_VISITS.with(|count| count.set(0));
1301        let mut work = super::super::control::WorkCheck::new(|| {
1302            if PREPARED_SIZE_VISITS.with(std::cell::Cell::get) >= 64 {
1303                Err("cancelled")
1304            } else {
1305                Ok(())
1306            }
1307        });
1308        let result = PreparedTermQuery::allocation_bytes_controlled(&query, &mut work);
1309        let visited = PREPARED_SIZE_VISITS.with(std::cell::Cell::get);
1310        println!("prepared size terms visited={visited}");
1311        assert_eq!(result, Err("cancelled"));
1312        assert!((64..=128).contains(&visited));
1313        let mut clean = super::super::control::WorkCheck::new(|| Ok::<(), Infallible>(()));
1314        // Fixed 64-bit fixture: 4096 owned five-byte terms, their Vec headers,
1315        // u32 frequencies and 32-byte scorers, plus one eight-byte field entry.
1316        assert_eq!(
1317            PreparedTermQuery::allocation_bytes_controlled(&query, &mut clean),
1318            Ok(Some(266_248))
1319        );
1320    }
1321    use crate::fts::index::{DEFAULT_FIELD, Document, SegmentIndex};
1322    use crate::fts::tokenizer::{Analyzer, Profile};
1323    use crate::meta::DocBitmap;
1324
1325    #[test]
1326    fn astra_18_weighted_length_preparation_is_cancelable() {
1327        let index = index_of(&vec!["alpha beta"; 4_096]);
1328        let segment = index.segments().first().unwrap();
1329        let weights = FieldWeights::new(&[(DEFAULT_FIELD, 500)]);
1330        let expected = vec![1_u32; 4_096];
1331        assert_eq!(weighted_lengths(segment, &weights).as_ref(), expected);
1332        WEIGHTED_LENGTH_ROWS.with(|count| count.set(0));
1333        let mut checks = 0;
1334        let mut work = super::super::control::WorkCheck::new(|| {
1335            checks += 1;
1336            if checks == 3 {
1337                Err("cancelled")
1338            } else {
1339                Ok(())
1340            }
1341        });
1342        let result = weighted_lengths_controlled(segment, &weights, &mut work);
1343        let rows = WEIGHTED_LENGTH_ROWS.with(std::cell::Cell::get);
1344        println!("checks={checks}, length_rows={rows}");
1345        assert_eq!(result, Err("cancelled"));
1346        assert!(
1347            rows > 0 && rows <= 128,
1348            "preparation must stop inside the row loop"
1349        );
1350        let mut work = super::super::control::WorkCheck::new(|| Ok::<(), ()>(()));
1351        assert_eq!(
1352            weighted_lengths_controlled(segment, &weights, &mut work)
1353                .unwrap()
1354                .as_ref(),
1355            expected
1356        );
1357    }
1358
1359    fn analyzer() -> Analyzer {
1360        Analyzer::new(Profile::Code.config()).expect("valid config")
1361    }
1362
1363    fn index_of(texts: &[&str]) -> LexicalIndex {
1364        let analyzer = analyzer();
1365        let mut segment = SegmentIndex::new();
1366        for text in texts {
1367            segment
1368                .push_document(&analyzer, &Document::with_text(text))
1369                .expect("indexable");
1370        }
1371        let mut index = LexicalIndex::new();
1372        index.push_segment(segment).expect("seals");
1373        index
1374    }
1375
1376    fn candidate_batch(
1377        index: &LexicalIndex,
1378        query: &TermQuery,
1379        ordinal: usize,
1380        rows: &[u32],
1381    ) -> (CandidateScores, CandidateScoringWork) {
1382        let prepared =
1383            CandidateScoring::new(index, query, Bm25Params::default()).expect("frozen statistics");
1384        let mut work = CandidateScoringWork::default();
1385        let batch =
1386            match prepared.score_rows(ordinal, &index.segments()[ordinal], rows, &mut work, |_| {
1387                Ok::<(), Infallible>(())
1388            }) {
1389                Ok(batch) => batch,
1390                Err(ControlledSearchError::Index(error)) => panic!("candidate scoring: {error}"),
1391                Err(ControlledSearchError::Control(never)) => match never {},
1392            };
1393        (batch, work)
1394    }
1395
1396    #[test]
1397    fn astra_18_candidate_routing_cancels_inside_sort() {
1398        let index = index_of(&["alpha", "omega"]);
1399        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
1400        let prepared =
1401            CandidateScoring::new(&index, &query, Bm25Params::default()).expect("statistics");
1402        let rows = (0..4_096)
1403            .map(|i| ((i * 7_919) % 2) as u32)
1404            .collect::<Vec<_>>();
1405        CANDIDATE_ROUTING_COMPARISONS.with(|value| value.set(0));
1406        let mut work = CandidateScoringWork::default();
1407        let result = prepared.score_rows(0, &index.segments()[0], &rows, &mut work, |_| {
1408            if CANDIDATE_ROUTING_COMPARISONS.with(std::cell::Cell::get) >= 64 {
1409                Err("cancelled")
1410            } else {
1411                Ok(())
1412            }
1413        });
1414        let comparisons = CANDIDATE_ROUTING_COMPARISONS.with(std::cell::Cell::get);
1415        println!(
1416            "candidate routing comparisons={comparisons}, streams={}",
1417            work.streams_opened
1418        );
1419        assert!(matches!(
1420            result,
1421            Err(ControlledSearchError::Control("cancelled"))
1422        ));
1423        assert!(
1424            (64..=128).contains(&comparisons),
1425            "routing sort must stop within its comparisons"
1426        );
1427        assert_eq!(work.streams_opened, 0);
1428        let (clean, _) = candidate_batch(&index, &query, 0, &rows);
1429        assert_eq!(clean.scores.len(), rows.len());
1430        for (&row, score) in rows.iter().zip(clean.scores) {
1431            if row == 0 {
1432                assert!((score.expect("matching literal row") - 2.0_f64.ln()).abs() < 1e-12);
1433            } else {
1434                assert_eq!(score, None);
1435            }
1436        }
1437    }
1438
1439    #[test]
1440    fn astra_18_candidate_duplicates_cancel_inside_output_walk() {
1441        let index = index_of(&["alpha", "omega"]);
1442        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
1443        let prepared =
1444            CandidateScoring::new(&index, &query, Bm25Params::default()).expect("statistics");
1445        let rows = vec![0; 4_096];
1446        CANDIDATE_DUPLICATE_COPIES.with(|value| value.set(0));
1447        let mut work = CandidateScoringWork::default();
1448        let result = prepared.score_rows(0, &index.segments()[0], &rows, &mut work, |_| {
1449            if CANDIDATE_DUPLICATE_COPIES.with(std::cell::Cell::get) >= 64 {
1450                Err("cancelled")
1451            } else {
1452                Ok(())
1453            }
1454        });
1455        let copies = CANDIDATE_DUPLICATE_COPIES.with(std::cell::Cell::get);
1456        println!(
1457            "candidate duplicate copies={copies}, rows_scored={}",
1458            work.rows_visited
1459        );
1460        assert!(matches!(
1461            result,
1462            Err(ControlledSearchError::Control("cancelled"))
1463        ));
1464        assert!(
1465            (64..=128).contains(&copies),
1466            "one repeated row must not hide an unbounded output walk"
1467        );
1468        assert_eq!(work.rows_visited, 1);
1469        let (clean, clean_work) = candidate_batch(&index, &query, 0, &rows);
1470        assert_eq!(clean.scores.len(), rows.len());
1471        assert_eq!(clean_work.rows_visited, 1);
1472        assert!(
1473            clean
1474                .scores
1475                .iter()
1476                .all(|score| score.is_some_and(|score| (score - 2.0_f64.ln()).abs() < 1e-12))
1477        );
1478    }
1479
1480    #[test]
1481    fn astra_02_candidate_bm25_opens_one_stream_per_term_segment() {
1482        let index = index_of(&vec!["alpha beta"; 1_000]);
1483        let query = TermQuery::flat(vec![b"alpha".to_vec(), b"beta".to_vec()], &[DEFAULT_FIELD]);
1484        let (batch, work) = candidate_batch(&index, &query, 0, &[999, 7, 511, 42]);
1485        assert!(batch.scores.iter().all(Option::is_some));
1486        assert_eq!(
1487            work.streams_opened, 2,
1488            "prepare once per term, not once per row"
1489        );
1490        assert_eq!(work.scorers_prepared, 2);
1491        assert_eq!(work.seeks, 8);
1492    }
1493
1494    #[test]
1495    fn astra_02_candidate_nonmatch_is_explicit() {
1496        let index = index_of(&["alpha", "omega"]);
1497        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
1498        let (batch, _) = candidate_batch(&index, &query, 0, &[1, 0, 1]);
1499        assert_eq!(batch.scores.len(), 3);
1500        assert_eq!(batch.scores[0], None);
1501        assert!(batch.scores[1].is_some());
1502        assert_eq!(batch.scores[2], None);
1503        let prepared =
1504            CandidateScoring::new(&index, &query, Bm25Params::default()).expect("prepare");
1505        let mut work = CandidateScoringWork::default();
1506        let result = prepared.score_rows(7, &index.segments()[0], &[0, 2], &mut work, |_| {
1507            Ok::<(), Infallible>(())
1508        });
1509        assert!(matches!(
1510            result,
1511            Err(ControlledSearchError::Index(
1512                IndexError::LiveRowOutOfRange {
1513                    segment: 7,
1514                    row: 2,
1515                    row_count: 2
1516                }
1517            ))
1518        ));
1519        assert_eq!(
1520            work.streams_opened, 0,
1521            "reject the entire request before scoring a valid prefix"
1522        );
1523    }
1524
1525    #[test]
1526    fn astra_02_candidate_scoring_cancels_during_sparse_seeks() {
1527        let index = index_of(&["alpha", "omega"]);
1528        let query = TermQuery::flat(vec![b"alpha".to_vec(); 128], &[DEFAULT_FIELD]);
1529        let prepared =
1530            CandidateScoring::new(&index, &query, Bm25Params::default()).expect("prepare");
1531        let mut work = CandidateScoringWork::default();
1532        let result = prepared.score_rows(0, &index.segments()[0], &[1], &mut work, |receipt| {
1533            if receipt.seeks >= 64 {
1534                Err("cancelled")
1535            } else {
1536                Ok(())
1537            }
1538        });
1539        assert!(matches!(
1540            result,
1541            Err(ControlledSearchError::Control("cancelled"))
1542        ));
1543        assert_eq!(
1544            work.seeks, 64,
1545            "unsuccessful seeks must reach cancellation checkpoints"
1546        );
1547        let (clean, clean_work) = candidate_batch(&index, &query, 0, &[1]);
1548        assert_eq!(clean.scores, vec![None]);
1549        assert_eq!(clean_work.seeks, 128);
1550    }
1551
1552    #[test]
1553    fn astra_02_candidate_work_screen() {
1554        let index = index_of(&vec!["alpha beta"; 1_000]);
1555        let query = TermQuery::flat(vec![b"alpha".to_vec(), b"beta".to_vec()], &[DEFAULT_FIELD]);
1556        for count in [1, 10, 50, 400] {
1557            let rows = (0..count).map(|row| row * 2).collect::<Vec<_>>();
1558            let (batch, work) = candidate_batch(&index, &query, 0, &rows);
1559            assert_eq!(batch.scores.len(), count as usize);
1560            assert!(batch.scores.iter().all(Option::is_some));
1561            println!(
1562                "candidate_rows={count} work={work:?} counters={:?}",
1563                batch.counters
1564            );
1565        }
1566    }
1567
1568    fn astra_weighted_live_fixture() -> (LexicalIndex, TermQuery) {
1569        let mut index = LexicalIndex::new();
1570        for (documents, live) in [
1571            (
1572                vec![
1573                    ("alpha alpha", "beta gamma"),
1574                    ("alpha deleted", "beta beta"),
1575                    ("beta", "alpha alpha"),
1576                ],
1577                vec![0, 2],
1578            ),
1579            (
1580                vec![
1581                    ("gamma", "alpha beta"),
1582                    ("alpha beta", "beta beta"),
1583                    ("omega", "omega"),
1584                    ("alpha", "beta beta beta"),
1585                ],
1586                vec![0, 1, 3],
1587            ),
1588        ] {
1589            let mut segment = SegmentIndex::new();
1590            for (title, body) in documents {
1591                let mut document = Document::new();
1592                document.set(FieldId(0), title);
1593                document.set(FieldId(1), body);
1594                segment
1595                    .push_document(&analyzer(), &document)
1596                    .expect("literal fields");
1597            }
1598            index
1599                .push_sealed_with_live_rows(
1600                    SealedSegment::seal(&segment).expect("sealed fixture"),
1601                    &DocBitmap::from_ids(live),
1602                )
1603                .expect("frozen live statistics");
1604        }
1605        let query = TermQuery {
1606            terms: vec![b"alpha".to_vec(), b"beta".to_vec(), b"alpha".to_vec()],
1607            fields: FieldWeights::new(&[
1608                (FieldId(0), 2_000),
1609                (FieldId(1), 500),
1610                (FieldId(0), 9_000),
1611            ]),
1612        };
1613        (index, query)
1614    }
1615
1616    #[test]
1617    fn astra_02_candidate_bm25_matches_exhaustive_scores() {
1618        let (index, query) = astra_weighted_live_fixture();
1619        let exhaustive =
1620            search(&index, &query, usize::MAX, Bm25Params::default()).expect("exhaustive scoring");
1621        let expected = exhaustive
1622            .hits
1623            .iter()
1624            .map(|hit| (hit.doc, hit.score.to_bits()))
1625            .collect::<std::collections::BTreeMap<_, _>>();
1626        for (ordinal, rows) in [vec![2, 0, 2], vec![3, 0, 1, 3]].iter().enumerate() {
1627            let (batch, work) = candidate_batch(&index, &query, ordinal, rows);
1628            for (&row, score) in rows.iter().zip(batch.scores) {
1629                assert_eq!(
1630                    score.map(f64::to_bits),
1631                    expected
1632                        .get(&GlobalDocId {
1633                            segment: ordinal as u32,
1634                            row
1635                        })
1636                        .copied()
1637                );
1638            }
1639            assert_eq!(
1640                work.streams_opened, 3,
1641                "preserve repeated query-term occurrences"
1642            );
1643            assert_eq!(
1644                work.rows_visited,
1645                [2, 3][ordinal],
1646                "route duplicate caller rows once"
1647            );
1648        }
1649    }
1650
1651    #[test]
1652    fn astra_02_candidate_scoring_preserves_repeated_terms_and_fields() {
1653        let (index, query) = astra_weighted_live_fixture();
1654        assert_eq!(
1655            index.corpus_stats().expect("live stats").document_count(),
1656            5
1657        );
1658        assert_eq!(index.corpus_stats().expect("live stats").total_tokens(), 18);
1659        // Literal live rows have weighted (length, alpha tf, beta tf) below.
1660        // Both terms occur in all five live documents before tf weighting.
1661        // N=5, raw avgdl=18/5, title weight=2, body weight=1/2 (floor after
1662        // merging). Query order is alpha, beta, alpha; no product scorer is
1663        // used by this independent BM25 formula.
1664        let idf = (1.0_f64 + 0.5 / 5.5).ln();
1665        for (ordinal, row, length, alpha, beta) in [
1666            (0, 0, 5, 4, 0),
1667            (0, 2, 3, 1, 2),
1668            (1, 0, 3, 0, 0),
1669            (1, 1, 5, 2, 3),
1670            (1, 3, 3, 2, 1),
1671        ] {
1672            let mut expected = 0.0;
1673            let mut matched = false;
1674            for tf in [alpha, beta, alpha] {
1675                if tf != 0 {
1676                    let tf = f64::from(tf);
1677                    expected += idf * (tf * 2.2)
1678                        / (tf + 1.2 * (0.25 + 0.75 * f64::from(length) / (18.0 / 5.0)));
1679                    matched = true;
1680                }
1681            }
1682            let (batch, _) = candidate_batch(&index, &query, ordinal, &[row]);
1683            assert_eq!(
1684                batch.scores[0].map(f64::to_bits),
1685                matched.then_some(expected.to_bits()),
1686                "segment {ordinal} row {row}"
1687            );
1688        }
1689    }
1690
1691    #[test]
1692    fn astra_09_prepared_lexical_paths_keep_weighted_repeated_term_scores() {
1693        let (index, query) = astra_weighted_live_fixture();
1694        let prepared =
1695            PreparedTermQuery::new(&index, &query, Bm25Params::beir()).expect("prepared query");
1696        let idf = (1.0_f64 + 0.5 / 5.5).ln();
1697        let mut expected = Vec::new();
1698        let mut work = CandidateScoringWork::default();
1699        for (ordinal, row, length, alpha, beta) in [
1700            (0, 0, 5, 4, 0),
1701            (0, 2, 3, 1, 2),
1702            (1, 0, 3, 0, 0),
1703            (1, 1, 5, 2, 3),
1704            (1, 3, 3, 2, 1),
1705        ] {
1706            let mut score = 0.0_f64;
1707            let mut matched = false;
1708            for tf in [alpha, beta, alpha] {
1709                if tf != 0 {
1710                    let tf = f64::from(tf);
1711                    score += idf * (tf * 2.2)
1712                        / (tf + 1.2 * (0.25 + 0.75 * f64::from(length) / (18.0 / 5.0)));
1713                    matched = true;
1714                }
1715            }
1716            let batch = CandidateScoring::prepared(&prepared)
1717                .score_rows(
1718                    ordinal,
1719                    &index.segments()[ordinal],
1720                    &[row],
1721                    &mut work,
1722                    |_| Ok::<(), Infallible>(()),
1723                )
1724                .unwrap_or_else(|error| match error {
1725                    ControlledSearchError::Index(error) => {
1726                        panic!("prepared scoring failed: {error}")
1727                    }
1728                    ControlledSearchError::Control(never) => match never {},
1729                });
1730            assert_eq!(
1731                batch.scores[0].map(f64::to_bits),
1732                matched.then_some(score.to_bits())
1733            );
1734            if matched {
1735                expected.push(ScoredDoc {
1736                    doc: GlobalDocId {
1737                        segment: ordinal as u32,
1738                        row,
1739                    },
1740                    score,
1741                });
1742            }
1743        }
1744        assert_eq!(
1745            work.scorers_prepared, 0,
1746            "all constants came from query preparation"
1747        );
1748        expected.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.doc.cmp(&b.doc)));
1749        let allow = [
1750            crate::meta::DocBitmap::from_ids([0, 2]),
1751            crate::meta::DocBitmap::from_ids([0, 1, 3]),
1752        ];
1753        let allow = [&allow[0], &allow[1]];
1754        for branch in [
1755            AllowListScoringBranch::RowDriven,
1756            AllowListScoringBranch::PostingDriven,
1757        ] {
1758            let result = search_allow_list_driven_controlled_with_branch(
1759                &index,
1760                prepared.query(),
1761                10,
1762                Bm25Params::beir(),
1763                &allow,
1764                branch,
1765                || Ok::<(), Infallible>(()),
1766                Some(&prepared),
1767            )
1768            .unwrap_or_else(|error| match error {
1769                ControlledSearchError::Index(error) => panic!("prepared scoring failed: {error}"),
1770                ControlledSearchError::Control(never) => match never {},
1771            });
1772            assert_eq!(result.hits, expected);
1773        }
1774        for strategy in [
1775            crate::fts::prune::Strategy::BlockMaxWand,
1776            crate::fts::prune::Strategy::BlockMaxMaxscore,
1777        ] {
1778            for k in [1, 3, 10] {
1779                let result = crate::fts::prune::search_pruned_prepared_filtered(
1780                    &index, &prepared, k, strategy, &allow,
1781                )
1782                .expect("fresh wider cursors");
1783                assert_eq!(result.hits, expected[..k.min(expected.len())]);
1784            }
1785        }
1786    }
1787
1788    #[test]
1789    #[ignore = "explicit paired release-process candidate screen"]
1790    fn astra_02_selective_batch_screen() {
1791        let documents = (0..65_536)
1792            .map(|row| ["alpha beta", "alpha gamma", "beta gamma", "omega"][row % 4])
1793            .collect::<Vec<_>>();
1794        let index = index_of(&documents);
1795        let queries = [
1796            vec!["alpha", "beta"],
1797            vec!["alpha", "alpha", "beta"],
1798            vec!["omega", "beta"],
1799            vec!["gamma"],
1800        ]
1801        .into_iter()
1802        .map(|terms| {
1803            TermQuery::flat(
1804                terms
1805                    .into_iter()
1806                    .map(|term| term.as_bytes().to_vec())
1807                    .collect(),
1808                &[DEFAULT_FIELD],
1809            )
1810        })
1811        .collect::<Vec<_>>();
1812        let references = queries
1813            .iter()
1814            .map(|query| {
1815                search(&index, query, usize::MAX, Bm25Params::default())
1816                    .expect("exhaustive reference")
1817                    .hits
1818                    .into_iter()
1819                    .map(|hit| (hit.doc.row, hit.score.to_bits()))
1820                    .collect::<std::collections::BTreeMap<_, _>>()
1821            })
1822            .collect::<Vec<_>>();
1823        println!(
1824            "astra02_fixture,v1,rows=65536,queries=64,warm=20,seed=0x5eed,core_candidate_only=true"
1825        );
1826        for position in 0..84 {
1827            let case = position % 64;
1828            let count = [1, 10, 50, 400][case % 4];
1829            // Fixed arithmetic routing, including unsorted sparse inputs. No
1830            // random generator or qrels affect this exact-scoring treatment.
1831            let rows = (0..count)
1832                .rev()
1833                .map(|row| ((row * (65_536 / count) + case * 37 + 0x5eed) % 65_536) as u32)
1834                .collect::<Vec<_>>();
1835            let started = std::time::Instant::now();
1836            let (batch, work) = candidate_batch(&index, &queries[case % queries.len()], 0, &rows);
1837            let elapsed = started.elapsed().as_nanos();
1838            for (row, score) in rows.iter().zip(&batch.scores) {
1839                assert_eq!(
1840                    score.map(f64::to_bits),
1841                    references[case % queries.len()].get(row).copied()
1842                );
1843            }
1844            if position >= 20 {
1845                println!(
1846                    "astra02_sample,{case},{count},{elapsed},{},{},{},{},{},{}",
1847                    work.streams_opened,
1848                    work.scorers_prepared,
1849                    work.seeks,
1850                    work.rows_visited,
1851                    batch.counters.blocks_decoded,
1852                    batch.counters.docs_evaluated
1853                );
1854            }
1855        }
1856    }
1857
1858    fn forced_allow_list_search(
1859        index: &LexicalIndex,
1860        query: &TermQuery,
1861        k: usize,
1862        allow_lists: &[&DocBitmap],
1863        branch: AllowListScoringBranch,
1864    ) -> SearchResult {
1865        match search_allow_list_driven_controlled_with_branch(
1866            index,
1867            query,
1868            k,
1869            Bm25Params::default(),
1870            allow_lists,
1871            branch,
1872            || Ok::<(), Infallible>(()),
1873            None,
1874        ) {
1875            Ok(result) => result,
1876            Err(ControlledSearchError::Index(error)) => {
1877                panic!("forced allow-list search failed: {error:?}")
1878            }
1879            Err(ControlledSearchError::Control(error)) => match error {},
1880        }
1881    }
1882
1883    #[test]
1884    fn dense_allow_list_search_is_posting_driven() {
1885        let analyzer = analyzer();
1886        let mut segment = SegmentIndex::new();
1887        for _ in 0..256 {
1888            segment
1889                .push_document(&analyzer, &Document::with_text("alpha beta"))
1890                .expect("indexable");
1891        }
1892        let mut index = LexicalIndex::new();
1893        index.push_segment(segment).expect("seals");
1894        let query = TermQuery::flat(vec![b"alpha".to_vec(), b"beta".to_vec()], &[DEFAULT_FIELD]);
1895        let allow_list = DocBitmap::full(256);
1896        let result =
1897            search_allow_list_driven(&index, &query, 256, Bm25Params::default(), &[&allow_list])
1898                .expect("scores");
1899
1900        assert_eq!(result.hits.len(), 256);
1901        assert_eq!(result.counters.blocks_decoded, 8);
1902    }
1903
1904    #[test]
1905    fn posting_driven_matches_row_driven_hit_for_hit() {
1906        let name = "fts::search::tests::posting_driven_matches_row_driven_hit_for_hit";
1907        let mut seeded = crate::test_support::seeded_rng(name);
1908        let mut runner = TestRunner::new(Config {
1909            cases: 128,
1910            rng_seed: RngSeed::Fixed(seeded.next_u64()),
1911            ..Config::default()
1912        });
1913        let cases = (
1914            prop::collection::vec((0_u8..16, any::<bool>()), 2..130),
1915            prop::collection::vec(0_u8..6, 1..7),
1916            any::<u16>(),
1917            any::<u8>(),
1918        );
1919
1920        let result = runner.run(&cases, |(rows, query_terms, cut_seed, k_seed)| {
1921            let analyzer = analyzer();
1922            let cut = 1 + usize::from(cut_seed) % rows.len().saturating_sub(1);
1923            let mut first = SegmentIndex::new();
1924            let mut second = SegmentIndex::new();
1925
1926            for (position, (mask, _)) in rows.iter().enumerate() {
1927                let mut words = Vec::new();
1928                for (bit, term) in [(1_u8, "alpha"), (2, "beta"), (4, "gamma"), (8, "delta")] {
1929                    if mask & bit != 0 {
1930                        words.push(term);
1931                        if position.is_multiple_of(3) {
1932                            words.push(term);
1933                        }
1934                    }
1935                }
1936                if words.is_empty() {
1937                    words.push("filler");
1938                }
1939                let document = Document::with_text(&words.join(" "));
1940                if position < cut {
1941                    first
1942                        .push_document(&analyzer, &document)
1943                        .expect("indexable first segment");
1944                } else {
1945                    second
1946                        .push_document(&analyzer, &document)
1947                        .expect("indexable second segment");
1948                }
1949            }
1950
1951            let mut index = LexicalIndex::new();
1952            index.push_segment(first).expect("seals first segment");
1953            index.push_segment(second).expect("seals second segment");
1954            let terms = query_terms
1955                .iter()
1956                .map(|term| match term {
1957                    0 => b"alpha".to_vec(),
1958                    1 => b"beta".to_vec(),
1959                    2 => b"gamma".to_vec(),
1960                    3 => b"delta".to_vec(),
1961                    4 => b"missing".to_vec(),
1962                    _ => b"alpha".to_vec(),
1963                })
1964                .collect();
1965            let query = TermQuery::flat(terms, &[DEFAULT_FIELD]);
1966            let k = usize::from(k_seed) % rows.len().saturating_add(1);
1967
1968            for dense in [false, true] {
1969                let mut first_allowed = Vec::new();
1970                let mut second_allowed = Vec::new();
1971                for (position, (_, selected)) in rows.iter().enumerate() {
1972                    let allowed = if dense {
1973                        *selected || !position.is_multiple_of(4)
1974                    } else {
1975                        *selected && position.is_multiple_of(16)
1976                    };
1977                    if allowed && position < cut {
1978                        first_allowed.push(u32::try_from(position).unwrap_or(u32::MAX));
1979                    } else if allowed {
1980                        second_allowed
1981                            .push(u32::try_from(position.saturating_sub(cut)).unwrap_or(u32::MAX));
1982                    }
1983                }
1984                let first_allow_list = DocBitmap::from_ids(first_allowed);
1985                let second_allow_list = DocBitmap::from_ids(second_allowed);
1986                let allow_lists = [&first_allow_list, &second_allow_list];
1987                let row_driven = forced_allow_list_search(
1988                    &index,
1989                    &query,
1990                    k,
1991                    &allow_lists,
1992                    AllowListScoringBranch::RowDriven,
1993                );
1994                let posting_driven = forced_allow_list_search(
1995                    &index,
1996                    &query,
1997                    k,
1998                    &allow_lists,
1999                    AllowListScoringBranch::PostingDriven,
2000                );
2001
2002                prop_assert_eq!(row_driven.hits.len(), posting_driven.hits.len());
2003                for (row_hit, posting_hit) in row_driven.hits.iter().zip(&posting_driven.hits) {
2004                    prop_assert_eq!(row_hit.doc, posting_hit.doc);
2005                    prop_assert_eq!(row_hit.score.to_bits(), posting_hit.score.to_bits());
2006                }
2007            }
2008            Ok(())
2009        });
2010        assert!(result.is_ok(), "property result: {result:?}");
2011    }
2012
2013    #[test]
2014    fn divisor_rule_keeps_the_selective_boundary_row_driven() {
2015        assert_eq!(
2016            allow_list_scoring_branch(1, 64),
2017            AllowListScoringBranch::RowDriven
2018        );
2019        assert_eq!(
2020            allow_list_scoring_branch(2, 128),
2021            AllowListScoringBranch::RowDriven
2022        );
2023        assert_eq!(
2024            allow_list_scoring_branch(2, 127),
2025            AllowListScoringBranch::PostingDriven
2026        );
2027    }
2028
2029    #[test]
2030    fn a_repeated_field_is_collapsed_rather_than_counted_twice() {
2031        // Weight tables are normalized at construction, so iterating the
2032        // table and looking each field up by id cannot disagree. Before
2033        // normalization a repeated field added its length to the document
2034        // twice, which lowered every score for that document.
2035        let weights = FieldWeights::flat(&[FieldId(1), FieldId(1), FieldId(0)]);
2036        assert_eq!(weights.fields(), vec![FieldId(0), FieldId(1)]);
2037        let table: Vec<(FieldId, u32)> = weights.iter().collect();
2038        assert_eq!(table, vec![(FieldId(0), 1_000), (FieldId(1), 1_000)]);
2039        for (field, weight) in weights.iter() {
2040            assert_eq!(weight, weights.weight(field));
2041        }
2042        let explicit = FieldWeights::new(&[(FieldId(2), 7), (FieldId(2), 9)]);
2043        assert_eq!(explicit.iter().count(), 1);
2044        assert_eq!(explicit.weight(FieldId(2)), 7);
2045    }
2046
2047    #[test]
2048    fn the_flat_single_field_length_array_is_borrowed_not_rebuilt() {
2049        // The flat scorer asks for `length * 1000 / 1000`, which is the
2050        // field's own dense array. Borrowing it is what makes the per-query
2051        // length preparation free; the identity must be exact, not close.
2052        let index = index_of(&["a b c", "d", "e f"]);
2053        let segment = index.segments().first().expect("one segment");
2054        let weights = FieldWeights::flat(&[DEFAULT_FIELD]);
2055        let lengths = weighted_lengths(segment, &weights);
2056        assert!(
2057            matches!(lengths, std::borrow::Cow::Borrowed(_)),
2058            "the flat single-field case must borrow"
2059        );
2060        assert_eq!(lengths.as_ref(), &[3, 1, 2]);
2061        for row in 0..segment.row_count() {
2062            assert_eq!(
2063                lengths.get(row as usize).copied(),
2064                Some(row_length(segment, row, &weights))
2065            );
2066        }
2067        // The flat MULTI-field case must borrow too: it is what every BEIR
2068        // run uses, and rebuilding it per query cost a row-count-long pass
2069        // on every one of them.
2070        let mut two = SegmentIndex::new();
2071        let analyzer = analyzer();
2072        for (title, body) in [("alpha one", "body two three"), ("beta", "more body text")] {
2073            let mut document = Document::new();
2074            document.set(FieldId(0), title);
2075            document.set(FieldId(1), body);
2076            two.push_document(&analyzer, &document).expect("indexable");
2077        }
2078        let two = crate::fts::sealed::SealedSegment::seal(&two).expect("seals");
2079        let flat = FieldWeights::flat(&[FieldId(0), FieldId(1)]);
2080        let borrowed = weighted_lengths(&two, &flat);
2081        assert!(
2082            matches!(borrowed, std::borrow::Cow::Borrowed(_)),
2083            "the flat multi-field case must borrow, not rebuild per query"
2084        );
2085        for row in 0..two.row_count() {
2086            assert_eq!(
2087                borrowed.get(row as usize).copied(),
2088                Some(row_length(&two, row, &flat)),
2089                "the borrowed total must equal the weighted length it replaces"
2090            );
2091        }
2092        // A field the segment has but the query does not name must NOT be
2093        // folded into the borrowed total.
2094        let partial = FieldWeights::flat(&[FieldId(1)]);
2095        let only_body = weighted_lengths(&two, &partial);
2096        for row in 0..two.row_count() {
2097            assert_eq!(
2098                only_body.get(row as usize).copied(),
2099                Some(row_length(&two, row, &partial))
2100            );
2101        }
2102
2103        // A weight other than unity, or a second field, must materialize.
2104        let scaled = FieldWeights::new(&[(DEFAULT_FIELD, 500)]);
2105        assert!(matches!(
2106            weighted_lengths(segment, &scaled),
2107            std::borrow::Cow::Owned(_)
2108        ));
2109        assert_eq!(weighted_lengths(segment, &scaled).as_ref(), &[1, 0, 1]);
2110    }
2111
2112    #[test]
2113    fn a_query_returns_only_matching_documents_best_first() {
2114        let index = index_of(&["alpha beta", "alpha alpha gamma", "delta"]);
2115        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
2116        let result = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2117        assert_eq!(result.hits.len(), 2);
2118        // Row 1 has two occurrences of alpha, so it must rank first.
2119        assert_eq!(result.hits[0].doc, GlobalDocId { segment: 0, row: 1 });
2120        assert_eq!(result.hits[1].doc, GlobalDocId { segment: 0, row: 0 });
2121        assert!(result.hits[0].score > result.hits[1].score);
2122    }
2123
2124    #[test]
2125    fn top_k_truncates_without_reordering() {
2126        let index = index_of(&["a x", "a a x", "a a a x"]);
2127        let query = TermQuery::flat(vec![b"a".to_vec()], &[DEFAULT_FIELD]);
2128        let full = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2129        let truncated = search(&index, &query, 2, Bm25Params::default()).expect("scores");
2130        assert_eq!(truncated.hits, full.hits[..2]);
2131    }
2132
2133    #[test]
2134    fn the_tie_break_is_ascending_document_id() {
2135        // Three identical documents must come back in row order.
2136        let index = index_of(&["same", "same", "same"]);
2137        let query = TermQuery::flat(vec![b"same".to_vec()], &[DEFAULT_FIELD]);
2138        let result = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2139        let rows: Vec<u32> = result.hits.iter().map(|hit| hit.doc.row).collect();
2140        assert_eq!(rows, vec![0, 1, 2]);
2141    }
2142
2143    #[test]
2144    fn scores_are_comparable_across_segments() {
2145        // The same document text in two different segments must score
2146        // identically. If any statistic were segment-local it would not.
2147        let analyzer = analyzer();
2148        let mut first = SegmentIndex::new();
2149        first
2150            .push_document(&analyzer, &Document::with_text("engine tuning"))
2151            .expect("indexable");
2152        first
2153            .push_document(&analyzer, &Document::with_text("unrelated text here"))
2154            .expect("indexable");
2155        let mut second = SegmentIndex::new();
2156        second
2157            .push_document(&analyzer, &Document::with_text("engine tuning"))
2158            .expect("indexable");
2159        second
2160            .push_document(&analyzer, &Document::with_text("unrelated text here"))
2161            .expect("indexable");
2162
2163        let mut index = LexicalIndex::new();
2164        index.push_segment(first).expect("seals");
2165        index.push_segment(second).expect("seals");
2166
2167        let query = TermQuery::flat(vec![b"engine".to_vec()], &[DEFAULT_FIELD]);
2168        let result = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2169        assert_eq!(result.hits.len(), 2);
2170        assert!(
2171            (result.hits[0].score - result.hits[1].score).abs() < 1e-12,
2172            "identical documents scored differently across segments: {:?}",
2173            result.hits
2174        );
2175    }
2176
2177    #[test]
2178    fn an_unknown_term_matches_nothing_rather_than_everything() {
2179        let index = index_of(&["alpha", "beta"]);
2180        let query = TermQuery::flat(vec![b"omega".to_vec()], &[DEFAULT_FIELD]);
2181        let result = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2182        assert!(result.hits.is_empty());
2183        assert_eq!(result.counters.docs_evaluated, 0);
2184    }
2185
2186    #[test]
2187    fn an_empty_query_matches_nothing() {
2188        let index = index_of(&["alpha"]);
2189        let query = TermQuery::flat(Vec::new(), &[DEFAULT_FIELD]);
2190        let result = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2191        assert!(result.hits.is_empty());
2192    }
2193
2194    #[test]
2195    fn an_empty_index_is_a_typed_error_not_a_division_by_zero() {
2196        let index = LexicalIndex::new();
2197        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
2198        assert!(search(&index, &query, 10, Bm25Params::default()).is_err());
2199    }
2200
2201    #[test]
2202    fn the_exhaustive_path_never_skips_a_block() {
2203        let index = index_of(&["alpha beta", "alpha gamma"]);
2204        let query = TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]);
2205        let result = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2206        assert_eq!(result.counters.blocks_skipped, 0);
2207        assert!(result.counters.postings_decoded >= 2);
2208    }
2209
2210    #[test]
2211    fn multiple_query_terms_sum_their_contributions() {
2212        let index = index_of(&["alpha", "alpha beta"]);
2213        let query = TermQuery::flat(vec![b"alpha".to_vec(), b"beta".to_vec()], &[DEFAULT_FIELD]);
2214        let result = search(&index, &query, 10, Bm25Params::default()).expect("scores");
2215        // The document matching both terms must outrank the one matching one.
2216        assert_eq!(result.hits[0].doc.row, 1);
2217    }
2218
2219    #[test]
2220    fn field_weights_change_ranking_and_flat_weights_do_not() {
2221        let analyzer = analyzer();
2222        let mut segment = SegmentIndex::new();
2223        let mut first = Document::new();
2224        first.set(FieldId(0), "engine");
2225        first.set(FieldId(1), "unrelated");
2226        segment.push_document(&analyzer, &first).expect("indexable");
2227        let mut second = Document::new();
2228        second.set(FieldId(0), "unrelated");
2229        second.set(FieldId(1), "engine");
2230        segment
2231            .push_document(&analyzer, &second)
2232            .expect("indexable");
2233        let mut index = LexicalIndex::new();
2234        index.push_segment(segment).expect("seals");
2235
2236        let flat = TermQuery {
2237            terms: vec![b"engine".to_vec()],
2238            fields: FieldWeights::flat(&[FieldId(0), FieldId(1)]),
2239        };
2240        assert!(flat.fields.is_flat());
2241        let flat_result = search(&index, &flat, 10, Bm25Params::default()).expect("scores");
2242        assert!(
2243            (flat_result.hits[0].score - flat_result.hits[1].score).abs() < 1e-12,
2244            "flat weights must not prefer a field"
2245        );
2246
2247        let weighted = TermQuery {
2248            terms: vec![b"engine".to_vec()],
2249            fields: FieldWeights::new(&[(FieldId(0), 4_000), (FieldId(1), 1_000)]),
2250        };
2251        assert!(!weighted.fields.is_flat());
2252        let weighted_result = search(&index, &weighted, 10, Bm25Params::default()).expect("scores");
2253        assert_eq!(
2254            weighted_result.hits[0].doc.row, 0,
2255            "the heavier field must win"
2256        );
2257    }
2258
2259    #[test]
2260    fn field_weight_lookup_reports_zero_for_unweighted_fields() {
2261        let weights = FieldWeights::new(&[(FieldId(1), 2_000)]);
2262        assert_eq!(weights.weight(FieldId(1)), 2_000);
2263        assert_eq!(weights.weight(FieldId(7)), 0);
2264        assert_eq!(weights.fields(), vec![FieldId(1)]);
2265        assert!(FieldWeights::new(&[]).is_flat());
2266    }
2267}