1use crate::fuzz;
4use crate::utils::full_process;
5use std::marker::PhantomData;
6
7#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9pub struct Match<'a> {
10 pub item: &'a str,
12 pub score: u8,
14}
15
16#[derive(Clone, Copy)]
18pub enum Scorer<'s> {
19 Ratio,
21 PartialRatio,
23 TokenSortRatio,
25 TokenSetRatio,
27 PartialTokenSortRatio,
29 PartialTokenSetRatio,
31 WRatio,
33 QRatio,
35 UwRatio,
37 UqRatio,
39 Custom(&'s dyn Fn(&str, &str) -> u8),
41}
42
43impl std::fmt::Debug for Scorer<'_> {
44 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45 let name = match self {
46 Self::Ratio => "Ratio",
47 Self::PartialRatio => "PartialRatio",
48 Self::TokenSortRatio => "TokenSortRatio",
49 Self::TokenSetRatio => "TokenSetRatio",
50 Self::PartialTokenSortRatio => "PartialTokenSortRatio",
51 Self::PartialTokenSetRatio => "PartialTokenSetRatio",
52 Self::WRatio => "WRatio",
53 Self::QRatio => "QRatio",
54 Self::UwRatio => "UwRatio",
55 Self::UqRatio => "UqRatio",
56 Self::Custom(_) => "Custom",
57 };
58 formatter.write_str(name)
59 }
60}
61
62impl Scorer<'_> {
63 fn force_ascii(self) -> bool {
64 !matches!(self, Self::UwRatio | Self::UqRatio)
65 }
66
67 fn score_processed_raw(self, a: &str, b: &str) -> f64 {
68 match self {
69 Self::Ratio | Self::QRatio | Self::UqRatio => fuzz::ratio_raw(a, b) * 100.0,
70 Self::PartialRatio => fuzz::partial_ratio_raw(a, b) * 100.0,
71 Self::TokenSortRatio => fuzz::token_sort_raw(a, b, false, false, false) * 100.0,
72 Self::TokenSetRatio => fuzz::token_set_raw(a, b, false, false, false) * 100.0,
73 Self::PartialTokenSortRatio => fuzz::token_sort_raw(a, b, false, false, true) * 100.0,
74 Self::PartialTokenSetRatio => fuzz::token_set_raw(a, b, false, false, true) * 100.0,
75 Self::WRatio | Self::UwRatio => fuzz::wratio_raw(a, b, false, false) * 100.0,
76 Self::Custom(score) => f64::from(score(a, b)),
77 }
78 }
79}
80
81pub fn extract<'a, I>(
84 query: &str,
85 choices: I,
86 scorer: Scorer<'_>,
87 limit: Option<usize>,
88 score_cutoff: u8,
89) -> Vec<Match<'a>>
90where
91 I: IntoIterator<Item = &'a str>,
92{
93 let query = full_process(query, scorer.force_ascii());
94 let mut matches = Vec::new();
95 for (index, item) in choices.into_iter().enumerate() {
96 let raw_score = score_choice_raw(&query, item, scorer);
97 if raw_score >= f64::from(score_cutoff) {
98 matches.push((
99 index,
100 raw_score,
101 Match {
102 item,
103 score: fuzz::round_score(raw_score),
104 },
105 ));
106 }
107 }
108 matches.sort_by(|left, right| {
109 right
110 .1
111 .total_cmp(&left.1)
112 .then_with(|| left.0.cmp(&right.0))
113 });
114 matches.truncate(limit.unwrap_or(usize::MAX));
115 matches.into_iter().map(|(_, _, item)| item).collect()
116}
117
118pub fn extract_one<'a, I>(
120 query: &str,
121 choices: I,
122 scorer: Scorer<'_>,
123 score_cutoff: u8,
124) -> Option<Match<'a>>
125where
126 I: IntoIterator<Item = &'a str>,
127{
128 let query = full_process(query, scorer.force_ascii());
129 choices
130 .into_iter()
131 .filter_map(|item| {
132 let raw_score = score_choice_raw(&query, item, scorer);
133 (raw_score >= f64::from(score_cutoff)).then_some((raw_score, item))
134 })
135 .fold(None, |best, (raw_score, item)| match best {
136 Some((best_score, _)) if best_score >= raw_score => best,
137 _ => Some((raw_score, item)),
138 })
139 .map(|(raw_score, item)| Match {
140 item,
141 score: fuzz::round_score(raw_score),
142 })
143}
144
145pub fn extract_without_order<'a, 's, I>(
147 query: &str,
148 choices: I,
149 scorer: Scorer<'s>,
150 score_cutoff: u8,
151) -> ExtractWithoutOrder<'a, 's, I::IntoIter>
152where
153 I: IntoIterator<Item = &'a str>,
154{
155 ExtractWithoutOrder {
156 choices: choices.into_iter(),
157 query: full_process(query, scorer.force_ascii()),
158 scorer,
159 score_cutoff,
160 item_lifetime: PhantomData,
161 }
162}
163
164pub struct ExtractWithoutOrder<'a, 's, I> {
166 choices: I,
167 query: String,
168 scorer: Scorer<'s>,
169 score_cutoff: u8,
170 item_lifetime: PhantomData<&'a str>,
171}
172
173impl<'a, I> Iterator for ExtractWithoutOrder<'a, '_, I>
174where
175 I: Iterator<Item = &'a str>,
176{
177 type Item = Match<'a>;
178
179 fn next(&mut self) -> Option<Self::Item> {
180 for item in self.choices.by_ref() {
181 let raw_score = score_choice_raw(&self.query, item, self.scorer);
182 if raw_score >= f64::from(self.score_cutoff) {
183 return Some(Match {
184 item,
185 score: fuzz::round_score(raw_score),
186 });
187 }
188 }
189 None
190 }
191}
192
193#[must_use]
195pub fn dedupe(items: &[&str], threshold: u8, scorer: Scorer<'_>) -> Vec<String> {
196 let mut representatives = Vec::with_capacity(items.len());
197 for &item in items {
198 let representative = items
199 .iter()
200 .copied()
201 .filter(|&candidate| score_pair(item, candidate, scorer) >= threshold)
202 .max_by(|left, right| {
203 left.chars()
204 .count()
205 .cmp(&right.chars().count())
206 .then_with(|| left.cmp(right))
207 })
208 .unwrap_or(item);
209 if !representatives.contains(&representative) {
210 representatives.push(representative);
211 }
212 }
213
214 if representatives.len() == items.len() {
215 items.iter().map(|item| (*item).to_owned()).collect()
216 } else {
217 representatives.into_iter().map(str::to_owned).collect()
218 }
219}
220
221pub fn extract_default<'a, I>(query: &str, choices: I) -> Vec<Match<'a>>
223where
224 I: IntoIterator<Item = &'a str>,
225{
226 extract(query, choices, Scorer::WRatio, Some(5), 0)
227}
228
229pub fn extract_one_default<'a, I>(query: &str, choices: I) -> Option<Match<'a>>
231where
232 I: IntoIterator<Item = &'a str>,
233{
234 extract_one(query, choices, Scorer::WRatio, 0)
235}
236
237pub fn extract_without_order_default<'a, I>(
239 query: &str,
240 choices: I,
241) -> ExtractWithoutOrder<'a, 'static, I::IntoIter>
242where
243 I: IntoIterator<Item = &'a str>,
244{
245 extract_without_order(query, choices, Scorer::WRatio, 0)
246}
247
248#[must_use]
250pub fn dedupe_default(items: &[&str]) -> Vec<String> {
251 dedupe(items, 70, Scorer::TokenSetRatio)
252}
253
254fn score_pair(a: &str, b: &str, scorer: Scorer<'_>) -> u8 {
255 let force_ascii = scorer.force_ascii();
256 let left = full_process(a, force_ascii);
257 let right = full_process(b, force_ascii);
258 if left.is_empty() {
259 0
260 } else {
261 fuzz::round_score(scorer.score_processed_raw(&left, &right))
262 }
263}
264
265fn score_choice_raw(query: &str, item: &str, scorer: Scorer<'_>) -> f64 {
266 if query.is_empty() {
267 0.0
268 } else {
269 let processed = full_process(item, scorer.force_ascii());
270 scorer.score_processed_raw(query, &processed)
271 }
272}