Skip to main content

zeppelin_embed/fts/
query.rs

1//! Structured lexical queries and their deterministic expansion plan.
2//!
3//! Persisted source text is deliberately not interpreted here. `StoredText`
4//! remains a lossless row-addressed storage primitive; this module compiles
5//! query behavior from the indexed vocabulary. Adding another operator (for
6//! example regex) therefore does not require a segment-format change.
7
8use super::vocabulary::Vocabulary;
9
10use super::fuzzy;
11use super::index::FieldId;
12use super::phonetic;
13#[cfg(test)]
14use super::phrase;
15use super::search::{FieldWeights, TermQuery};
16use super::snippet::Highlight;
17#[cfg(test)]
18use super::tokenizer::Analyzer;
19
20/// Expansion order and one set of pinned BM25 statistics, shared by the
21/// combined top-k producer and exact scoring of supplied hybrid candidates.
22pub(crate) struct PreparedWeightedQuery {
23    expansions: Vec<LexicalExpansion>,
24    scoring: super::search::PreparedTermQuery,
25}
26
27#[cfg(test)]
28std::thread_local! {
29    static WEIGHTED_SIZE_VISITS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
30}
31
32impl PreparedWeightedQuery {
33    pub(crate) fn new_controlled<E>(
34        index: &super::index::LexicalIndex,
35        expansions: Vec<LexicalExpansion>,
36        fields: FieldWeights,
37        mut check: impl FnMut() -> Result<(), E>,
38    ) -> Result<Self, super::search::ControlledSearchError<E>> {
39        use super::search::ControlledSearchError;
40        let mut work =
41            super::control::WorkCheck::new(|| check().map_err(ControlledSearchError::Control));
42        work.check_now()?;
43        let mut terms = Vec::with_capacity(expansions.len());
44        for entry in &expansions {
45            work.step()?;
46            terms.push(entry.term.clone());
47        }
48        let query = TermQuery { terms, fields };
49        Ok(Self {
50            scoring: super::search::PreparedTermQuery::from_owned_controlled(
51                index,
52                query,
53                super::bm25::Bm25Params::beir(),
54                check,
55            )?,
56            expansions,
57        })
58    }
59
60    #[cfg(test)]
61    pub(crate) fn new(
62        index: &super::index::LexicalIndex,
63        expansions: Vec<LexicalExpansion>,
64        fields: FieldWeights,
65    ) -> Result<Self, super::index::IndexError> {
66        let query = TermQuery {
67            terms: expansions.iter().map(|entry| entry.term.clone()).collect(),
68            fields,
69        };
70        Ok(Self {
71            scoring: super::search::PreparedTermQuery::from_owned(
72                index,
73                query,
74                super::bm25::Bm25Params::beir(),
75            )?,
76            expansions,
77        })
78    }
79
80    #[cfg(test)]
81    pub(crate) fn allocation_bytes(
82        expansions: &Vec<LexicalExpansion>,
83        fields: &FieldWeights,
84    ) -> Option<usize> {
85        let mut bytes = expansions
86            .capacity()
87            .checked_mul(std::mem::size_of::<LexicalExpansion>())?;
88        for entry in expansions {
89            #[cfg(test)]
90            WEIGHTED_SIZE_VISITS.with(|count| count.set(count.get() + 1));
91            bytes = bytes.checked_add(entry.term.capacity())?.checked_add(
92                super::search::PreparedTermQuery::single_allocation_bytes(
93                    &FieldWeights::flat(&[]),
94                    entry.term.len(),
95                )?,
96            )?;
97        }
98        bytes.checked_add(
99            fields
100                .iter()
101                .count()
102                .checked_mul(std::mem::size_of::<(FieldId, u32)>())?,
103        )
104    }
105
106    pub(crate) fn allocation_bytes_controlled<E>(
107        expansions: &Vec<LexicalExpansion>,
108        fields: &FieldWeights,
109        work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
110    ) -> Result<Option<usize>, E> {
111        work.check_now()?;
112        let Some(mut bytes) = expansions
113            .capacity()
114            .checked_mul(std::mem::size_of::<LexicalExpansion>())
115        else {
116            return Ok(None);
117        };
118        for entry in expansions {
119            work.step()?;
120            #[cfg(test)]
121            WEIGHTED_SIZE_VISITS.with(|count| count.set(count.get() + 1));
122            let Some(next) = bytes.checked_add(entry.term.capacity()).and_then(|bytes| {
123                bytes.checked_add(super::search::PreparedTermQuery::single_allocation_bytes(
124                    &FieldWeights::flat(&[]),
125                    entry.term.len(),
126                )?)
127            }) else {
128                return Ok(None);
129            };
130            bytes = next;
131        }
132        work.check_now()?;
133        Ok(fields
134            .allocation_bytes()
135            .and_then(|field_bytes| bytes.checked_add(field_bytes)))
136    }
137
138    pub(crate) fn expansions(&self) -> &[LexicalExpansion] {
139        &self.expansions
140    }
141    pub(crate) fn take_expansions(&mut self) -> Vec<LexicalExpansion> {
142        std::mem::take(&mut self.expansions)
143    }
144    pub(crate) fn scoring(&self) -> &super::search::PreparedTermQuery {
145        &self.scoring
146    }
147}
148
149/// A structured lexical query over already-analyzed terms.
150#[derive(Clone, Debug, Eq, PartialEq)]
151pub enum LexicalQuery {
152    /// Existing OR-of-terms semantics.
153    Term(TermQuery),
154    /// Terms must occur in positional phrase order within the declared slop.
155    Phrase {
156        /// Analyzed phrase terms in order.
157        terms: Vec<Vec<u8>>,
158        /// Maximum total positional displacement.
159        slop: u32,
160        /// Field to match.
161        field: FieldId,
162    },
163    /// Every indexed term beginning with this analyzed byte prefix.
164    Prefix {
165        /// Non-empty analyzed prefix.
166        prefix: Vec<u8>,
167        /// Field to match.
168        field: FieldId,
169    },
170    /// Leading exact terms plus a prefix on the trailing term. Used for type-ahead.
171    /// Store queries retain at most 32 prefix contributions by highest live
172    /// document frequency (ties by term bytes), reducing recall above that
173    /// limit. Exact contributions are unchanged; `Prefix` remains exhaustive.
174    TermsWithPrefix {
175        /// Analyzed leading terms matched exactly.
176        terms: Vec<Vec<u8>>,
177        /// Non-empty analyzed trailing prefix.
178        prefix: Vec<u8>,
179        /// Fields and weights shared by the exact and prefix legs.
180        fields: FieldWeights,
181    },
182    /// Every indexed term within a bounded edit distance.
183    Fuzzy {
184        /// Analyzed source term.
185        term: Vec<u8>,
186        /// Requested edit-distance ceiling, at most two.
187        max_distance: u32,
188        /// Field to match.
189        field: FieldId,
190    },
191    /// Terms sharing the pinned phonetic code.
192    Phonetic {
193        /// UTF-8 term whose primary phonetic code is used.
194        term: Vec<u8>,
195        /// Field to match.
196        field: FieldId,
197    },
198}
199
200impl LexicalQuery {
201    /// Wraps the compatibility term query.
202    #[must_use]
203    pub const fn term(query: TermQuery) -> Self {
204        Self::Term(query)
205    }
206
207    /// Builds a phrase query.
208    #[must_use]
209    pub const fn phrase(terms: Vec<Vec<u8>>, slop: u32, field: FieldId) -> Self {
210        Self::Phrase { terms, slop, field }
211    }
212
213    /// Builds a prefix query.
214    #[must_use]
215    pub const fn prefix(prefix: Vec<u8>, field: FieldId) -> Self {
216        Self::Prefix { prefix, field }
217    }
218
219    /// Builds a bounded fuzzy query.
220    #[must_use]
221    pub const fn fuzzy(term: Vec<u8>, max_distance: u32, field: FieldId) -> Self {
222        Self::Fuzzy {
223            term,
224            max_distance,
225            field,
226        }
227    }
228
229    /// Builds a phonetic query.
230    #[must_use]
231    pub const fn phonetic(term: Vec<u8>, field: FieldId) -> Self {
232        Self::Phonetic { term, field }
233    }
234
235    pub(crate) fn needs_vocabulary(&self) -> bool {
236        matches!(
237            self,
238            Self::Prefix { .. }
239                | Self::TermsWithPrefix { .. }
240                | Self::Fuzzy { .. }
241                | Self::Phonetic { .. }
242        )
243    }
244
245    pub(crate) fn fields_controlled<E>(
246        &self,
247        work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
248    ) -> Result<FieldWeights, E> {
249        work.check_now()?;
250        Ok(match self {
251            Self::Term(query) => return query.fields.clone_controlled(work),
252            Self::TermsWithPrefix { fields, .. } => return fields.clone_controlled(work),
253            Self::Phrase { field, .. }
254            | Self::Prefix { field, .. }
255            | Self::Fuzzy { field, .. }
256            | Self::Phonetic { field, .. } => FieldWeights::flat(&[*field]),
257        })
258    }
259
260    #[cfg(test)]
261    pub(crate) fn phrase_constraint(&self) -> Option<(&[Vec<u8>], u32)> {
262        match self {
263            Self::Phrase { terms, slop, .. } => Some((terms, *slop)),
264            _ => None,
265        }
266    }
267}
268
269/// Why one analyzed term entered a structured query.
270#[derive(Clone, Copy, Debug, Eq, PartialEq)]
271pub enum LexicalMatchKind {
272    /// The caller supplied this exact term.
273    Term,
274    /// The term is one member of a phrase.
275    Phrase,
276    /// The term expanded from a prefix.
277    Prefix,
278    /// The term expanded by edit distance.
279    Fuzzy {
280        /// Exact Wagner-Fischer edit distance.
281        distance: u32,
282    },
283    /// The term shares the query's pinned phonetic code.
284    Phonetic,
285}
286
287/// One reported scoring expansion.
288#[derive(Clone, Debug, Eq, PartialEq)]
289pub struct LexicalExpansion {
290    /// Indexed analyzed term.
291    pub term: Vec<u8>,
292    /// Multiplicative BM25 boost in thousandths.
293    pub boost_thousandths: u16,
294    /// Expansion provenance.
295    pub kind: LexicalMatchKind,
296}
297
298/// Owned snippet text with absolute source ranges.
299#[derive(Clone, Debug, Eq, PartialEq)]
300pub struct OwnedLexicalSnippet {
301    /// UTF-8-valid text copied from `source`.
302    pub text: String,
303    /// Absolute byte range within the persisted source text.
304    pub source: Highlight,
305    /// Absolute matched byte ranges within the persisted source text.
306    pub highlights: Vec<Highlight>,
307    /// Byte length of the whole persisted source text, so a caller can tell
308    /// whether `source` stops before its end.
309    pub source_len: usize,
310}
311
312/// A malformed structured lexical query.
313#[derive(Clone, Debug, Eq, PartialEq)]
314pub enum LexicalQueryError {
315    /// A term, phrase, or prefix that must be non-empty was empty.
316    Empty,
317    /// Fuzzy edit distance exceeded the pinned maximum.
318    FuzzyDistance {
319        /// Requested edit-distance ceiling.
320        requested: u32,
321        /// Engine maximum.
322        maximum: u32,
323    },
324    /// A phonetic query was not UTF-8.
325    InvalidPhoneticUtf8,
326    /// The phonetic query contained no encodable letters.
327    EmptyPhoneticCode,
328}
329
330impl std::fmt::Display for LexicalQueryError {
331    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332        match self {
333            Self::Empty => formatter.write_str("structured lexical query must not be empty"),
334            Self::FuzzyDistance { requested, maximum } => write!(
335                formatter,
336                "fuzzy distance {requested} exceeds maximum {maximum}"
337            ),
338            Self::InvalidPhoneticUtf8 => formatter.write_str("phonetic query must be UTF-8"),
339            Self::EmptyPhoneticCode => {
340                formatter.write_str("phonetic query has no encodable letters")
341            }
342        }
343    }
344}
345
346impl std::error::Error for LexicalQueryError {}
347
348#[cfg(test)]
349fn vocabulary<'a>(query: &LexicalQuery, terms: impl Iterator<Item = &'a [u8]>) -> Vocabulary {
350    if query.needs_vocabulary() {
351        terms.map(<[u8]>::to_vec).collect()
352    } else {
353        Vocabulary::empty()
354    }
355}
356
357#[cfg(test)]
358pub(crate) fn expand_with_phonetic<E: From<LexicalQueryError>>(
359    query: &LexicalQuery,
360    vocabulary: &Vocabulary,
361    phonetic_lookup: impl FnOnce(&str) -> Result<Vec<LexicalExpansion>, E>,
362) -> Result<Vec<LexicalExpansion>, E> {
363    let mut work = super::control::WorkCheck::new(|| Ok(()));
364    expand_with_phonetic_controlled(query, vocabulary, &mut work, phonetic_lookup)
365}
366
367pub(crate) fn expand_with_phonetic_controlled<E: From<LexicalQueryError>>(
368    query: &LexicalQuery,
369    vocabulary: &Vocabulary,
370    work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
371    phonetic_lookup: impl FnOnce(&str) -> Result<Vec<LexicalExpansion>, E>,
372) -> Result<Vec<LexicalExpansion>, E> {
373    work.check_now()?;
374    #[cfg(any(test, feature = "test-seams"))]
375    super::preparation_observer::expansion();
376    let expansions = match query {
377        LexicalQuery::Term(TermQuery { terms, .. }) | LexicalQuery::Phrase { terms, .. } => {
378            if terms.is_empty() {
379                return Err(LexicalQueryError::Empty.into());
380            }
381            // Preserve occurrence order and the validation of every term.
382            for term in terms {
383                work.step()?;
384                if term.is_empty() {
385                    return Err(LexicalQueryError::Empty.into());
386                }
387            }
388            let kind = if matches!(query, LexicalQuery::Term(_)) {
389                LexicalMatchKind::Term
390            } else {
391                LexicalMatchKind::Phrase
392            };
393            terms
394                .iter()
395                .map(|term| {
396                    work.step()?;
397                    Ok(LexicalExpansion {
398                        term: term.clone(),
399                        boost_thousandths: 1_000,
400                        kind,
401                    })
402                })
403                .collect::<Result<Vec<_>, E>>()?
404        }
405        LexicalQuery::Prefix { prefix, .. } => {
406            if prefix.is_empty() {
407                return Err(LexicalQueryError::Empty.into());
408            }
409            vocabulary
410                .prefix(prefix)
411                .map(|term| {
412                    work.step()?;
413                    Ok(LexicalExpansion {
414                        term: term.to_vec(),
415                        boost_thousandths: 1_000,
416                        kind: LexicalMatchKind::Prefix,
417                    })
418                })
419                .collect::<Result<Vec<_>, E>>()?
420        }
421        LexicalQuery::TermsWithPrefix { terms, prefix, .. } => {
422            if prefix.is_empty() {
423                return Err(LexicalQueryError::Empty.into());
424            }
425            let mut expansions = Vec::new();
426            for term in terms {
427                work.step()?;
428                if term.is_empty() {
429                    return Err(LexicalQueryError::Empty.into());
430                }
431                expansions.push(LexicalExpansion {
432                    term: term.clone(),
433                    boost_thousandths: 1_000,
434                    kind: LexicalMatchKind::Term,
435                });
436            }
437            let mut prefix_terms = Vec::new();
438            for term in vocabulary.prefix(prefix) {
439                work.step()?;
440                prefix_terms.push(term);
441            }
442            for length in 3..prefix.len() {
443                work.step()?;
444                let Some(candidate) = prefix.get(..length) else {
445                    continue;
446                };
447                if let Some(term) = vocabulary.exact(candidate) {
448                    prefix_terms.push(term);
449                }
450            }
451            let boost_thousandths = if prefix_terms.is_empty() {
452                1_000
453            } else {
454                u16::try_from((1_000 / prefix_terms.len()).max(1)).unwrap_or(1)
455            };
456            expansions.extend(prefix_terms.into_iter().map(|term| LexicalExpansion {
457                term: term.to_vec(),
458                boost_thousandths,
459                kind: LexicalMatchKind::Prefix,
460            }));
461            expansions
462        }
463        LexicalQuery::Fuzzy {
464            term, max_distance, ..
465        } => {
466            if term.is_empty() {
467                return Err(LexicalQueryError::Empty.into());
468            }
469            if *max_distance > fuzzy::MAX_EDIT_DISTANCE {
470                return Err(LexicalQueryError::FuzzyDistance {
471                    requested: *max_distance,
472                    maximum: fuzzy::MAX_EDIT_DISTANCE,
473                }
474                .into());
475            }
476            let mut scratch = fuzzy::BoundedDistance::new();
477            let mut matches = Vec::new();
478            for candidate in vocabulary.fuzzy_candidates(term.len(), *max_distance) {
479                work.step()?;
480                if let Some(distance) =
481                    scratch.distance_controlled(term, candidate, *max_distance, work)?
482                {
483                    matches.push(LexicalExpansion {
484                        term: candidate.to_vec(),
485                        boost_thousandths: match distance {
486                            0 => 1_000,
487                            1 => 500,
488                            _ => 250,
489                        },
490                        kind: LexicalMatchKind::Fuzzy { distance },
491                    });
492                }
493            }
494            super::control::sort_by(&mut matches, |left, right| left.term.cmp(&right.term), work)?;
495            matches
496        }
497        LexicalQuery::Phonetic { term, .. } => {
498            let term =
499                std::str::from_utf8(term).map_err(|_| LexicalQueryError::InvalidPhoneticUtf8)?;
500            let code = phonetic::encode_controlled(term, work)?;
501            if code.is_empty() {
502                return Err(LexicalQueryError::EmptyPhoneticCode.into());
503            }
504            phonetic_lookup(&code)?
505        }
506    };
507    work.check_now()?;
508    Ok(expansions)
509}
510
511#[cfg(test)]
512pub(crate) fn phonetic_expansions(
513    vocabulary: &Vocabulary,
514    index: &super::phonetic_index::PhoneticIndex,
515    code: &str,
516) -> Vec<LexicalExpansion> {
517    let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
518    match phonetic_expansions_controlled(vocabulary, index, code, &mut work) {
519        Ok(expansions) => expansions,
520        Err(never) => match never {},
521    }
522}
523
524pub(crate) fn phonetic_expansions_controlled<E>(
525    vocabulary: &Vocabulary,
526    index: &super::phonetic_index::PhoneticIndex,
527    code: &str,
528    work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
529) -> Result<Vec<LexicalExpansion>, E> {
530    work.check_now()?;
531    let expansions = vocabulary
532        .select(index.terms(code))
533        .map(|term| {
534            work.step()?;
535            Ok(LexicalExpansion {
536                term: {
537                    #[cfg(test)]
538                    PHONETIC_COPIES.with(|count| count.set(count.get() + 1));
539                    term.to_vec()
540                },
541                boost_thousandths: 250,
542                kind: LexicalMatchKind::Phonetic,
543            })
544        })
545        .collect::<Result<Vec<_>, E>>()?;
546    work.check_now()?;
547    Ok(expansions)
548}
549
550#[cfg(test)]
551thread_local! { static PHONETIC_COPIES:std::cell::Cell<usize>=const {std::cell::Cell::new(0)}; }
552
553#[cfg(test)]
554fn expand(
555    query: &LexicalQuery,
556    vocabulary: &Vocabulary,
557) -> Result<Vec<LexicalExpansion>, LexicalQueryError> {
558    // Existing small query tests retain the independent full encoder-scan control.
559    expand_with_phonetic(query, vocabulary, |code| {
560        Ok(vocabulary
561            .iter()
562            .filter(|term| {
563                std::str::from_utf8(term).is_ok_and(|term| phonetic::encode(term) == code)
564            })
565            .map(|term| LexicalExpansion {
566                term: term.to_vec(),
567                boost_thousandths: 250,
568                kind: LexicalMatchKind::Phonetic,
569            })
570            .collect())
571    })
572}
573
574#[cfg(test)]
575pub(crate) fn phrase_matches(
576    analyzer: &Analyzer,
577    text: &str,
578    terms: &[Vec<u8>],
579    slop: u32,
580) -> bool {
581    #[cfg(any(test, feature = "test-seams"))]
582    super::preparation_observer::phrase_reanalysis(text.len());
583    let analyzed = analyzer.analyze(text);
584    let streams = terms
585        .iter()
586        .map(|term| {
587            analyzed
588                .iter()
589                .filter(|token| token.term.as_bytes() == term)
590                .map(|token| token.position)
591                .collect::<Vec<_>>()
592        })
593        .collect::<Vec<_>>();
594    phrase::streams_match(&streams, slop)
595}
596
597#[cfg(test)]
598fn wagner_fischer(left: &[u8], right: &[u8]) -> u32 {
599    #[cfg(any(test, feature = "test-seams"))]
600    let (mut cells, mut allocations) = (0, 1);
601    let mut previous = (0..=right.len())
602        .map(|value| u32::try_from(value).unwrap_or(u32::MAX))
603        .collect::<Vec<_>>();
604    for (row, left_byte) in left.iter().enumerate() {
605        let mut current = Vec::with_capacity(right.len() + 1);
606        #[cfg(any(test, feature = "test-seams"))]
607        {
608            allocations += 1;
609        }
610        current.push(u32::try_from(row + 1).unwrap_or(u32::MAX));
611        for (column, right_byte) in right.iter().enumerate() {
612            #[cfg(any(test, feature = "test-seams"))]
613            {
614                cells += 1;
615            }
616            let (Some(up), Some(left), Some(diagonal)) = (
617                previous.get(column + 1).copied(),
618                current.get(column).copied(),
619                previous.get(column).copied(),
620            ) else {
621                return u32::MAX;
622            };
623            let deletion = up.saturating_add(1);
624            let insertion = left.saturating_add(1);
625            let substitution = diagonal.saturating_add(u32::from(left_byte != right_byte));
626            current.push(deletion.min(insertion).min(substitution));
627        }
628        previous = current;
629    }
630    #[cfg(any(test, feature = "test-seams"))]
631    super::preparation_observer::fuzzy_distance(cells, allocations);
632    previous.last().copied().unwrap_or(u32::MAX)
633}
634
635#[cfg(test)]
636#[allow(clippy::expect_used)]
637mod tests {
638    use super::*;
639
640    #[test]
641    fn astra_18_weighted_query_size_cancels_inside_walk() {
642        let expansions = (0..4_096)
643            .map(|_| LexicalExpansion {
644                term: b"alpha".to_vec(),
645                boost_thousandths: 1_000,
646                kind: LexicalMatchKind::Term,
647            })
648            .collect::<Vec<_>>();
649        let fields = FieldWeights::flat(&[field()]);
650        let expected = PreparedWeightedQuery::allocation_bytes(&expansions, &fields)
651            .expect("fixed allocation fits");
652        WEIGHTED_SIZE_VISITS.with(|count| count.set(0));
653        let mut work = super::super::control::WorkCheck::new(|| {
654            if WEIGHTED_SIZE_VISITS.with(std::cell::Cell::get) >= 64 {
655                Err("cancelled")
656            } else {
657                Ok(())
658            }
659        });
660        let result =
661            PreparedWeightedQuery::allocation_bytes_controlled(&expansions, &fields, &mut work);
662        let visited = WEIGHTED_SIZE_VISITS.with(std::cell::Cell::get);
663        println!("weighted size visits={visited}, clean bytes={expected}");
664        assert_eq!(result, Err("cancelled"));
665        assert!((64..=128).contains(&visited));
666        let mut clean = super::super::control::WorkCheck::new(|| Ok::<(), ()>(()));
667        assert_eq!(
668            PreparedWeightedQuery::allocation_bytes_controlled(&expansions, &fields, &mut clean),
669            Ok(Some(expected))
670        );
671    }
672
673    #[test]
674    fn astra_18_phonetic_selection_cancels_inside_one_bucket() {
675        let dictionary = (0..4_096)
676            .map(|i| format!("night{i:04}").into_bytes())
677            .collect::<Vocabulary>();
678        let index =
679            super::super::phonetic_index::PhoneticIndex::build(
680                &dictionary,
681                |_, _| Ok::<(), ()>(()),
682            )
683            .expect("index");
684        PHONETIC_COPIES.with(|count| count.set(0));
685        let mut work = super::super::control::WorkCheck::new(|| {
686            if PHONETIC_COPIES.with(std::cell::Cell::get) >= 64 {
687                Err("cancelled")
688            } else {
689                Ok(())
690            }
691        });
692        let result = phonetic_expansions_controlled(&dictionary, &index, "NT", &mut work);
693        let copies = PHONETIC_COPIES.with(std::cell::Cell::get);
694        println!("copied_terms={copies}");
695        assert!(matches!(result, Err("cancelled")));
696        assert!(
697            copies <= 128,
698            "selection must not copy the full bucket after cancellation"
699        );
700        let mut work = super::super::control::WorkCheck::new(|| Ok::<(), ()>(()));
701        let clean = phonetic_expansions_controlled(&dictionary, &index, "NT", &mut work)
702            .expect("clean selection");
703        assert_eq!(
704            clean.iter().map(|e| e.term.as_slice()).collect::<Vec<_>>(),
705            dictionary.iter().collect::<Vec<_>>()
706        );
707    }
708
709    #[test]
710    fn astra_18_expansion_cancels_inside_prefix_and_rejected_fuzzy_candidates() {
711        #[derive(Debug, PartialEq)]
712        enum Error {
713            Shape(LexicalQueryError),
714            Cancelled,
715        }
716        impl From<LexicalQueryError> for Error {
717            fn from(error: LexicalQueryError) -> Self {
718                Self::Shape(error)
719            }
720        }
721        let dictionary = (0..4_096)
722            .map(|i| format!("zz{i:06}").into_bytes())
723            .collect::<Vocabulary>();
724        let queries = [
725            LexicalQuery::prefix(b"zz".to_vec(), field()),
726            LexicalQuery::fuzzy(b"aaaaaaaa".to_vec(), 1, field()),
727        ];
728        let mut results = Vec::new();
729        for (slot, query) in queries.iter().enumerate() {
730            let clean = expand(query, &dictionary).expect("clean expansion");
731            assert_eq!(clean.len(), if slot == 0 { 4_096 } else { 0 });
732            super::super::preparation_observer::begin();
733            let mut checks = 0;
734            let mut work = super::super::control::WorkCheck::new(|| {
735                checks += 1;
736                if checks == 3 {
737                    Err(Error::Cancelled)
738                } else {
739                    Ok(())
740                }
741            });
742            let result =
743                expand_with_phonetic_controlled(query, &dictionary, &mut work, |_| Ok(vec![]));
744            let terms = if slot == 0 {
745                super::super::preparation_observer::vocabulary_work().terms_visited
746            } else {
747                super::super::preparation_observer::fuzzy_work().candidates
748            };
749            super::super::preparation_observer::take();
750            println!(
751                "operator={slot}, checks={checks}, examined={terms}, cancelled={}",
752                matches!(result, Err(Error::Cancelled))
753            );
754            results.push((matches!(result, Err(Error::Cancelled)), terms));
755            let mut work = super::super::control::WorkCheck::new(|| Ok::<(), Error>(()));
756            assert_eq!(
757                expand_with_phonetic_controlled(query, &dictionary, &mut work, |_| Ok(vec![]))
758                    .expect("controlled clean"),
759                clean
760            );
761        }
762        for (cancelled, terms) in results {
763            assert!(cancelled, "must stop even when no candidate matches");
764            assert!(terms > 0 && terms <= 128);
765        }
766    }
767
768    #[test]
769    fn astra_14_phonetic_expansion_cost_screen() {
770        use std::time::Instant;
771        let mut terms = vec![b"night".to_vec(), b"knight".to_vec()];
772        for mut n in 0..4096 {
773            let mut term = if n < 2048 {
774                b"smithson".to_vec()
775            } else {
776                b"robert".to_vec()
777            };
778            for _ in 0..4 {
779                term.push(b'a' + (n % 26) as u8);
780                n /= 26;
781            }
782            terms.push(term);
783        }
784        let dictionary = terms.into_iter().collect::<Vocabulary>();
785        let index = super::super::phonetic_index::PhoneticIndex::build(&dictionary, |_, _| {
786            Ok::<_, std::convert::Infallible>(())
787        })
788        .expect("cached map");
789        let cached_expand = |query: &LexicalQuery| {
790            expand_with_phonetic::<LexicalQueryError>(query, &dictionary, |code| {
791                Ok(phonetic_expansions(&dictionary, &index, code))
792            })
793        };
794        for (name, term, count) in [
795            ("rare", b"night".as_slice(), 2),
796            ("collision", b"smithson".as_slice(), 2048),
797        ] {
798            let query = LexicalQuery::phonetic(term.to_vec(), field());
799            for _ in 0..2 {
800                drop(cached_expand(&query).expect("warm"));
801            }
802            super::super::preparation_observer::begin();
803            let expected = cached_expand(&query).expect("probe");
804            let encodings = super::super::preparation_observer::phonetic_encoding_calls();
805            super::super::preparation_observer::take();
806            assert_eq!(expected.len(), count);
807            let mut samples = Vec::new();
808            for _ in 0..8 {
809                let start = Instant::now();
810                let actual = cached_expand(&query).expect("measure");
811                samples.push(start.elapsed().as_secs_f64() * 1e6);
812                assert_eq!(actual, expected);
813            }
814            samples.sort_by(f64::total_cmp);
815            println!(
816                "PHONETIC_SCREEN {name} vocabulary=4098 warmups=2 samples=8 p50_us={} p95_us={} expansions={} encodings={encodings}",
817                (samples.get(3).expect("sample") + samples.get(4).expect("sample")) / 2.0,
818                samples.last().expect("sample"),
819                expected.len()
820            );
821        }
822    }
823
824    #[test]
825    fn astra_13_fuzzy_scratch_is_reused_per_query() {
826        let query = LexicalQuery::fuzzy(b"abcdefgh".to_vec(), 2, field());
827        let terms = (b'a'..=b'z')
828            .map(|last| {
829                let mut term = b"abcdefg".to_vec();
830                term.push(last);
831                term
832            })
833            .collect::<Vec<_>>();
834        let dictionary = vocabulary(&query, terms.iter().map(Vec::as_slice));
835        for _ in 0..2 {
836            super::super::preparation_observer::begin();
837            assert_eq!(expand(&query, &dictionary).expect("query").len(), 26);
838            let work = super::super::preparation_observer::fuzzy_work();
839            super::super::preparation_observer::take();
840            assert_eq!(work.scratch_constructions, 1, "one per query: {work:?}");
841            assert_eq!(work.row_allocations, 0, "stack rows: {work:?}");
842            assert_eq!(work.candidates, 26);
843            assert!(work.dp_cells < 26 * 8 * 8, "banded cells: {work:?}");
844        }
845    }
846
847    #[test]
848    fn astra_13_fuzzy_bounded_matches_full_distance_expansions() {
849        use rand::RngCore;
850        use std::collections::{BTreeMap, BTreeSet};
851        let mut tiny = vec![Vec::new()];
852        for len in 1..=4 {
853            for bits in 0..(1 << len) {
854                tiny.push((0..len).map(|bit| b'a' + ((bits >> bit) & 1)).collect());
855            }
856        }
857        let dictionary = tiny.iter().cloned().collect::<Vocabulary>();
858        // Primary oracle enumerates actual edit operations, with no DP recurrence.
859        for term in tiny.iter().filter(|term| !term.is_empty()) {
860            let mut distances = BTreeMap::from([(term.clone(), 0_u32)]);
861            let mut frontier = BTreeSet::from([term.clone()]);
862            for distance in 0..=2 {
863                let query = LexicalQuery::fuzzy(term.clone(), distance, field());
864                let expected = dictionary
865                    .iter()
866                    .filter_map(|candidate| {
867                        distances.get(candidate).map(|&distance| LexicalExpansion {
868                            term: candidate.to_vec(),
869                            boost_thousandths: match distance {
870                                0 => 1000,
871                                1 => 500,
872                                _ => 250,
873                            },
874                            kind: LexicalMatchKind::Fuzzy { distance },
875                        })
876                    })
877                    .collect::<Vec<_>>();
878                assert_eq!(
879                    expand(&query, &dictionary).expect("tiny query"),
880                    expected,
881                    "{query:?}"
882                );
883                if distance == 2 {
884                    break;
885                }
886                let mut next = BTreeSet::new();
887                for word in frontier {
888                    for at in 0..=word.len() {
889                        for byte in b"ab" {
890                            let mut inserted = word.clone();
891                            inserted.insert(at, *byte);
892                            next.insert(inserted);
893                            if at < word.len() {
894                                let mut replaced = word.clone();
895                                *replaced.get_mut(at).expect("within word") = *byte;
896                                next.insert(replaced);
897                            }
898                        }
899                        if at < word.len() {
900                            let mut removed = word.clone();
901                            removed.remove(at);
902                            next.insert(removed);
903                        }
904                    }
905                }
906                for word in &next {
907                    distances.entry(word.clone()).or_insert(distance + 1);
908                }
909                frontier = next;
910            }
911        }
912        // Secondary full-distance oracle covers arbitrary bytes, lengths and mutations.
913        let mut rng = crate::test_support::seeded_rng(
914            "astra_13_fuzzy_bounded_matches_full_distance_expansions",
915        );
916        for _ in 0..256 {
917            let len = 1 + (rng.next_u32() % 80) as usize;
918            let term = (0..len)
919                .map(|_| (rng.next_u32() % 256) as u8)
920                .collect::<Vec<_>>();
921            let mut terms = vec![term.clone(), Vec::new()];
922            for _ in 0..12 {
923                let mut candidate = term.clone();
924                for _ in 0..rng.next_u32() % 5 {
925                    let at = (rng.next_u32() as usize) % (candidate.len() + 1);
926                    match rng.next_u32() % 3 {
927                        0 => candidate.insert(at, rng.next_u32() as u8),
928                        1 if at < candidate.len() => {
929                            candidate.remove(at);
930                        }
931                        _ => {
932                            if let Some(byte) = candidate.get_mut(at) {
933                                *byte = rng.next_u32() as u8;
934                            }
935                        }
936                    }
937                }
938                terms.push(candidate);
939            }
940            let dictionary = terms.into_iter().collect::<Vocabulary>();
941            for maximum in 0..=2 {
942                let query = LexicalQuery::fuzzy(term.clone(), maximum, field());
943                let expected = dictionary
944                    .iter()
945                    .filter_map(|candidate| {
946                        let distance = wagner_fischer(&term, candidate);
947                        (distance <= maximum).then(|| LexicalExpansion {
948                            term: candidate.to_vec(),
949                            boost_thousandths: match distance {
950                                0 => 1000,
951                                1 => 500,
952                                _ => 250,
953                            },
954                            kind: LexicalMatchKind::Fuzzy { distance },
955                        })
956                    })
957                    .collect::<Vec<_>>();
958                assert_eq!(expand(&query, &dictionary).expect("seeded query"), expected);
959            }
960        }
961    }
962
963    #[test]
964    fn astra_13_fuzzy_distance_boundaries_preserve_boosts() {
965        let dictionary = [
966            b"ab".to_vec(),
967            b"ba".to_vec(),
968            b"xb".to_vec(),
969            b"zzab".to_vec(),
970            b"zzzab".to_vec(),
971        ]
972        .into_iter()
973        .collect::<Vocabulary>();
974        let query = LexicalQuery::fuzzy(b"ab".to_vec(), 2, field());
975        let actual = expand(&query, &dictionary).expect("boundary query");
976        assert_eq!(
977            actual
978                .iter()
979                .map(|e| (e.term.as_slice(), e.boost_thousandths))
980                .collect::<Vec<_>>(),
981            vec![
982                (b"ab".as_slice(), 1000),
983                (b"ba".as_slice(), 250),
984                (b"xb".as_slice(), 500),
985                (b"zzab".as_slice(), 250)
986            ]
987        );
988        let utf8 = [
989            b"e".to_vec(),
990            "é".as_bytes().to_vec(),
991            "ê".as_bytes().to_vec(),
992        ]
993        .into_iter()
994        .collect::<Vocabulary>();
995        let query = LexicalQuery::fuzzy("é".as_bytes().to_vec(), 1, FieldId(99));
996        let actual = expand(&query, &utf8).expect("byte distance, global fields");
997        assert_eq!(
998            actual
999                .iter()
1000                .map(|e| (e.term.as_slice(), e.boost_thousandths))
1001                .collect::<Vec<_>>(),
1002            vec![("é".as_bytes(), 1000), ("ê".as_bytes(), 500)]
1003        );
1004        assert_eq!(
1005            expand(&LexicalQuery::fuzzy(Vec::new(), 2, field()), &utf8),
1006            Err(LexicalQueryError::Empty)
1007        );
1008        assert_eq!(
1009            expand(&LexicalQuery::fuzzy(b"ab".to_vec(), 3, field()), &utf8),
1010            Err(LexicalQueryError::FuzzyDistance {
1011                requested: 3,
1012                maximum: 2
1013            })
1014        );
1015    }
1016
1017    #[test]
1018    fn astra_13_fuzzy_length_filter_avoids_impossible_dp_work() {
1019        let query = LexicalQuery::fuzzy(b"ab".to_vec(), 2, field());
1020        let terms = [vec![b'x'; 64], vec![b'y'; 128]];
1021        let dictionary = vocabulary(&query, terms.iter().map(Vec::as_slice));
1022        super::super::preparation_observer::begin();
1023        let actual = expand(&query, &dictionary).expect("valid fuzzy query");
1024        let work = super::super::preparation_observer::fuzzy_work();
1025        super::super::preparation_observer::take();
1026        assert!(actual.is_empty());
1027        assert_eq!(
1028            work.dp_cells, 0,
1029            "impossible byte lengths performed DP: {work:?}"
1030        );
1031        assert_eq!(work.row_allocations, 0);
1032        assert_eq!(work.candidates, 0);
1033    }
1034
1035    #[test]
1036    fn astra_13_fuzzy_expansion_cost_screen() {
1037        use std::time::Instant;
1038        let terms = (0..4096)
1039            .map(|mut n| {
1040                let mut term = vec![b'a'; 60];
1041                for _ in 0..4 {
1042                    term.push(b'a' + (n % 26) as u8);
1043                    n /= 26;
1044                }
1045                term
1046            })
1047            .collect::<Vec<_>>();
1048        let seed_query = LexicalQuery::fuzzy(vec![b'a'; 64], 2, field());
1049        let dictionary = vocabulary(&seed_query, terms.iter().map(Vec::as_slice));
1050        for (name, term) in [
1051            ("length_mismatch", vec![b'z'; 4]),
1052            ("same_length_far", vec![b'z'; 64]),
1053            ("same_length_near", vec![b'a'; 64]),
1054        ] {
1055            let query = LexicalQuery::fuzzy(term, 2, field());
1056            for _ in 0..2 {
1057                drop(expand(&query, &dictionary).expect("warm"));
1058            }
1059            super::super::preparation_observer::begin();
1060            let expected = expand(&query, &dictionary).expect("work probe");
1061            let work = super::super::preparation_observer::fuzzy_work();
1062            super::super::preparation_observer::take();
1063            if name != "same_length_near" {
1064                assert!(expected.is_empty());
1065            }
1066            let mut samples = Vec::new();
1067            for _ in 0..8 {
1068                let start = Instant::now();
1069                let actual = expand(&query, &dictionary).expect("timed expansion");
1070                samples.push(start.elapsed().as_secs_f64() * 1e6);
1071                assert_eq!(actual, expected);
1072            }
1073            samples.sort_by(f64::total_cmp);
1074            println!(
1075                "FUZZY_SCREEN {name} vocabulary=4096 warmups=2 samples=8 p50_us={} p95_us={} expansions={} work={work:?}",
1076                (samples.get(3).expect("sample") + samples.get(4).expect("sample")) / 2.0,
1077                samples.last().expect("sample"),
1078                expected.len()
1079            );
1080        }
1081    }
1082
1083    #[test]
1084    fn astra_12_vocabulary_groups_memberships_with_linear_work() {
1085        let terms = (0..10_000)
1086            .map(|n| format!("unique{n:05}").into_bytes())
1087            .collect::<Vec<_>>();
1088        super::super::preparation_observer::begin();
1089        let dictionary = vocabulary(
1090            &LexicalQuery::prefix(b"unique".to_vec(), field()),
1091            terms.iter().map(Vec::as_slice),
1092        );
1093        let checks = super::super::preparation_observer::vocabulary_group_checks();
1094        super::super::preparation_observer::take();
1095        assert_eq!(dictionary.iter().count(), terms.len());
1096        assert!(
1097            checks <= 2 * terms.len(),
1098            "group boundary comparisons: {checks}"
1099        );
1100    }
1101
1102    #[test]
1103    fn astra_12_prefix_visits_only_matching_vocabulary_range() {
1104        let mut terms = (0..20_000)
1105            .map(|n| format!("{}{n:05}", if n < 10_000 { "aaa" } else { "zzz" }).into_bytes())
1106            .collect::<Vec<_>>();
1107        terms.extend([
1108            b"middlea".to_vec(),
1109            b"middleb".to_vec(),
1110            b"middlec".to_vec(),
1111        ]);
1112        let query = LexicalQuery::prefix(b"middle".to_vec(), field());
1113        let dictionary = vocabulary(&query, terms.iter().map(Vec::as_slice));
1114        super::super::preparation_observer::begin();
1115        let actual = expand(&query, &dictionary).expect("expand");
1116        let work = super::super::preparation_observer::vocabulary_work();
1117        super::super::preparation_observer::take();
1118        let expected = [b"middlea", b"middleb", b"middlec"]
1119            .into_iter()
1120            .map(|term| LexicalExpansion {
1121                term: term.to_vec(),
1122                boost_thousandths: 1000,
1123                kind: LexicalMatchKind::Prefix,
1124            })
1125            .collect::<Vec<_>>();
1126        assert_eq!(actual, expected);
1127        assert!(
1128            work.terms_visited <= 4,
1129            "prefix walked unrelated vocabulary: {work:?}"
1130        );
1131        assert!(work.seek_steps <= 16, "lower-bound work: {work:?}");
1132    }
1133
1134    fn field() -> FieldId {
1135        FieldId(7)
1136    }
1137
1138    #[test]
1139    fn astra_12_prefix_field_membership_and_byte_edges_preserve_expansion_order() {
1140        use std::collections::{BTreeMap, BTreeSet};
1141        let source = [
1142            (b"alpha".as_slice(), FieldId(7)),
1143            (b"alpine", FieldId(8)),
1144            (b"alpha", FieldId(8)),
1145            (b"alpha", FieldId(7)),
1146            (&[0], FieldId(3)),
1147            (&[0, 255], FieldId(4)),
1148            (&[255], FieldId(7)),
1149            (&[255, 255], FieldId(8)),
1150            (&[255, 255, 0], FieldId(7)),
1151            ("éclair".as_bytes(), FieldId(8)),
1152        ];
1153        let mut oracle = BTreeMap::<Vec<u8>, BTreeSet<FieldId>>::new();
1154        for (term, field) in source {
1155            oracle.entry(term.to_vec()).or_default().insert(field);
1156        }
1157        let dictionary =
1158            Vocabulary::build(source.into_iter(), |_, _| Ok::<_, ()>(())).expect("dictionary");
1159        assert_eq!(
1160            dictionary
1161                .entries()
1162                .map(|(term, fields)| (term.to_vec(), fields.to_vec()))
1163                .collect::<Vec<_>>(),
1164            oracle
1165                .iter()
1166                .map(|(term, fields)| (term.clone(), fields.iter().copied().collect()))
1167                .collect::<Vec<_>>()
1168        );
1169        for prefix in [
1170            b"al".as_slice(),
1171            &[0],
1172            &[255],
1173            &[255, 255],
1174            &[255, 255, 255],
1175            &[0xc3],
1176            b"missing",
1177        ] {
1178            for field in [FieldId(7), FieldId(8), FieldId(99)] {
1179                // Vocabulary visibility is global across fields, just as before;
1180                // field restrictions belong to scoring, not expansion policy.
1181                let query = LexicalQuery::prefix(prefix.to_vec(), field);
1182                let expected = oracle
1183                    .keys()
1184                    .filter(|term| term.starts_with(prefix))
1185                    .map(|term| LexicalExpansion {
1186                        term: term.clone(),
1187                        boost_thousandths: 1000,
1188                        kind: LexicalMatchKind::Prefix,
1189                    })
1190                    .collect::<Vec<_>>();
1191                assert_eq!(
1192                    expand(&query, &dictionary).expect("expand arbitrary bytes"),
1193                    expected
1194                );
1195            }
1196        }
1197    }
1198
1199    #[cfg(feature = "allocation-audit")]
1200    #[test]
1201    fn term_query_expansion_builds_no_vocabulary() {
1202        let source_terms = [
1203            b"cat".to_vec(),
1204            b"cats".to_vec(),
1205            b"cut".to_vec(),
1206            b"dog".to_vec(),
1207            b"night".to_vec(),
1208            b"nite".to_vec(),
1209        ];
1210        let query = LexicalQuery::term(TermQuery::flat(vec![b"cat".to_vec()], &[field()]));
1211        let empty_vocabulary = Vocabulary::empty();
1212        drop(expand(&query, &empty_vocabulary).expect("warm term expansion"));
1213        let (_, baseline) =
1214            crate::allocation_audit::audit_engine_path(|| expand(&query, &empty_vocabulary));
1215
1216        let (expansions, report) = crate::allocation_audit::audit_engine_path(|| {
1217            let vocabulary = vocabulary(&query, source_terms.iter().map(Vec::as_slice));
1218            expand(&query, &vocabulary)
1219        });
1220
1221        assert_eq!(
1222            expansions.expect("term expansion"),
1223            vec![LexicalExpansion {
1224                term: b"cat".to_vec(),
1225                boost_thousandths: 1_000,
1226                kind: LexicalMatchKind::Term,
1227            }]
1228        );
1229        assert_eq!(
1230            report.allocations, baseline.allocations,
1231            "Term expansion built the index vocabulary"
1232        );
1233    }
1234
1235    #[test]
1236    fn terms_with_prefix_expands_exact_terms_before_deterministic_prefix_terms() {
1237        let query = LexicalQuery::TermsWithPrefix {
1238            terms: vec![b"meeting".to_vec(), b"notes".to_vec()],
1239            prefix: b"meetings".to_vec(),
1240            fields: FieldWeights::flat(&[field()]),
1241        };
1242        let dictionary = [
1243            b"mee".as_slice(),
1244            b"meet",
1245            b"meeting",
1246            b"meetings",
1247            b"meetup",
1248        ]
1249        .into_iter()
1250        .map(<[u8]>::to_vec)
1251        .collect::<Vocabulary>();
1252        let expected = vec![
1253            LexicalExpansion {
1254                term: b"meeting".to_vec(),
1255                boost_thousandths: 1_000,
1256                kind: LexicalMatchKind::Term,
1257            },
1258            LexicalExpansion {
1259                term: b"notes".to_vec(),
1260                boost_thousandths: 1_000,
1261                kind: LexicalMatchKind::Term,
1262            },
1263            LexicalExpansion {
1264                term: b"meetings".to_vec(),
1265                boost_thousandths: 250,
1266                kind: LexicalMatchKind::Prefix,
1267            },
1268            LexicalExpansion {
1269                term: b"mee".to_vec(),
1270                boost_thousandths: 250,
1271                kind: LexicalMatchKind::Prefix,
1272            },
1273            LexicalExpansion {
1274                term: b"meet".to_vec(),
1275                boost_thousandths: 250,
1276                kind: LexicalMatchKind::Prefix,
1277            },
1278            LexicalExpansion {
1279                term: b"meeting".to_vec(),
1280                boost_thousandths: 250,
1281                kind: LexicalMatchKind::Prefix,
1282            },
1283        ];
1284
1285        assert_eq!(expand(&query, &dictionary), Ok(expected.clone()));
1286        assert_eq!(expand(&query, &dictionary), Ok(expected));
1287    }
1288
1289    #[test]
1290    fn terms_with_prefix_rejects_an_empty_prefix_but_allows_no_leading_terms() {
1291        let dictionary = [b"meeting".to_vec()].into_iter().collect::<Vocabulary>();
1292        let empty_prefix = LexicalQuery::TermsWithPrefix {
1293            terms: Vec::new(),
1294            prefix: Vec::new(),
1295            fields: FieldWeights::flat(&[field()]),
1296        };
1297        assert_eq!(
1298            expand(&empty_prefix, &dictionary),
1299            Err(LexicalQueryError::Empty)
1300        );
1301
1302        let prefix_only = LexicalQuery::TermsWithPrefix {
1303            terms: Vec::new(),
1304            prefix: b"mee".to_vec(),
1305            fields: FieldWeights::flat(&[field()]),
1306        };
1307        assert_eq!(
1308            expand(&prefix_only, &dictionary),
1309            Ok(vec![LexicalExpansion {
1310                term: b"meeting".to_vec(),
1311                boost_thousandths: 1_000,
1312                kind: LexicalMatchKind::Prefix,
1313            }])
1314        );
1315    }
1316
1317    #[test]
1318    fn terms_with_prefix_requires_vocabulary_while_term_does_not() {
1319        let prefixed = LexicalQuery::TermsWithPrefix {
1320            terms: Vec::new(),
1321            prefix: b"mee".to_vec(),
1322            fields: FieldWeights::flat(&[field()]),
1323        };
1324        let term = LexicalQuery::term(TermQuery::flat(vec![b"meeting".to_vec()], &[field()]));
1325        assert!(prefixed.needs_vocabulary());
1326        assert!(!term.needs_vocabulary());
1327    }
1328
1329    #[test]
1330    fn terms_with_prefix_cancels_inside_the_prefix_walk() {
1331        #[derive(Debug, PartialEq)]
1332        enum Error {
1333            Shape(LexicalQueryError),
1334            Cancelled,
1335        }
1336        impl From<LexicalQueryError> for Error {
1337            fn from(error: LexicalQueryError) -> Self {
1338                Self::Shape(error)
1339            }
1340        }
1341
1342        let dictionary = (0..4_096)
1343            .map(|index| format!("mee{index:04}").into_bytes())
1344            .collect::<Vocabulary>();
1345        let query = LexicalQuery::TermsWithPrefix {
1346            terms: vec![b"notes".to_vec()],
1347            prefix: b"mee".to_vec(),
1348            fields: FieldWeights::flat(&[field()]),
1349        };
1350        let mut checks = 0;
1351        let mut work = super::super::control::WorkCheck::new(|| {
1352            checks += 1;
1353            if checks == 3 {
1354                Err(Error::Cancelled)
1355            } else {
1356                Ok(())
1357            }
1358        });
1359        let result =
1360            expand_with_phonetic_controlled(&query, &dictionary, &mut work, |_| Ok(vec![]));
1361        assert_eq!(result, Err(Error::Cancelled));
1362        assert_eq!(checks, 3);
1363    }
1364
1365    #[test]
1366    fn vocabulary_expansions_preserve_terms_and_order() {
1367        let terms = [
1368            b"cat".as_slice(),
1369            b"cats",
1370            b"cut",
1371            b"dog",
1372            b"nite",
1373            b"night",
1374        ];
1375        let prefix = LexicalQuery::prefix(b"ca".to_vec(), field());
1376        let prefix_vocabulary = vocabulary(&prefix, terms.iter().copied());
1377
1378        assert_eq!(
1379            expand(&prefix, &prefix_vocabulary).expect("prefix expansion"),
1380            vec![
1381                LexicalExpansion {
1382                    term: b"cat".to_vec(),
1383                    boost_thousandths: 1_000,
1384                    kind: LexicalMatchKind::Prefix,
1385                },
1386                LexicalExpansion {
1387                    term: b"cats".to_vec(),
1388                    boost_thousandths: 1_000,
1389                    kind: LexicalMatchKind::Prefix,
1390                },
1391            ]
1392        );
1393        let fuzzy = LexicalQuery::fuzzy(b"cat".to_vec(), 1, field());
1394        let fuzzy_vocabulary = vocabulary(&fuzzy, terms.iter().copied());
1395        assert_eq!(
1396            expand(&fuzzy, &fuzzy_vocabulary).expect("fuzzy expansion"),
1397            vec![
1398                LexicalExpansion {
1399                    term: b"cat".to_vec(),
1400                    boost_thousandths: 1_000,
1401                    kind: LexicalMatchKind::Fuzzy { distance: 0 },
1402                },
1403                LexicalExpansion {
1404                    term: b"cats".to_vec(),
1405                    boost_thousandths: 500,
1406                    kind: LexicalMatchKind::Fuzzy { distance: 1 },
1407                },
1408                LexicalExpansion {
1409                    term: b"cut".to_vec(),
1410                    boost_thousandths: 500,
1411                    kind: LexicalMatchKind::Fuzzy { distance: 1 },
1412                },
1413            ]
1414        );
1415        let phonetic = LexicalQuery::phonetic(b"night".to_vec(), field());
1416        let phonetic_vocabulary = vocabulary(&phonetic, terms.iter().copied());
1417        assert_eq!(
1418            expand(&phonetic, &phonetic_vocabulary).expect("phonetic expansion"),
1419            vec![
1420                LexicalExpansion {
1421                    term: b"night".to_vec(),
1422                    boost_thousandths: 250,
1423                    kind: LexicalMatchKind::Phonetic,
1424                },
1425                LexicalExpansion {
1426                    term: b"nite".to_vec(),
1427                    boost_thousandths: 250,
1428                    kind: LexicalMatchKind::Phonetic,
1429                },
1430            ]
1431        );
1432    }
1433
1434    #[test]
1435    fn structured_query_validation_and_expansion_are_exhaustive() {
1436        let vocabulary = [
1437            b"cat".as_slice(),
1438            b"cats",
1439            b"cut",
1440            b"dog",
1441            b"nite",
1442            b"night",
1443        ]
1444        .into_iter()
1445        .map(<[u8]>::to_vec)
1446        .collect();
1447        let term = LexicalQuery::term(TermQuery::flat(vec![b"cat".to_vec()], &[field()]));
1448        assert_eq!(
1449            term.fields_controlled(&mut super::super::control::WorkCheck::new(|| Ok::<(), ()>(
1450                ()
1451            ))),
1452            Ok(FieldWeights::flat(&[field()]))
1453        );
1454        assert!(term.phrase_constraint().is_none());
1455        assert_eq!(expand(&term, &vocabulary).expect("term expansion").len(), 1);
1456
1457        let phrase = LexicalQuery::phrase(vec![b"cat".to_vec(), b"dog".to_vec()], 1, field());
1458        assert_eq!(
1459            phrase.phrase_constraint(),
1460            Some((&[b"cat".to_vec(), b"dog".to_vec()][..], 1))
1461        );
1462        assert_eq!(
1463            expand(&phrase, &vocabulary)
1464                .expect("phrase expansion")
1465                .len(),
1466            2
1467        );
1468        assert_eq!(
1469            expand(&LexicalQuery::prefix(b"ca".to_vec(), field()), &vocabulary)
1470                .expect("prefix expansion")
1471                .len(),
1472            2
1473        );
1474        let fuzzy = expand(
1475            &LexicalQuery::fuzzy(b"cat".to_vec(), 1, field()),
1476            &vocabulary,
1477        )
1478        .expect("fuzzy expansion");
1479        assert!(fuzzy.iter().any(|entry| entry.boost_thousandths == 1_000));
1480        assert!(fuzzy.iter().any(|entry| entry.boost_thousandths == 500));
1481        assert!(
1482            !expand(
1483                &LexicalQuery::phonetic(b"night".to_vec(), field()),
1484                &vocabulary,
1485            )
1486            .expect("phonetic expansion")
1487            .is_empty()
1488        );
1489
1490        for query in [
1491            LexicalQuery::term(TermQuery::flat(Vec::new(), &[field()])),
1492            LexicalQuery::phrase(vec![Vec::new()], 0, field()),
1493            LexicalQuery::prefix(Vec::new(), field()),
1494            LexicalQuery::fuzzy(Vec::new(), 1, field()),
1495        ] {
1496            assert_eq!(expand(&query, &vocabulary), Err(LexicalQueryError::Empty));
1497        }
1498        let distance = expand(
1499            &LexicalQuery::fuzzy(b"cat".to_vec(), 3, field()),
1500            &vocabulary,
1501        )
1502        .expect_err("distance ceiling");
1503        assert!(distance.to_string().contains("exceeds maximum"));
1504        let invalid = expand(&LexicalQuery::phonetic(vec![0xff], field()), &vocabulary)
1505            .expect_err("invalid phonetic UTF-8");
1506        assert_eq!(invalid, LexicalQueryError::InvalidPhoneticUtf8);
1507        assert_eq!(invalid.to_string(), "phonetic query must be UTF-8");
1508        let empty_code = expand(
1509            &LexicalQuery::phonetic(b"123".to_vec(), field()),
1510            &vocabulary,
1511        )
1512        .expect_err("empty phonetic code");
1513        assert_eq!(empty_code, LexicalQueryError::EmptyPhoneticCode);
1514        assert_eq!(
1515            LexicalQueryError::Empty.to_string(),
1516            "structured lexical query must not be empty"
1517        );
1518        assert_eq!(wagner_fischer(b"kitten", b"sitting"), 3);
1519    }
1520}