Skip to main content

hyphae_retrieval/
lexical.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Deterministic provider-free lexical retrieval under semantics v1.
4
5use std::{
6    collections::{BTreeMap, BTreeSet},
7    convert::Infallible,
8    time::{Duration, Instant},
9};
10
11use hyphae_core::VectorSpaceName;
12use hyphae_query::{FieldPath, Record, Value};
13use thiserror::Error;
14use unicode_casefold::UnicodeCaseFold;
15use unicode_normalization::UnicodeNormalization;
16
17/// Maximum UTF-8 token length retained by tokenizer v1.
18pub const MAX_LEXICAL_TOKEN_BYTES: usize = 256;
19/// Maximum positive field weight.
20pub const MAX_LEXICAL_FIELD_WEIGHT_MICROS: u32 = 1_000_000_000;
21/// Maximum fields in one lexical definition.
22pub const MAX_LEXICAL_FIELDS: usize = 64;
23/// Maximum exact segments in one configured path.
24pub const MAX_LEXICAL_PATH_SEGMENTS: usize = 32;
25/// Maximum UTF-8 bytes in one path segment.
26pub const MAX_LEXICAL_PATH_SEGMENT_BYTES: usize = 1_024;
27const WEIGHT_SCALE: f64 = 1_000_000.0;
28const K1: f64 = 1.2;
29const B: f64 = 0.75;
30
31/// One configured document field.
32#[derive(Clone, Debug, Eq, PartialEq)]
33pub struct LexicalField {
34    /// Exact canonical document field path.
35    pub path: FieldPath,
36    /// Positive weight in millionths.
37    pub weight_micros: u32,
38}
39
40/// Immutable named lexical-index definition.
41#[derive(Clone, Debug, Eq, PartialEq)]
42pub struct LexicalIndexDefinition {
43    /// Canonical index identifier.
44    pub name: VectorSpaceName,
45    /// Unique fields sorted by exact path.
46    pub fields: Vec<LexicalField>,
47}
48
49impl LexicalIndexDefinition {
50    /// Constructs and canonicalizes one definition.
51    ///
52    /// # Errors
53    ///
54    /// Rejects empty definitions, empty paths, duplicate paths, and invalid
55    /// weights.
56    pub fn new(name: VectorSpaceName, mut fields: Vec<LexicalField>) -> Result<Self, LexicalError> {
57        if fields.is_empty() {
58            return Err(LexicalError::EmptyFields);
59        }
60        if fields.len() > MAX_LEXICAL_FIELDS {
61            return Err(LexicalError::TooManyFields);
62        }
63        if fields.iter().any(|field| field.path.segments().is_empty()) {
64            return Err(LexicalError::EmptyFieldPath);
65        }
66        if fields.iter().any(|field| {
67            field.path.segments().len() > MAX_LEXICAL_PATH_SEGMENTS
68                || field.path.segments().iter().any(|segment| {
69                    segment.is_empty() || segment.len() > MAX_LEXICAL_PATH_SEGMENT_BYTES
70                })
71        }) {
72            return Err(LexicalError::InvalidFieldSegment);
73        }
74        if fields
75            .iter()
76            .any(|field| !(1..=MAX_LEXICAL_FIELD_WEIGHT_MICROS).contains(&field.weight_micros))
77        {
78            return Err(LexicalError::InvalidFieldWeight);
79        }
80        fields.sort_by(|left, right| left.path.cmp(&right.path));
81        if fields.windows(2).any(|pair| pair[0].path == pair[1].path) {
82            return Err(LexicalError::DuplicateFieldPath);
83        }
84        Ok(Self { name, fields })
85    }
86}
87
88/// Complete lexical retrieval request.
89#[derive(Clone, Debug, Eq, PartialEq)]
90pub struct LexicalRequest {
91    /// Named durable definition.
92    pub index: VectorSpaceName,
93    /// UTF-8 query analyzed by tokenizer v1.
94    pub query: String,
95    /// Maximum returned documents.
96    pub limit: usize,
97}
98
99/// Complete bounded execution policy.
100#[derive(Clone, Debug, Eq, PartialEq)]
101pub struct LexicalLimits {
102    /// Maximum documents inspected.
103    pub max_documents: u64,
104    /// Maximum normalized tokens across corpus and query.
105    pub max_tokens: u64,
106    /// Maximum documents retained after matching.
107    pub max_candidates: u64,
108    /// Maximum returned documents.
109    pub max_returned: usize,
110    /// Cooperative timeout.
111    pub timeout: Duration,
112}
113
114impl Default for LexicalLimits {
115    fn default() -> Self {
116        Self {
117            max_documents: 1_000_000,
118            max_tokens: 10_000_000,
119            max_candidates: 100_000,
120            max_returned: 1_000,
121            timeout: Duration::from_secs(30),
122        }
123    }
124}
125
126/// One field contribution for one query term.
127#[derive(Clone, Debug, Eq, PartialEq)]
128pub struct LexicalFieldContribution {
129    /// Canonical field path.
130    pub path: FieldPath,
131    /// Raw term frequency.
132    pub term_frequency: u64,
133    /// Field token length.
134    pub field_length: u64,
135}
136
137/// One canonical query-term explanation.
138#[derive(Clone, Debug, Eq, PartialEq)]
139pub struct LexicalTermContribution {
140    /// Canonical token.
141    pub token: String,
142    /// Corpus document frequency.
143    pub document_frequency: u64,
144    /// Quantized contribution to the final score.
145    pub score_nanos: i64,
146    /// Configured fields in canonical order.
147    pub fields: Vec<LexicalFieldContribution>,
148}
149
150/// One canonical lexical match.
151#[derive(Clone, Debug, Eq, PartialEq)]
152pub struct LexicalMatch {
153    /// Binary object key.
154    pub key: Vec<u8>,
155    /// Canonical BM25F score in nanos.
156    pub score_nanos: i64,
157    /// Per-term deterministic explanation.
158    pub terms: Vec<LexicalTermContribution>,
159}
160
161/// Stable normal abstention reason.
162#[derive(Clone, Copy, Debug, Eq, PartialEq)]
163pub enum LexicalAbstentionReason {
164    /// No document contains any normalized query token.
165    NoCandidates,
166}
167
168/// Stable normal abstention evidence.
169#[derive(Clone, Debug, Eq, PartialEq)]
170pub struct LexicalAbstention {
171    /// Stable reason.
172    pub reason: LexicalAbstentionReason,
173    /// Documents inspected.
174    pub scanned_documents: u64,
175    /// Canonical unique query tokens.
176    pub query_tokens: Vec<String>,
177}
178
179/// Complete lexical outcome.
180#[derive(Clone, Debug, Eq, PartialEq)]
181pub enum LexicalOutcome {
182    /// Accepted ranked documents.
183    Matches {
184        /// Final matches.
185        matches: Vec<LexicalMatch>,
186        /// Documents inspected.
187        scanned_documents: u64,
188        /// Documents containing a query token.
189        matched_documents: u64,
190        /// Canonical unique query tokens.
191        query_tokens: Vec<String>,
192    },
193    /// Typed normal abstention.
194    Abstained(LexicalAbstention),
195}
196
197/// Complete lexical execution failure.
198#[derive(Clone, Debug, Error, Eq, PartialEq)]
199pub enum LexicalError {
200    /// At least one field is required.
201    #[error("lexical definition requires at least one field")]
202    EmptyFields,
203    /// The definition exceeds the field-count bound.
204    #[error("lexical definition exceeds 64 fields")]
205    TooManyFields,
206    /// Root/empty field paths are not accepted.
207    #[error("lexical field path must be nonempty")]
208    EmptyFieldPath,
209    /// Path segments must be nonempty and bounded.
210    #[error("lexical field path contains an invalid segment")]
211    InvalidFieldSegment,
212    /// Field paths must be unique.
213    #[error("lexical field paths must be unique")]
214    DuplicateFieldPath,
215    /// Field weights must be positive and bounded.
216    #[error("lexical field weight is outside 1..=1000000000")]
217    InvalidFieldWeight,
218    /// Request and definition names differ.
219    #[error("lexical request index does not match the definition")]
220    IndexMismatch,
221    /// The normalized query has no tokens.
222    #[error("lexical query has no retained normalized tokens")]
223    EmptyQuery,
224    /// At least one result must be requested.
225    #[error("lexical result limit must be nonzero")]
226    ZeroLimit,
227    /// Requested result count exceeds policy.
228    #[error("lexical result limit {requested} exceeds maximum {maximum}")]
229    ResultLimitExceeded {
230        /// Requested count.
231        requested: usize,
232        /// Maximum count.
233        maximum: usize,
234    },
235    /// Record keys must be nonempty.
236    #[error("lexical document key must be nonempty")]
237    EmptyDocumentKey,
238    /// Record keys must be unique.
239    #[error("duplicate lexical document key")]
240    DuplicateDocumentKey,
241    /// Document budget exhausted.
242    #[error("lexical document budget exceeded: {maximum}")]
243    DocumentBudgetExceeded {
244        /// Maximum documents.
245        maximum: u64,
246    },
247    /// Token budget exhausted.
248    #[error("lexical token budget exceeded: {maximum}")]
249    TokenBudgetExceeded {
250        /// Maximum tokens.
251        maximum: u64,
252    },
253    /// Candidate budget exhausted.
254    #[error("lexical candidate budget exceeded: {maximum}")]
255    CandidateBudgetExceeded {
256        /// Maximum candidates.
257        maximum: u64,
258    },
259    /// Rebuildable lexical statistics are structurally inconsistent.
260    #[error("materialized lexical projection is malformed")]
261    MalformedProjection,
262    /// Cooperative deadline elapsed.
263    #[error("lexical retrieval timed out")]
264    TimedOut,
265    /// Canonical numeric operation failed.
266    #[error("lexical score arithmetic overflow or non-finite result")]
267    ArithmeticOverflow,
268}
269
270/// Applies tokenizer semantics `hyphae-unicode-tokenizer-v1`.
271pub fn tokenize_v1(input: &str) -> Vec<String> {
272    match tokenize_v1_checked(
273        input,
274        || Ok::<(), Infallible>(()),
275        || Ok::<(), Infallible>(()),
276    ) {
277        Ok(tokens) => tokens,
278        Err(never) => match never {},
279    }
280}
281
282/// Applies tokenizer v1 while cooperatively checking work and retained-token policy.
283///
284/// `checkpoint` runs throughout normalization, including within one very long
285/// token. `accept_token` runs immediately before each retained token is added.
286///
287/// # Errors
288///
289/// Returns the first error produced by either callback without returning a
290/// partial token list.
291pub fn tokenize_v1_checked<E>(
292    input: &str,
293    mut checkpoint: impl FnMut() -> Result<(), E>,
294    mut accept_token: impl FnMut() -> Result<(), E>,
295) -> Result<Vec<String>, E> {
296    const CHECKPOINT_INTERVAL: usize = 256;
297
298    checkpoint()?;
299    let mut tokens = Vec::new();
300    let mut token = String::new();
301    let mut discarding_oversized_token = false;
302    for (index, character) in input.nfkc().case_fold().enumerate() {
303        if index % CHECKPOINT_INTERVAL == 0 {
304            checkpoint()?;
305        }
306        if character.is_alphanumeric() {
307            if !discarding_oversized_token {
308                let next_length = token.len().saturating_add(character.len_utf8());
309                if next_length <= MAX_LEXICAL_TOKEN_BYTES {
310                    token.push(character);
311                } else {
312                    token.clear();
313                    discarding_oversized_token = true;
314                }
315            }
316        } else {
317            push_token_checked(
318                &mut tokens,
319                &mut token,
320                &mut discarding_oversized_token,
321                &mut accept_token,
322            )?;
323        }
324    }
325    checkpoint()?;
326    push_token_checked(
327        &mut tokens,
328        &mut token,
329        &mut discarding_oversized_token,
330        &mut accept_token,
331    )?;
332    Ok(tokens)
333}
334
335fn push_token_checked<E>(
336    tokens: &mut Vec<String>,
337    token: &mut String,
338    discarding_oversized_token: &mut bool,
339    accept_token: &mut impl FnMut() -> Result<(), E>,
340) -> Result<(), E> {
341    if !*discarding_oversized_token && !token.is_empty() {
342        accept_token()?;
343        tokens.push(std::mem::take(token));
344    } else {
345        token.clear();
346    }
347    *discarding_oversized_token = false;
348    Ok(())
349}
350
351struct AnalyzedDocument {
352    key: Vec<u8>,
353    fields: Vec<Vec<String>>,
354}
355
356#[derive(Clone, Copy)]
357struct LexicalDeadline {
358    started: Instant,
359    timeout: Duration,
360}
361
362impl LexicalDeadline {
363    fn check(self) -> Result<(), LexicalError> {
364        check_timeout(self.started, self.timeout)
365    }
366}
367
368struct ScoringContext<'a> {
369    document_count: u64,
370    averages: &'a [f64],
371    frequencies: &'a BTreeMap<String, u64>,
372    definition: &'a LexicalIndexDefinition,
373    query_tokens: &'a [String],
374    deadline: LexicalDeadline,
375}
376
377/// Rebuildable lexical statistics for one document that contains at least one
378/// normalized query term.
379#[derive(Clone, Debug, Eq, PartialEq)]
380pub struct LexicalMaterializedDocument {
381    /// Binary document key.
382    pub key: Vec<u8>,
383    /// Token count for every configured field in canonical field order.
384    pub field_lengths: Vec<u64>,
385    /// Per-field frequencies for each canonical query token.
386    pub term_frequencies: BTreeMap<String, Vec<u64>>,
387}
388
389/// Complete bounded view read from a rebuildable lexical projection.
390#[derive(Clone, Debug, Eq, PartialEq)]
391pub struct LexicalMaterializedCorpus {
392    /// Number of authoritative documents represented by the projection.
393    pub document_count: u64,
394    /// Total normalized token count across all configured fields.
395    pub token_count: u64,
396    /// Per-field token totals across the complete corpus.
397    pub total_field_lengths: Vec<u64>,
398    /// Candidate documents containing at least one query token.
399    pub documents: Vec<LexicalMaterializedDocument>,
400}
401
402/// Executes the deterministic BM25F-compatible reference algorithm.
403///
404/// # Errors
405///
406/// Returns an invalid-input, budget, timeout, or numeric error and never a
407/// partial ranking.
408pub fn retrieve_lexical(
409    records: &[Record],
410    definition: &LexicalIndexDefinition,
411    request: &LexicalRequest,
412    limits: &LexicalLimits,
413) -> Result<LexicalOutcome, LexicalError> {
414    validate_request(definition, request, limits)?;
415    let started = Instant::now();
416    let query_tokens = tokenize_before_deadline(&request.query, started, limits.timeout)?
417        .into_iter()
418        .collect::<BTreeSet<_>>()
419        .into_iter()
420        .collect::<Vec<_>>();
421    if query_tokens.is_empty() {
422        return Err(LexicalError::EmptyQuery);
423    }
424    let mut token_count = u64::try_from(query_tokens.len()).unwrap_or(u64::MAX);
425    if token_count > limits.max_tokens {
426        return Err(LexicalError::TokenBudgetExceeded {
427            maximum: limits.max_tokens,
428        });
429    }
430    let mut keys = BTreeSet::new();
431    let mut documents = Vec::with_capacity(records.len());
432    let mut total_lengths = vec![0_u64; definition.fields.len()];
433    for record in records {
434        check_timeout(started, limits.timeout)?;
435        if u64::try_from(documents.len()).unwrap_or(u64::MAX) >= limits.max_documents {
436            return Err(LexicalError::DocumentBudgetExceeded {
437                maximum: limits.max_documents,
438            });
439        }
440        if record.key.is_empty() {
441            return Err(LexicalError::EmptyDocumentKey);
442        }
443        if !keys.insert(record.key.as_slice()) {
444            return Err(LexicalError::DuplicateDocumentKey);
445        }
446        let mut fields = Vec::with_capacity(definition.fields.len());
447        for (field_index, field) in definition.fields.iter().enumerate() {
448            let tokens = match field.path.resolve(&record.value) {
449                Some(Value::String(value)) => {
450                    tokenize_with_limits(value, &mut token_count, started, limits)?
451                }
452                _ => Vec::new(),
453            };
454            let length = u64::try_from(tokens.len()).unwrap_or(u64::MAX);
455            total_lengths[field_index] = total_lengths[field_index]
456                .checked_add(length)
457                .ok_or(LexicalError::ArithmeticOverflow)?;
458            fields.push(tokens);
459        }
460        documents.push(AnalyzedDocument {
461            key: record.key.clone(),
462            fields,
463        });
464    }
465    score_documents(
466        &documents,
467        &total_lengths,
468        definition,
469        request,
470        limits,
471        &query_tokens,
472        started,
473    )
474}
475
476/// Executes BM25F from a rebuildable materialized lexical projection.
477///
478/// # Errors
479///
480/// Returns an invalid-input, malformed projection, budget, timeout, or
481/// numeric error and never a partial ranking.
482pub fn retrieve_lexical_materialized(
483    corpus: &LexicalMaterializedCorpus,
484    definition: &LexicalIndexDefinition,
485    request: &LexicalRequest,
486    limits: &LexicalLimits,
487) -> Result<LexicalOutcome, LexicalError> {
488    validate_request(definition, request, limits)?;
489    let started = Instant::now();
490    let query_tokens = tokenize_before_deadline(&request.query, started, limits.timeout)?
491        .into_iter()
492        .collect::<BTreeSet<_>>()
493        .into_iter()
494        .collect::<Vec<_>>();
495    if query_tokens.is_empty() {
496        return Err(LexicalError::EmptyQuery);
497    }
498    validate_materialized_corpus(corpus, definition, &query_tokens, limits, started)?;
499    let averages = corpus
500        .total_field_lengths
501        .iter()
502        .map(|length| {
503            if corpus.document_count == 0 {
504                0.0
505            } else {
506                bounded_count_as_f64(*length) / bounded_count_as_f64(corpus.document_count)
507            }
508        })
509        .collect::<Vec<_>>();
510    let mut frequencies = BTreeMap::new();
511    for token in &query_tokens {
512        check_timeout(started, limits.timeout)?;
513        let mut frequency = 0_u64;
514        for document in &corpus.documents {
515            check_timeout(started, limits.timeout)?;
516            if document
517                .term_frequencies
518                .get(token)
519                .is_some_and(|fields| fields.iter().any(|value| *value > 0))
520            {
521                frequency = frequency
522                    .checked_add(1)
523                    .ok_or(LexicalError::ArithmeticOverflow)?;
524            }
525        }
526        frequencies.insert(token.clone(), frequency);
527    }
528    let scoring = ScoringContext {
529        document_count: corpus.document_count,
530        averages: &averages,
531        frequencies: &frequencies,
532        definition,
533        query_tokens: &query_tokens,
534        deadline: LexicalDeadline {
535            started,
536            timeout: limits.timeout,
537        },
538    };
539    let mut matches = Vec::with_capacity(corpus.documents.len().min(request.limit));
540    for document in &corpus.documents {
541        scoring.deadline.check()?;
542        if let Some(matched) = score_materialized_document(document, &scoring)? {
543            matches.push(matched);
544        }
545    }
546    finish_ranking(
547        matches,
548        corpus.document_count,
549        &query_tokens,
550        request.limit,
551        started,
552        limits.timeout,
553    )
554}
555
556fn validate_materialized_corpus(
557    corpus: &LexicalMaterializedCorpus,
558    definition: &LexicalIndexDefinition,
559    query_tokens: &[String],
560    limits: &LexicalLimits,
561    started: Instant,
562) -> Result<(), LexicalError> {
563    check_timeout(started, limits.timeout)?;
564    if corpus.document_count > limits.max_documents {
565        return Err(LexicalError::DocumentBudgetExceeded {
566            maximum: limits.max_documents,
567        });
568    }
569    let total_tokens = corpus
570        .token_count
571        .checked_add(u64::try_from(query_tokens.len()).unwrap_or(u64::MAX))
572        .ok_or(LexicalError::TokenBudgetExceeded {
573            maximum: limits.max_tokens,
574        })?;
575    if total_tokens > limits.max_tokens {
576        return Err(LexicalError::TokenBudgetExceeded {
577            maximum: limits.max_tokens,
578        });
579    }
580    if u64::try_from(corpus.documents.len()).unwrap_or(u64::MAX) > limits.max_candidates {
581        return Err(LexicalError::CandidateBudgetExceeded {
582            maximum: limits.max_candidates,
583        });
584    }
585    if corpus.total_field_lengths.len() != definition.fields.len() {
586        return Err(LexicalError::MalformedProjection);
587    }
588    let mut keys = BTreeSet::new();
589    for document in &corpus.documents {
590        check_timeout(started, limits.timeout)?;
591        if document.key.is_empty() {
592            return Err(LexicalError::EmptyDocumentKey);
593        }
594        if !keys.insert(document.key.as_slice()) {
595            return Err(LexicalError::DuplicateDocumentKey);
596        }
597        if document.field_lengths.len() != definition.fields.len()
598            || document.term_frequencies.len() != query_tokens.len()
599        {
600            return Err(LexicalError::MalformedProjection);
601        }
602        for token in query_tokens {
603            check_timeout(started, limits.timeout)?;
604            if document
605                .term_frequencies
606                .get(token)
607                .is_none_or(|frequencies| frequencies.len() != definition.fields.len())
608            {
609                return Err(LexicalError::MalformedProjection);
610            }
611        }
612    }
613    Ok(())
614}
615
616fn tokenize_before_deadline(
617    input: &str,
618    started: Instant,
619    timeout: Duration,
620) -> Result<Vec<String>, LexicalError> {
621    tokenize_v1_checked(
622        input,
623        || check_timeout(started, timeout),
624        || check_timeout(started, timeout),
625    )
626}
627
628fn tokenize_with_limits(
629    input: &str,
630    token_count: &mut u64,
631    started: Instant,
632    limits: &LexicalLimits,
633) -> Result<Vec<String>, LexicalError> {
634    tokenize_v1_checked(
635        input,
636        || check_timeout(started, limits.timeout),
637        || {
638            *token_count = token_count
639                .checked_add(1)
640                .ok_or(LexicalError::TokenBudgetExceeded {
641                    maximum: limits.max_tokens,
642                })?;
643            if *token_count > limits.max_tokens {
644                return Err(LexicalError::TokenBudgetExceeded {
645                    maximum: limits.max_tokens,
646                });
647            }
648            Ok(())
649        },
650    )
651}
652
653fn validate_request(
654    definition: &LexicalIndexDefinition,
655    request: &LexicalRequest,
656    limits: &LexicalLimits,
657) -> Result<(), LexicalError> {
658    if request.index != definition.name {
659        return Err(LexicalError::IndexMismatch);
660    }
661    if request.limit == 0 {
662        return Err(LexicalError::ZeroLimit);
663    }
664    if request.limit > limits.max_returned {
665        return Err(LexicalError::ResultLimitExceeded {
666            requested: request.limit,
667            maximum: limits.max_returned,
668        });
669    }
670    Ok(())
671}
672
673fn score_documents(
674    documents: &[AnalyzedDocument],
675    total_lengths: &[u64],
676    definition: &LexicalIndexDefinition,
677    request: &LexicalRequest,
678    limits: &LexicalLimits,
679    query_tokens: &[String],
680    started: Instant,
681) -> Result<LexicalOutcome, LexicalError> {
682    check_timeout(started, limits.timeout)?;
683    let document_count = u64::try_from(documents.len()).unwrap_or(u64::MAX);
684    let averages = total_lengths
685        .iter()
686        .map(|length| {
687            if document_count == 0 {
688                0.0
689            } else {
690                bounded_count_as_f64(*length) / bounded_count_as_f64(document_count)
691            }
692        })
693        .collect::<Vec<_>>();
694    let mut frequencies = BTreeMap::new();
695    for token in query_tokens {
696        check_timeout(started, limits.timeout)?;
697        let mut count = 0_u64;
698        for document in documents {
699            check_timeout(started, limits.timeout)?;
700            let mut present = false;
701            'fields: for field in &document.fields {
702                for candidate in field {
703                    check_timeout(started, limits.timeout)?;
704                    if candidate == token {
705                        present = true;
706                        break 'fields;
707                    }
708                }
709            }
710            if present {
711                count = count
712                    .checked_add(1)
713                    .ok_or(LexicalError::ArithmeticOverflow)?;
714            }
715        }
716        frequencies.insert(token.clone(), count);
717    }
718    let scoring = ScoringContext {
719        document_count,
720        averages: &averages,
721        frequencies: &frequencies,
722        definition,
723        query_tokens,
724        deadline: LexicalDeadline {
725            started,
726            timeout: limits.timeout,
727        },
728    };
729    let mut matches = Vec::new();
730    for document in documents {
731        scoring.deadline.check()?;
732        let document_match = score_document(document, &scoring)?;
733        if let Some(matched) = document_match {
734            if u64::try_from(matches.len()).unwrap_or(u64::MAX) >= limits.max_candidates {
735                return Err(LexicalError::CandidateBudgetExceeded {
736                    maximum: limits.max_candidates,
737                });
738            }
739            matches.push(matched);
740        }
741    }
742    finish_ranking(
743        matches,
744        document_count,
745        query_tokens,
746        request.limit,
747        started,
748        limits.timeout,
749    )
750}
751
752fn score_document(
753    document: &AnalyzedDocument,
754    context: &ScoringContext<'_>,
755) -> Result<Option<LexicalMatch>, LexicalError> {
756    context.deadline.check()?;
757    let field_lengths = document
758        .fields
759        .iter()
760        .map(|field| u64::try_from(field.len()).unwrap_or(u64::MAX))
761        .collect::<Vec<_>>();
762    let mut term_frequencies = BTreeMap::new();
763    for token in context.query_tokens {
764        context.deadline.check()?;
765        let mut per_field = Vec::with_capacity(document.fields.len());
766        for field in &document.fields {
767            let mut frequency = 0_u64;
768            for candidate in field {
769                context.deadline.check()?;
770                if candidate == token {
771                    frequency = frequency
772                        .checked_add(1)
773                        .ok_or(LexicalError::ArithmeticOverflow)?;
774                }
775            }
776            per_field.push(frequency);
777        }
778        term_frequencies.insert(token.clone(), per_field);
779    }
780    score_statistics(&document.key, &field_lengths, &term_frequencies, context)
781}
782
783fn score_materialized_document(
784    document: &LexicalMaterializedDocument,
785    context: &ScoringContext<'_>,
786) -> Result<Option<LexicalMatch>, LexicalError> {
787    score_statistics(
788        &document.key,
789        &document.field_lengths,
790        &document.term_frequencies,
791        context,
792    )
793}
794
795fn score_statistics(
796    key: &[u8],
797    field_lengths: &[u64],
798    term_frequencies: &BTreeMap<String, Vec<u64>>,
799    context: &ScoringContext<'_>,
800) -> Result<Option<LexicalMatch>, LexicalError> {
801    context.deadline.check()?;
802    let mut terms = Vec::new();
803    let mut score_nanos = 0_i64;
804    for token in context.query_tokens {
805        context.deadline.check()?;
806        let document_frequency = context.frequencies[token];
807        if document_frequency == 0 {
808            continue;
809        }
810        let mut combined_tf = 0.0_f64;
811        let mut fields = Vec::with_capacity(context.definition.fields.len());
812        for (index, definition_field) in context.definition.fields.iter().enumerate() {
813            context.deadline.check()?;
814            let term_frequency = term_frequencies[token][index];
815            let field_length = field_lengths[index];
816            fields.push(LexicalFieldContribution {
817                path: definition_field.path.clone(),
818                term_frequency,
819                field_length,
820            });
821            if term_frequency > 0 && context.averages[index] > 0.0 {
822                let normalization =
823                    1.0 - B + B * bounded_count_as_f64(field_length) / context.averages[index];
824                combined_tf += (f64::from(definition_field.weight_micros) / WEIGHT_SCALE)
825                    * bounded_count_as_f64(term_frequency)
826                    / normalization;
827            }
828        }
829        if combined_tf == 0.0 {
830            continue;
831        }
832        let numerator =
833            bounded_count_as_f64(context.document_count.saturating_sub(document_frequency)) + 0.5;
834        let denominator = bounded_count_as_f64(document_frequency) + 0.5;
835        let idf = libm::log(1.0 + numerator / denominator);
836        let term_score = quantize_score(idf * combined_tf * (K1 + 1.0) / (combined_tf + K1))?;
837        score_nanos = score_nanos
838            .checked_add(term_score)
839            .ok_or(LexicalError::ArithmeticOverflow)?;
840        terms.push(LexicalTermContribution {
841            token: token.clone(),
842            document_frequency,
843            score_nanos: term_score,
844            fields,
845        });
846    }
847    Ok((score_nanos > 0).then(|| LexicalMatch {
848        key: key.to_vec(),
849        score_nanos,
850        terms,
851    }))
852}
853
854fn finish_ranking(
855    mut matches: Vec<LexicalMatch>,
856    document_count: u64,
857    query_tokens: &[String],
858    limit: usize,
859    started: Instant,
860    timeout: Duration,
861) -> Result<LexicalOutcome, LexicalError> {
862    check_timeout(started, timeout)?;
863    matches.sort_by(|left, right| {
864        right
865            .score_nanos
866            .cmp(&left.score_nanos)
867            .then_with(|| left.key.cmp(&right.key))
868    });
869    check_timeout(started, timeout)?;
870    let matched_documents = u64::try_from(matches.len()).unwrap_or(u64::MAX);
871    matches.truncate(limit);
872    Ok(if matches.is_empty() {
873        LexicalOutcome::Abstained(LexicalAbstention {
874            reason: LexicalAbstentionReason::NoCandidates,
875            scanned_documents: document_count,
876            query_tokens: query_tokens.to_vec(),
877        })
878    } else {
879        LexicalOutcome::Matches {
880            matches,
881            scanned_documents: document_count,
882            matched_documents,
883            query_tokens: query_tokens.to_vec(),
884        }
885    })
886}
887
888fn quantize_score(value: f64) -> Result<i64, LexicalError> {
889    if !value.is_finite() || value < 0.0 {
890        return Err(LexicalError::ArithmeticOverflow);
891    }
892    let scaled = value * 1_000_000_000.0;
893    if !scaled.is_finite() {
894        return Err(LexicalError::ArithmeticOverflow);
895    }
896    if scaled >= maximum_i64_as_f64() {
897        return Ok(i64::MAX);
898    }
899    Ok(rounded_nonnegative_f64_as_i64(scaled))
900}
901
902/// Counts accepted by lexical execution are bounded far below the 53-bit
903/// integer precision of `f64`, so this conversion is exact.
904#[allow(clippy::cast_precision_loss)]
905fn bounded_count_as_f64(value: u64) -> f64 {
906    value as f64
907}
908
909#[allow(clippy::cast_precision_loss)]
910fn maximum_i64_as_f64() -> f64 {
911    i64::MAX as f64
912}
913
914/// The caller has already checked finiteness, non-negativity, and the i64
915/// upper bound.
916#[allow(clippy::cast_possible_truncation)]
917fn rounded_nonnegative_f64_as_i64(value: f64) -> i64 {
918    libm::floor(value + 0.5) as i64
919}
920
921fn check_timeout(started: Instant, timeout: Duration) -> Result<(), LexicalError> {
922    if started.elapsed() >= timeout {
923        Err(LexicalError::TimedOut)
924    } else {
925        Ok(())
926    }
927}
928
929#[cfg(test)]
930mod tests {
931    use std::collections::{BTreeMap, BTreeSet};
932
933    use super::*;
934    use proptest::prelude::*;
935
936    fn record(key: &[u8], title: &str, body: &str) -> Record {
937        Record::new(
938            key,
939            Value::Object(BTreeMap::from([
940                ("title".into(), Value::String(title.into())),
941                ("body".into(), Value::String(body.into())),
942            ])),
943        )
944    }
945
946    fn definition() -> Result<LexicalIndexDefinition, LexicalError> {
947        LexicalIndexDefinition::new(
948            VectorSpaceName::new("docs").map_err(|_| LexicalError::EmptyFields)?,
949            vec![
950                LexicalField {
951                    path: FieldPath::field("body"),
952                    weight_micros: 1_000_000,
953                },
954                LexicalField {
955                    path: FieldPath::field("title"),
956                    weight_micros: 2_000_000,
957                },
958            ],
959        )
960    }
961
962    fn materialize_reference_corpus(
963        records: &[Record],
964        definition: &LexicalIndexDefinition,
965        query: &str,
966    ) -> LexicalMaterializedCorpus {
967        let query_tokens = tokenize_v1(query)
968            .into_iter()
969            .collect::<BTreeSet<_>>()
970            .into_iter()
971            .collect::<Vec<_>>();
972        let mut token_count = 0_u64;
973        let mut total_field_lengths = vec![0_u64; definition.fields.len()];
974        let mut documents = Vec::new();
975
976        for record in records {
977            let fields = definition
978                .fields
979                .iter()
980                .map(|field| match field.path.resolve(&record.value) {
981                    Some(Value::String(value)) => tokenize_v1(value),
982                    _ => Vec::new(),
983                })
984                .collect::<Vec<_>>();
985            let field_lengths = fields
986                .iter()
987                .map(|field| u64::try_from(field.len()).unwrap_or(u64::MAX))
988                .collect::<Vec<_>>();
989            for (total, length) in total_field_lengths.iter_mut().zip(&field_lengths) {
990                *total = total.saturating_add(*length);
991                token_count = token_count.saturating_add(*length);
992            }
993            let term_frequencies = query_tokens
994                .iter()
995                .map(|token| {
996                    let frequencies = fields
997                        .iter()
998                        .map(|field| {
999                            u64::try_from(
1000                                field.iter().filter(|candidate| *candidate == token).count(),
1001                            )
1002                            .unwrap_or(u64::MAX)
1003                        })
1004                        .collect::<Vec<_>>();
1005                    (token.clone(), frequencies)
1006                })
1007                .collect::<BTreeMap<_, _>>();
1008            if term_frequencies
1009                .values()
1010                .any(|frequencies| frequencies.iter().any(|frequency| *frequency > 0))
1011            {
1012                documents.push(LexicalMaterializedDocument {
1013                    key: record.key.clone(),
1014                    field_lengths,
1015                    term_frequencies,
1016                });
1017            }
1018        }
1019
1020        LexicalMaterializedCorpus {
1021            document_count: u64::try_from(records.len()).unwrap_or(u64::MAX),
1022            token_count,
1023            total_field_lengths,
1024            documents,
1025        }
1026    }
1027
1028    #[test]
1029    fn tokenizer_pins_nfkc_casefold_and_alphanumeric_runs() {
1030        assert_eq!(
1031            tokenize_v1("Straße ABC—café"),
1032            vec!["strasse", "abc", "café"]
1033        );
1034    }
1035
1036    #[test]
1037    fn checked_tokenizer_matches_v1_and_checks_inside_one_long_token() {
1038        let input = "Straße ABC—café";
1039        let mut accepted = 0_usize;
1040        let checked = match tokenize_v1_checked(
1041            input,
1042            || Ok::<(), Infallible>(()),
1043            || {
1044                accepted += 1;
1045                Ok::<(), Infallible>(())
1046            },
1047        ) {
1048            Ok(tokens) => tokens,
1049            Err(never) => match never {},
1050        };
1051        assert_eq!(checked, tokenize_v1(input));
1052        assert_eq!(accepted, checked.len());
1053        assert_eq!(
1054            tokenize_v1(&format!("{} tail", "a".repeat(257))),
1055            vec!["tail"]
1056        );
1057
1058        let long_token = "a".repeat(2_048);
1059        let mut checkpoints = 0_usize;
1060        let interrupted = tokenize_v1_checked(
1061            &long_token,
1062            || {
1063                checkpoints += 1;
1064                if checkpoints == 3 {
1065                    Err("stop")
1066                } else {
1067                    Ok(())
1068                }
1069            },
1070            || Ok::<(), &'static str>(()),
1071        );
1072        assert_eq!(interrupted, Err("stop"));
1073        assert_eq!(checkpoints, 3);
1074    }
1075
1076    #[test]
1077    fn zero_timeout_includes_query_tokenization_for_both_scorers() -> Result<(), LexicalError> {
1078        let definition = definition()?;
1079        let request = LexicalRequest {
1080            index: definition.name.clone(),
1081            query: "rust".into(),
1082            limit: 1,
1083        };
1084        let limits = LexicalLimits {
1085            timeout: Duration::ZERO,
1086            ..LexicalLimits::default()
1087        };
1088        let corpus = materialize_reference_corpus(&[], &definition, &request.query);
1089
1090        assert_eq!(
1091            retrieve_lexical(&[], &definition, &request, &limits),
1092            Err(LexicalError::TimedOut)
1093        );
1094        assert_eq!(
1095            retrieve_lexical_materialized(&corpus, &definition, &request, &limits),
1096            Err(LexicalError::TimedOut)
1097        );
1098        Ok(())
1099    }
1100
1101    #[test]
1102    fn duplicate_query_tokens_count_once_against_token_budget() -> Result<(), LexicalError> {
1103        let definition = definition()?;
1104        let request = LexicalRequest {
1105            index: definition.name.clone(),
1106            query: "rust rust".into(),
1107            limit: 1,
1108        };
1109        let limits = LexicalLimits {
1110            max_tokens: 1,
1111            ..LexicalLimits::default()
1112        };
1113        let corpus = materialize_reference_corpus(&[], &definition, &request.query);
1114
1115        let reference = retrieve_lexical(&[], &definition, &request, &limits)?;
1116        let materialized = retrieve_lexical_materialized(&corpus, &definition, &request, &limits)?;
1117        assert_eq!(reference, materialized);
1118        assert!(matches!(
1119            reference,
1120            LexicalOutcome::Abstained(LexicalAbstention {
1121                query_tokens,
1122                ..
1123            }) if query_tokens == vec!["rust"]
1124        ));
1125        Ok(())
1126    }
1127
1128    #[test]
1129    fn bm25f_is_deterministic_and_binary_key_breaks_ties() -> Result<(), LexicalError> {
1130        let definition = definition()?;
1131        let outcome = retrieve_lexical(
1132            &[
1133                record(b"b", "Rust memory", "durable engine"),
1134                record(b"a", "Rust memory", "durable engine"),
1135                record(b"z", "other", "nothing"),
1136            ],
1137            &definition,
1138            &LexicalRequest {
1139                index: definition.name.clone(),
1140                query: "RUST rust".into(),
1141                limit: 10,
1142            },
1143            &LexicalLimits::default(),
1144        )?;
1145        let LexicalOutcome::Matches { matches, .. } = outcome else {
1146            return Err(LexicalError::ArithmeticOverflow);
1147        };
1148        assert_eq!(matches[0].key, b"a");
1149        assert_eq!(matches[1].key, b"b");
1150        assert_eq!(matches[0].score_nanos, matches[1].score_nanos);
1151        Ok(())
1152    }
1153
1154    proptest! {
1155        #![proptest_config(ProptestConfig::with_cases(64))]
1156
1157        #[test]
1158        fn materialized_scorer_matches_reference_for_random_corpora(
1159            generated in prop::collection::vec(
1160                ("[a-z ]{0,24}", "[a-z ]{0,48}"),
1161                0..32
1162            ),
1163            query in "(rust|durable|engine|memory)( (rust|durable|engine|memory)){0,2}",
1164            limit in 1_usize..16
1165        ) {
1166            let definition = definition().map_err(|error| TestCaseError::fail(error.to_string()))?;
1167            let records = generated
1168                .iter()
1169                .enumerate()
1170                .map(|(index, (title, body))| {
1171                    record(&u64::try_from(index).unwrap_or(u64::MAX).to_be_bytes(), title, body)
1172                })
1173                .collect::<Vec<_>>();
1174            let request = LexicalRequest {
1175                index: definition.name.clone(),
1176                query: query.clone(),
1177                limit,
1178            };
1179            let limits = LexicalLimits::default();
1180            let reference = retrieve_lexical(&records, &definition, &request, &limits)
1181                .map_err(|error| TestCaseError::fail(error.to_string()))?;
1182            let corpus = materialize_reference_corpus(&records, &definition, &query);
1183            let materialized =
1184                retrieve_lexical_materialized(&corpus, &definition, &request, &limits)
1185                    .map_err(|error| TestCaseError::fail(error.to_string()))?;
1186
1187            prop_assert_eq!(materialized, reference);
1188        }
1189    }
1190}