yuru_core/
candidate.rs

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 korean_romanized(text: impl Into<String>) -> Self {
102        Self {
103            text: text.into(),
104            kind: KeyKind::KoreanRomanized,
105            weight: 1800,
106            source_map: None,
107        }
108    }
109
110    pub fn korean_initials(text: impl Into<String>) -> Self {
111        Self {
112            text: text.into(),
113            kind: KeyKind::KoreanInitials,
114            weight: 1850,
115            source_map: None,
116        }
117    }
118
119    pub fn korean_keyboard(text: impl Into<String>) -> Self {
120        Self {
121            text: text.into(),
122            kind: KeyKind::KoreanKeyboard,
123            weight: 1750,
124            source_map: None,
125        }
126    }
127
128    pub fn learned_alias(text: impl Into<String>) -> Self {
129        Self {
130            text: text.into(),
131            kind: KeyKind::LearnedAlias,
132            weight: 2500,
133            source_map: None,
134        }
135    }
136}
137
138pub fn build_candidate(
139    id: usize,
140    display: impl Into<String>,
141    backend: &dyn LanguageBackend,
142    config: &SearchConfig,
143) -> Candidate {
144    let display = display.into();
145    let mut keys = vec![SearchKey::original(display.clone())];
146    if config.normalize {
147        keys.push(SearchKey::normalized(backend.normalize_candidate(&display)));
148    }
149    keys.extend(backend.build_candidate_keys(&display));
150    let keys = dedup_and_limit_keys(keys, config);
151
152    Candidate { id, display, keys }
153}
154
155pub fn build_index<I, S>(
156    items: I,
157    backend: &dyn LanguageBackend,
158    config: &SearchConfig,
159) -> Vec<Candidate>
160where
161    I: IntoIterator<Item = S>,
162    S: Into<String>,
163{
164    let items: Vec<_> = items.into_iter().map(Into::into).collect();
165    if should_build_index_parallel(items.len()) {
166        return items
167            .into_par_iter()
168            .enumerate()
169            .map(|(id, item)| build_candidate(id, item, backend, config))
170            .collect();
171    }
172
173    items
174        .into_iter()
175        .enumerate()
176        .map(|(id, item)| build_candidate(id, item, backend, config))
177        .collect()
178}
179
180fn should_build_index_parallel(len: usize) -> bool {
181    len >= PARALLEL_INDEX_THRESHOLD && rayon::current_num_threads() > 1
182}
183
184pub fn dedup_and_limit_keys(keys: Vec<SearchKey>, config: &SearchConfig) -> Vec<SearchKey> {
185    let mut seen = HashSet::new();
186    let mut out = Vec::new();
187    let mut total_bytes = 0usize;
188
189    for key in keys {
190        if !seen.insert((key.kind, key.text.clone())) {
191            continue;
192        }
193
194        let required_base_key = matches!(key.kind, KeyKind::Original | KeyKind::Normalized);
195        let would_exceed_count = out.len() >= config.max_search_keys_per_candidate;
196        let would_exceed_bytes =
197            total_bytes + key.text.len() > config.max_total_key_bytes_per_candidate;
198
199        if !required_base_key && (would_exceed_count || would_exceed_bytes) {
200            continue;
201        }
202
203        total_bytes += key.text.len();
204        out.push(key);
205    }
206
207    out
208}
209
210#[cfg(test)]
211mod tests {
212    use crate::{query::PlainBackend, KeyKind};
213
214    use super::*;
215
216    #[test]
217    fn plain_mode_only_original_and_normalized() {
218        let cand = build_candidate(0, "東京駅", &PlainBackend, &SearchConfig::default());
219
220        assert!(cand.keys.iter().any(|k| k.kind == KeyKind::Original));
221        assert!(cand.keys.iter().any(|k| k.kind == KeyKind::Normalized));
222        assert!(!cand
223            .keys
224            .iter()
225            .any(|k| matches!(k.kind, KeyKind::KanaReading | KeyKind::RomajiReading)));
226    }
227
228    #[test]
229    fn original_key_is_always_present() {
230        let cand = build_candidate(0, "README.md", &PlainBackend, &SearchConfig::default());
231        assert!(cand.keys.iter().any(|k| k.kind == KeyKind::Original));
232    }
233
234    #[test]
235    fn search_keys_are_deduped_and_capped() {
236        let cfg = SearchConfig {
237            max_search_keys_per_candidate: 4,
238            ..SearchConfig::default()
239        };
240        let keys = vec![
241            SearchKey::original("a"),
242            SearchKey::normalized("a"),
243            SearchKey::normalized("a"),
244            SearchKey::learned_alias("b"),
245            SearchKey::learned_alias("c"),
246            SearchKey::learned_alias("d"),
247        ];
248
249        let out = dedup_and_limit_keys(keys, &cfg);
250
251        assert!(out.len() <= 4);
252        assert_eq!(
253            out.len(),
254            out.iter()
255                .map(|k| (k.kind, k.text.as_str()))
256                .collect::<HashSet<_>>()
257                .len()
258        );
259    }
260
261    #[test]
262    fn parallel_index_preserves_input_order_and_ids() {
263        let cfg = SearchConfig::default();
264        let cand = build_index(["one", "two", "three", "four"], &PlainBackend, &cfg);
265
266        assert_eq!(
267            cand.iter()
268                .map(|candidate| (candidate.id, candidate.display.as_str()))
269                .collect::<Vec<_>>(),
270            vec![(0, "one"), (1, "two"), (2, "three"), (3, "four")]
271        );
272    }
273}
yuru_core/candidate.rs

yuru_core/
candidate.rs
