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