zeppelin_embed/fts/
prefix.rs1use super::dict::{DictError, TermDictionary};
27
28#[derive(Clone, Debug, Eq, PartialEq)]
30pub struct PrefixMatch {
31 pub term_index: usize,
33 pub term: Vec<u8>,
35}
36
37pub fn search(dictionary: &TermDictionary, prefix: &[u8]) -> Result<Vec<PrefixMatch>, DictError> {
43 let range = dictionary.prefix_range(prefix)?;
44 Ok(range
45 .filter_map(|term_index| {
46 dictionary.term_at(term_index).map(|term| PrefixMatch {
47 term_index,
48 term: term.to_vec(),
49 })
50 })
51 .collect())
52}
53
54pub fn count(dictionary: &TermDictionary, prefix: &[u8]) -> Result<usize, DictError> {
60 Ok(dictionary.prefix_range(prefix)?.len())
61}
62
63#[cfg(test)]
64#[allow(
65 clippy::expect_used,
66 clippy::indexing_slicing,
67 clippy::panic,
68 clippy::unwrap_used
69)]
70mod tests {
71 use super::*;
72 use crate::fts::dict::TermInfo;
73 use crate::fts::tokenizer::{Analyzer, TokenizerConfig};
74
75 fn dictionary(terms: &[&str]) -> TermDictionary {
76 let mut sorted: Vec<&str> = terms.to_vec();
77 sorted.sort_unstable();
78 sorted.dedup();
79 let mut dictionary = TermDictionary::default();
80 for term in sorted {
81 dictionary
82 .push(term.as_bytes(), TermInfo::default())
83 .expect("sorted");
84 }
85 dictionary
86 }
87
88 fn matched(dictionary: &TermDictionary, prefix: &str) -> Vec<String> {
89 search(dictionary, prefix.as_bytes())
90 .expect("valid prefix")
91 .into_iter()
92 .filter_map(|entry| String::from_utf8(entry.term).ok())
93 .collect()
94 }
95
96 const VOCABULARY: [&str; 10] = [
97 "engine",
98 "engineer",
99 "engineering",
100 "engines",
101 "england",
102 "lexical",
103 "search",
104 "searching",
105 "seat",
106 "zebra",
107 ];
108
109 #[test]
110 fn prefix_results_equal_brute_force_over_the_dictionary() {
111 let dictionary = dictionary(&VOCABULARY);
112 for prefix in [
113 "e", "en", "eng", "engine", "s", "sea", "search", "z", "q", "engines",
114 ] {
115 let mut expected: Vec<String> = VOCABULARY
116 .iter()
117 .filter(|term| term.starts_with(prefix))
118 .map(|term| (*term).to_owned())
119 .collect();
120 expected.sort();
121 let mut actual = matched(&dictionary, prefix);
122 actual.sort();
123 assert_eq!(actual, expected, "prefix {prefix:?}");
124 assert_eq!(
125 count(&dictionary, prefix.as_bytes()).expect("valid"),
126 expected.len()
127 );
128 }
129 }
130
131 #[test]
132 fn empty_prefix_is_rejected_with_a_typed_error() {
133 let dictionary = dictionary(&VOCABULARY);
134 assert_eq!(search(&dictionary, b""), Err(DictError::EmptyTerm));
135 assert_eq!(count(&dictionary, b""), Err(DictError::EmptyTerm));
136 }
137
138 #[test]
139 fn a_prefix_matching_nothing_returns_an_empty_range_not_an_error() {
140 let dictionary = dictionary(&VOCABULARY);
141 assert!(matched(&dictionary, "qqq").is_empty());
142 assert!(matched(&dictionary, "zzzz").is_empty());
143 }
144
145 #[test]
146 fn a_term_is_a_prefix_of_itself() {
147 let dictionary = dictionary(&VOCABULARY);
148 assert!(matched(&dictionary, "zebra").contains(&String::from("zebra")));
149 }
150
151 #[test]
152 fn prefix_ending_mid_grapheme_matches_by_codepoint_rule() {
153 let analyzer = Analyzer::new(TokenizerConfig::text_default()).expect("valid config");
156 let folded: Vec<String> = analyzer
157 .analyze("cafe\u{0301}")
158 .into_iter()
159 .map(|token| token.term)
160 .collect();
161 let refs: Vec<&str> = folded.iter().map(String::as_str).collect();
162 let dictionary = dictionary(&refs);
163 assert!(!matched(&dictionary, "caf").is_empty());
165 assert!(
168 folded.iter().all(|term| !term.contains('\u{0301}')),
169 "a combining mark survived into the dictionary"
170 );
171 }
172
173 #[test]
174 fn a_prefix_can_match_a_term_whose_surface_form_it_does_not_prefix() {
175 let analyzer = Analyzer::new(TokenizerConfig::text_default()).expect("valid config");
179 let terms: Vec<String> = analyzer
180 .analyze("stra\u{00DF}e")
181 .into_iter()
182 .map(|token| token.term)
183 .collect();
184 let refs: Vec<&str> = terms.iter().map(String::as_str).collect();
185 let dictionary = dictionary(&refs);
186 assert!(!matched(&dictionary, "strass").is_empty());
187 }
188
189 #[test]
190 fn an_empty_dictionary_answers_every_prefix_with_nothing() {
191 let dictionary = TermDictionary::default();
192 assert!(search(&dictionary, b"anything").expect("valid").is_empty());
193 }
194
195 #[test]
196 fn matches_carry_the_index_the_dictionary_agrees_with() {
197 let dictionary = dictionary(&VOCABULARY);
198 for entry in search(&dictionary, b"engine").expect("valid") {
199 assert_eq!(
200 dictionary.term_at(entry.term_index),
201 Some(entry.term.as_slice())
202 );
203 }
204 }
205}