yuru_core/
query.rs

1use std::collections::HashMap;
2
3use crate::{
4    KeyBudget, KeyKind, LangMode, LanguageBackend, QueryBudget, QueryVariantKind, SearchConfig,
5};
6
7/// Search text variant produced from the user's query.
8#[derive(Clone, Debug, Eq, PartialEq)]
9pub struct QueryVariant {
10    /// Text that will be matched against compatible search keys.
11    pub text: String,
12    /// Variant type used to filter compatible key kinds.
13    pub kind: QueryVariantKind,
14    /// Score adjustment for this variant type.
15    pub weight: i32,
16}
17
18impl QueryVariant {
19    /// Creates an original query variant.
20    pub fn original(text: impl Into<String>) -> Self {
21        Self {
22            text: text.into(),
23            kind: QueryVariantKind::Original,
24            weight: 500,
25        }
26    }
27
28    /// Creates a normalized query variant.
29    pub fn normalized(text: impl Into<String>) -> Self {
30        Self {
31            text: text.into(),
32            kind: QueryVariantKind::Normalized,
33            weight: 450,
34        }
35    }
36
37    /// Creates a kana query variant.
38    pub fn kana(text: impl Into<String>) -> Self {
39        Self {
40            text: text.into(),
41            kind: QueryVariantKind::Kana,
42            weight: 350,
43        }
44    }
45
46    /// Creates a romaji-to-kana query variant.
47    pub fn romaji_to_kana(text: impl Into<String>) -> Self {
48        Self {
49            text: text.into(),
50            kind: QueryVariantKind::RomajiToKana,
51            weight: 200,
52        }
53    }
54
55    /// Creates a pinyin query variant.
56    pub fn pinyin(text: impl Into<String>) -> Self {
57        Self {
58            text: text.into(),
59            kind: QueryVariantKind::Pinyin,
60            weight: 250,
61        }
62    }
63
64    /// Creates an initials query variant.
65    pub fn initials(text: impl Into<String>) -> Self {
66        Self {
67            text: text.into(),
68            kind: QueryVariantKind::Initials,
69            weight: 250,
70        }
71    }
72}
73
74#[derive(Clone, Debug, Default)]
75/// Backend that only uses original and normalized text.
76pub struct PlainBackend;
77
78impl LanguageBackend for PlainBackend {
79    fn mode(&self) -> LangMode {
80        LangMode::Plain
81    }
82
83    fn build_candidate_keys(&self, _text: &str, _budget: KeyBudget) -> Vec<crate::SearchKey> {
84        Vec::new()
85    }
86
87    fn expand_query(&self, query: &str, _budget: QueryBudget) -> Vec<QueryVariant> {
88        base_query_variants(query)
89    }
90}
91
92/// Builds the language-neutral original and normalized query variants.
93pub fn base_query_variants(query: &str) -> Vec<QueryVariant> {
94    let mut variants = vec![QueryVariant::original(query)];
95    let normalized = crate::normalize::normalize(query);
96    if normalized != query {
97        variants.push(QueryVariant::normalized(normalized));
98    }
99    variants
100}
101
102/// Deduplicates variants by text and key coverage, then applies the query cap.
103pub fn dedup_and_limit_variants(
104    variants: Vec<QueryVariant>,
105    max_query_variants: usize,
106) -> Vec<QueryVariant> {
107    let mut seen_coverage_by_text = HashMap::new();
108    let mut out = Vec::new();
109
110    for variant in variants {
111        let coverage = key_kind_coverage(variant.kind);
112        let seen_coverage = seen_coverage_by_text
113            .entry(variant.text.clone())
114            .or_insert(0u16);
115        if coverage & !*seen_coverage != 0 {
116            *seen_coverage |= coverage;
117            out.push(variant);
118        }
119        if out.len() >= max_query_variants {
120            break;
121        }
122    }
123
124    out
125}
126
127/// Expands and caps query variants for one search run.
128pub(crate) fn prepare_query_variants(
129    query: &str,
130    backend: &dyn LanguageBackend,
131    config: &SearchConfig,
132) -> Vec<QueryVariant> {
133    dedup_and_limit_variants(
134        backend.expand_query(query, config.query_budget()),
135        config.max_query_variants,
136    )
137}
138
139/// Returns whether a key kind is disabled by case/normalization settings.
140pub(crate) fn key_blocked_by_config(kind: KeyKind, config: &SearchConfig) -> bool {
141    kind == KeyKind::Normalized && (config.case_sensitive || !config.normalize)
142}
143
144/// Returns whether a query variant is disabled by case/normalization settings.
145pub(crate) fn variant_blocked_by_config(kind: QueryVariantKind, config: &SearchConfig) -> bool {
146    kind == QueryVariantKind::Normalized && (config.case_sensitive || !config.normalize)
147}
148
149fn key_kind_coverage(kind: QueryVariantKind) -> u16 {
150    compatible_key_kinds(kind)
151        .iter()
152        .fold(0, |coverage, kind| coverage | key_kind_bit(*kind))
153}
154
155/// Returns whether a query variant may be scored against a key kind.
156pub fn key_kind_allowed(variant: &QueryVariant, kind: KeyKind) -> bool {
157    compatible_key_kinds(variant.kind).contains(&kind)
158}
159
160const ORIGINAL_QUERY_KEYS: &[KeyKind] = &[
161    KeyKind::Original,
162    KeyKind::Normalized,
163    KeyKind::RomajiReading,
164    KeyKind::PinyinFull,
165    KeyKind::PinyinJoined,
166    KeyKind::KoreanRomanized,
167    KeyKind::KoreanInitials,
168    KeyKind::KoreanKeyboard,
169    KeyKind::LearnedAlias,
170];
171const KANA_QUERY_KEYS: &[KeyKind] = &[KeyKind::KanaReading];
172const PINYIN_QUERY_KEYS: &[KeyKind] = &[KeyKind::PinyinFull, KeyKind::PinyinJoined];
173const INITIAL_QUERY_KEYS: &[KeyKind] = &[
174    KeyKind::PinyinInitials,
175    KeyKind::KoreanInitials,
176    KeyKind::LearnedAlias,
177];
178
179fn compatible_key_kinds(kind: QueryVariantKind) -> &'static [KeyKind] {
180    match kind {
181        QueryVariantKind::Original | QueryVariantKind::Normalized => ORIGINAL_QUERY_KEYS,
182        QueryVariantKind::Kana | QueryVariantKind::RomajiToKana => KANA_QUERY_KEYS,
183        QueryVariantKind::Pinyin => PINYIN_QUERY_KEYS,
184        QueryVariantKind::Initials => INITIAL_QUERY_KEYS,
185    }
186}
187
188fn key_kind_bit(kind: KeyKind) -> u16 {
189    1 << (kind as u16)
190}
191
192#[cfg(test)]
193mod tests;
yuru_core/query.rs

yuru_core/
query.rs
