use std::collections::HashMap;
use crate::{
KeyBudget, KeyKind, LangMode, LanguageBackend, QueryBudget, QueryVariantKind, SearchConfig,
};
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct QueryVariant {
pub text: String,
pub kind: QueryVariantKind,
pub weight: i32,
}
impl QueryVariant {
pub fn original(text: impl Into<String>) -> Self {
Self {
text: text.into(),
kind: QueryVariantKind::Original,
weight: 500,
}
}
pub fn normalized(text: impl Into<String>) -> Self {
Self {
text: text.into(),
kind: QueryVariantKind::Normalized,
weight: 450,
}
}
pub fn kana(text: impl Into<String>) -> Self {
Self {
text: text.into(),
kind: QueryVariantKind::Kana,
weight: 350,
}
}
pub fn romaji_to_kana(text: impl Into<String>) -> Self {
Self {
text: text.into(),
kind: QueryVariantKind::RomajiToKana,
weight: 200,
}
}
pub fn pinyin(text: impl Into<String>) -> Self {
Self {
text: text.into(),
kind: QueryVariantKind::Pinyin,
weight: 250,
}
}
pub fn initials(text: impl Into<String>) -> Self {
Self {
text: text.into(),
kind: QueryVariantKind::Initials,
weight: 250,
}
}
}
#[derive(Clone, Debug, Default)]
pub struct PlainBackend;
impl LanguageBackend for PlainBackend {
fn mode(&self) -> LangMode {
LangMode::Plain
}
fn build_candidate_keys(&self, _text: &str, _budget: KeyBudget) -> Vec<crate::SearchKey> {
Vec::new()
}
fn expand_query(&self, query: &str, _budget: QueryBudget) -> Vec<QueryVariant> {
base_query_variants(query)
}
}
pub fn base_query_variants(query: &str) -> Vec<QueryVariant> {
let mut variants = vec![QueryVariant::original(query)];
let normalized = crate::normalize::normalize(query);
if normalized != query {
variants.push(QueryVariant::normalized(normalized));
}
variants
}
pub fn dedup_and_limit_variants(
variants: Vec<QueryVariant>,
max_query_variants: usize,
) -> Vec<QueryVariant> {
let mut seen_coverage_by_text = HashMap::new();
let mut out = Vec::new();
for variant in variants {
let coverage = key_kind_coverage(variant.kind);
let seen_coverage = seen_coverage_by_text
.entry(variant.text.clone())
.or_insert(0u16);
if coverage & !*seen_coverage != 0 {
*seen_coverage |= coverage;
out.push(variant);
}
if out.len() >= max_query_variants {
break;
}
}
out
}
pub(crate) fn prepare_query_variants(
query: &str,
backend: &dyn LanguageBackend,
config: &SearchConfig,
) -> Vec<QueryVariant> {
dedup_and_limit_variants(
backend.expand_query(query, config.query_budget()),
config.max_query_variants,
)
}
pub(crate) fn key_blocked_by_config(
key: &crate::SearchKey,
config: &SearchConfig,
scorer_folds_case: bool,
) -> bool {
key.kind == KeyKind::Normalized
&& (config.case_sensitive || !config.normalize || (key.case_fold_only && scorer_folds_case))
}
pub(crate) fn variant_blocked_by_config(kind: QueryVariantKind, config: &SearchConfig) -> bool {
kind == QueryVariantKind::Normalized && (config.case_sensitive || !config.normalize)
}
fn key_kind_coverage(kind: QueryVariantKind) -> u16 {
compatible_key_kinds(kind)
.iter()
.fold(0, |coverage, kind| coverage | key_kind_bit(*kind))
}
pub fn key_kind_allowed(variant: &QueryVariant, kind: KeyKind) -> bool {
compatible_key_kinds(variant.kind).contains(&kind)
}
const ORIGINAL_QUERY_KEYS: &[KeyKind] = &[
KeyKind::Original,
KeyKind::Normalized,
KeyKind::RomajiReading,
KeyKind::PinyinFull,
KeyKind::PinyinJoined,
KeyKind::KoreanRomanized,
KeyKind::KoreanInitials,
KeyKind::KoreanKeyboard,
KeyKind::LearnedAlias,
];
const KANA_QUERY_KEYS: &[KeyKind] = &[KeyKind::KanaReading];
const PINYIN_QUERY_KEYS: &[KeyKind] = &[KeyKind::PinyinFull, KeyKind::PinyinJoined];
const INITIAL_QUERY_KEYS: &[KeyKind] = &[
KeyKind::PinyinInitials,
KeyKind::KoreanInitials,
KeyKind::LearnedAlias,
];
fn compatible_key_kinds(kind: QueryVariantKind) -> &'static [KeyKind] {
match kind {
QueryVariantKind::Original | QueryVariantKind::Normalized => ORIGINAL_QUERY_KEYS,
QueryVariantKind::Kana | QueryVariantKind::RomajiToKana => KANA_QUERY_KEYS,
QueryVariantKind::Pinyin => PINYIN_QUERY_KEYS,
QueryVariantKind::Initials => INITIAL_QUERY_KEYS,
}
}
fn key_kind_bit(kind: KeyKind) -> u16 {
1 << (kind as u16)
}
#[cfg(test)]
mod tests;