1use std::collections::HashSet;
2
3use crate::{KeyKind, LanguageBackend, SearchConfig};
4use rayon::prelude::*;
5
6#[cfg(not(test))]
7const PARALLEL_INDEX_THRESHOLD: usize = 50_000;
8#[cfg(test)]
9const PARALLEL_INDEX_THRESHOLD: usize = 4;
10
11#[derive(Clone, Copy, Debug, Eq, PartialEq)]
12pub struct SourceSpan {
13 pub start: usize,
14 pub end: usize,
15}
16
17#[derive(Clone, Debug, Eq, PartialEq)]
18pub struct Candidate {
19 pub id: usize,
20 pub display: String,
21 pub keys: Vec<SearchKey>,
22}
23
24#[derive(Clone, Debug, Eq, PartialEq)]
25pub struct SearchKey {
26 pub text: String,
27 pub kind: KeyKind,
28 pub weight: i32,
29 pub source_map: Option<Box<[Option<SourceSpan>]>>,
30}
31
32impl SearchKey {
33 pub fn with_source_map(mut self, source_map: Vec<Option<SourceSpan>>) -> Self {
34 self.source_map = Some(source_map.into_boxed_slice());
35 self
36 }
37
38 pub fn original(text: impl Into<String>) -> Self {
39 Self {
40 text: text.into(),
41 kind: KeyKind::Original,
42 weight: 3000,
43 source_map: None,
44 }
45 }
46
47 pub fn normalized(text: impl Into<String>) -> Self {
48 Self {
49 text: text.into(),
50 kind: KeyKind::Normalized,
51 weight: 2800,
52 source_map: None,
53 }
54 }
55
56 pub fn kana_reading(text: impl Into<String>) -> Self {
57 Self {
58 text: text.into(),
59 kind: KeyKind::KanaReading,
60 weight: 1700,
61 source_map: None,
62 }
63 }
64
65 pub fn romaji_reading(text: impl Into<String>) -> Self {
66 Self {
67 text: text.into(),
68 kind: KeyKind::RomajiReading,
69 weight: 1800,
70 source_map: None,
71 }
72 }
73
74 pub fn pinyin_full(text: impl Into<String>) -> Self {
75 Self {
76 text: text.into(),
77 kind: KeyKind::PinyinFull,
78 weight: 1750,
79 source_map: None,
80 }
81 }
82
83 pub fn pinyin_joined(text: impl Into<String>) -> Self {
84 Self {
85 text: text.into(),
86 kind: KeyKind::PinyinJoined,
87 weight: 1800,
88 source_map: None,
89 }
90 }
91
92 pub fn pinyin_initials(text: impl Into<String>) -> Self {
93 Self {
94 text: text.into(),
95 kind: KeyKind::PinyinInitials,
96 weight: 1850,
97 source_map: None,
98 }
99 }
100
101 pub fn learned_alias(text: impl Into<String>) -> Self {
102 Self {
103 text: text.into(),
104 kind: KeyKind::LearnedAlias,
105 weight: 2500,
106 source_map: None,
107 }
108 }
109}
110
111pub fn build_candidate(
112 id: usize,
113 display: impl Into<String>,
114 backend: &dyn LanguageBackend,
115 config: &SearchConfig,
116) -> Candidate {
117 let display = display.into();
118 let mut keys = vec![
119 SearchKey::original(display.clone()),
120 SearchKey::normalized(backend.normalize_candidate(&display)),
121 ];
122 keys.extend(backend.build_candidate_keys(&display));
123 let keys = dedup_and_limit_keys(keys, config);
124
125 Candidate { id, display, keys }
126}
127
128pub fn build_index<I, S>(
129 items: I,
130 backend: &dyn LanguageBackend,
131 config: &SearchConfig,
132) -> Vec<Candidate>
133where
134 I: IntoIterator<Item = S>,
135 S: Into<String>,
136{
137 let items: Vec<_> = items.into_iter().map(Into::into).collect();
138 if should_build_index_parallel(items.len()) {
139 return items
140 .into_par_iter()
141 .enumerate()
142 .map(|(id, item)| build_candidate(id, item, backend, config))
143 .collect();
144 }
145
146 items
147 .into_iter()
148 .enumerate()
149 .map(|(id, item)| build_candidate(id, item, backend, config))
150 .collect()
151}
152
153fn should_build_index_parallel(len: usize) -> bool {
154 len >= PARALLEL_INDEX_THRESHOLD && rayon::current_num_threads() > 1
155}
156
157pub fn dedup_and_limit_keys(keys: Vec<SearchKey>, config: &SearchConfig) -> Vec<SearchKey> {
158 let mut seen = HashSet::new();
159 let mut out = Vec::new();
160 let mut total_bytes = 0usize;
161
162 for key in keys {
163 if !seen.insert((key.kind, key.text.clone())) {
164 continue;
165 }
166
167 let required_base_key = matches!(key.kind, KeyKind::Original | KeyKind::Normalized);
168 let would_exceed_count = out.len() >= config.max_search_keys_per_candidate;
169 let would_exceed_bytes =
170 total_bytes + key.text.len() > config.max_total_key_bytes_per_candidate;
171
172 if !required_base_key && (would_exceed_count || would_exceed_bytes) {
173 continue;
174 }
175
176 total_bytes += key.text.len();
177 out.push(key);
178 }
179
180 out
181}
182
183#[cfg(test)]
184mod tests {
185 use crate::{query::PlainBackend, KeyKind};
186
187 use super::*;
188
189 #[test]
190 fn plain_mode_only_original_and_normalized() {
191 let cand = build_candidate(0, "東京駅", &PlainBackend, &SearchConfig::default());
192
193 assert!(cand.keys.iter().any(|k| k.kind == KeyKind::Original));
194 assert!(cand.keys.iter().any(|k| k.kind == KeyKind::Normalized));
195 assert!(!cand
196 .keys
197 .iter()
198 .any(|k| matches!(k.kind, KeyKind::KanaReading | KeyKind::RomajiReading)));
199 }
200
201 #[test]
202 fn original_key_is_always_present() {
203 let cand = build_candidate(0, "README.md", &PlainBackend, &SearchConfig::default());
204 assert!(cand.keys.iter().any(|k| k.kind == KeyKind::Original));
205 }
206
207 #[test]
208 fn search_keys_are_deduped_and_capped() {
209 let cfg = SearchConfig {
210 max_search_keys_per_candidate: 4,
211 ..SearchConfig::default()
212 };
213 let keys = vec![
214 SearchKey::original("a"),
215 SearchKey::normalized("a"),
216 SearchKey::normalized("a"),
217 SearchKey::learned_alias("b"),
218 SearchKey::learned_alias("c"),
219 SearchKey::learned_alias("d"),
220 ];
221
222 let out = dedup_and_limit_keys(keys, &cfg);
223
224 assert!(out.len() <= 4);
225 assert_eq!(
226 out.len(),
227 out.iter()
228 .map(|k| (k.kind, k.text.as_str()))
229 .collect::<HashSet<_>>()
230 .len()
231 );
232 }
233
234 #[test]
235 fn parallel_index_preserves_input_order_and_ids() {
236 let cfg = SearchConfig::default();
237 let cand = build_index(["one", "two", "three", "four"], &PlainBackend, &cfg);
238
239 assert_eq!(
240 cand.iter()
241 .map(|candidate| (candidate.id, candidate.display.as_str()))
242 .collect::<Vec<_>>(),
243 vec![(0, "one"), (1, "two"), (2, "three"), (3, "four")]
244 );
245 }
246}