use std::collections::HashSet;
use crate::{KeyKind, LanguageBackend, SearchConfig};
use rayon::prelude::*;
#[cfg(not(test))]
const PARALLEL_INDEX_THRESHOLD: usize = 50_000;
#[cfg(test)]
const PARALLEL_INDEX_THRESHOLD: usize = 4;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct SourceSpan {
pub start_char: usize,
pub end_char: usize,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Candidate {
pub id: usize,
pub display: String,
pub keys: Vec<SearchKey>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SearchKey {
pub text: String,
pub kind: KeyKind,
pub weight: i32,
pub case_fold_only: bool,
pub source_map: Option<Box<[Option<SourceSpan>]>>,
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct MappedText {
pub text: String,
pub source_map: Vec<Option<SourceSpan>>,
}
#[derive(Clone, Debug, Default)]
pub struct MappedTextBuilder {
mapped: MappedText,
}
impl MappedTextBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn push_str(&mut self, text: &str, source: Option<SourceSpan>) {
self.mapped.text.push_str(text);
self.mapped.source_map.extend(text.chars().map(|_| source));
}
pub fn push_char(&mut self, ch: char, source: Option<SourceSpan>) {
self.mapped.text.push(ch);
self.mapped.source_map.push(source);
}
pub fn push_unmapped_char(&mut self, ch: char) {
self.mapped.text.push(ch);
self.mapped.source_map.push(None);
}
pub fn finish(self) -> MappedText {
self.mapped
}
}
impl SearchKey {
pub fn new(kind: KeyKind, text: impl Into<String>) -> Self {
Self {
text: text.into(),
kind,
weight: Self::default_weight(kind),
case_fold_only: false,
source_map: None,
}
}
pub fn default_weight(kind: KeyKind) -> i32 {
match kind {
KeyKind::Original => 3000,
KeyKind::Normalized => 2800,
KeyKind::KanaReading => 1700,
KeyKind::RomajiReading => 1800,
KeyKind::PinyinFull => 1750,
KeyKind::PinyinJoined => 1800,
KeyKind::PinyinInitials => 1850,
KeyKind::KoreanRomanized => 1800,
KeyKind::KoreanInitials => 1850,
KeyKind::KoreanKeyboard => 1750,
KeyKind::LearnedAlias => 2500,
}
}
pub fn with_case_fold_only(mut self, case_fold_only: bool) -> Self {
self.case_fold_only = case_fold_only;
self
}
pub fn with_source_map(mut self, source_map: Vec<Option<SourceSpan>>) -> Self {
self.source_map = Some(source_map.into_boxed_slice());
self
}
pub fn original(text: impl Into<String>) -> Self {
Self::new(KeyKind::Original, text)
}
pub fn normalized(text: impl Into<String>) -> Self {
Self::new(KeyKind::Normalized, text)
}
pub fn kana_reading(text: impl Into<String>) -> Self {
Self::new(KeyKind::KanaReading, text)
}
pub fn romaji_reading(text: impl Into<String>) -> Self {
Self::new(KeyKind::RomajiReading, text)
}
pub fn pinyin_full(text: impl Into<String>) -> Self {
Self::new(KeyKind::PinyinFull, text)
}
pub fn pinyin_joined(text: impl Into<String>) -> Self {
Self::new(KeyKind::PinyinJoined, text)
}
pub fn pinyin_initials(text: impl Into<String>) -> Self {
Self::new(KeyKind::PinyinInitials, text)
}
pub fn korean_romanized(text: impl Into<String>) -> Self {
Self::new(KeyKind::KoreanRomanized, text)
}
pub fn korean_initials(text: impl Into<String>) -> Self {
Self::new(KeyKind::KoreanInitials, text)
}
pub fn korean_keyboard(text: impl Into<String>) -> Self {
Self::new(KeyKind::KoreanKeyboard, text)
}
pub fn learned_alias(text: impl Into<String>) -> Self {
Self::new(KeyKind::LearnedAlias, text)
}
}
pub fn build_candidate(
id: usize,
display: impl Into<String>,
backend: &dyn LanguageBackend,
config: &SearchConfig,
) -> Candidate {
let display = display.into();
let mut keys = vec![SearchKey::original(display.clone())];
if config.normalize {
keys.push(normalized_base_key(&display, backend));
}
keys.extend(backend.build_candidate_keys(&display, config.key_budget()));
let keys = dedup_and_limit_keys(keys, config);
Candidate { id, display, keys }
}
fn normalized_base_key(display: &str, backend: &dyn LanguageBackend) -> SearchKey {
let normalized = backend.normalize_candidate(display);
let case_fold_only = normalized
.as_bytes()
.eq_ignore_ascii_case(display.as_bytes())
|| normalized
.chars()
.eq(display.chars().map(crate::matcher::fold_case_char));
SearchKey::normalized(normalized).with_case_fold_only(case_fold_only)
}
pub fn build_index<I, S>(
items: I,
backend: &dyn LanguageBackend,
config: &SearchConfig,
) -> Vec<Candidate>
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
let items: Vec<_> = items.into_iter().map(Into::into).collect();
if should_build_index_parallel(items.len()) {
return items
.into_par_iter()
.enumerate()
.map(|(id, item)| build_candidate(id, item, backend, config))
.collect();
}
items
.into_iter()
.enumerate()
.map(|(id, item)| build_candidate(id, item, backend, config))
.collect()
}
fn should_build_index_parallel(len: usize) -> bool {
len >= PARALLEL_INDEX_THRESHOLD && rayon::current_num_threads() > 1
}
pub fn dedup_and_limit_keys(keys: Vec<SearchKey>, config: &SearchConfig) -> Vec<SearchKey> {
let mut seen = HashSet::new();
let mut out = Vec::new();
let mut total_bytes = 0usize;
for key in keys {
if !seen.insert((key.kind, key.text.clone())) {
continue;
}
let required_base_key = matches!(key.kind, KeyKind::Original | KeyKind::Normalized);
let would_exceed_count = out.len() >= config.max_search_keys_per_candidate;
let would_exceed_bytes =
total_bytes + key.text.len() > config.max_total_key_bytes_per_candidate;
if !required_base_key && (would_exceed_count || would_exceed_bytes) {
continue;
}
total_bytes += key.text.len();
out.push(key);
}
out
}
#[cfg(test)]
mod tests;