use decompound::{decompound, DecompositionOptions};
use once_cell::sync::Lazy;
use probe_code::ranking::get_stemmer;
use probe_code::search::term_exceptions::{is_exception_term, EXCEPTION_TERMS};
use std::collections::HashSet;
use std::sync::Mutex;
static DYNAMIC_SPECIAL_TERMS: Lazy<Mutex<HashSet<String>>> =
Lazy::new(|| Mutex::new(HashSet::new()));
pub fn add_special_term(term: &str) {
let mut special_terms = DYNAMIC_SPECIAL_TERMS.lock().unwrap();
special_terms.insert(term.to_lowercase());
if std::env::var("DEBUG").unwrap_or_default() == "1" {
println!("DEBUG: Added special term: {term}");
}
}
static ENGLISH_STOP_WORDS: Lazy<HashSet<String>> = Lazy::new(|| {
vec![
"a",
"about",
"above",
"after",
"again",
"against",
"all",
"am",
"an",
"and",
"any",
"are",
"aren't",
"as",
"at",
"be",
"because",
"been",
"before",
"being",
"below",
"between",
"both",
"but",
"by",
"can't",
"cannot",
"could",
"couldn't",
"did",
"didn't",
"do",
"does",
"doesn't",
"doing",
"don't",
"down",
"during",
"each",
"few",
"for",
"from",
"further",
"had",
"hadn't",
"has",
"hasn't",
"have",
"haven't",
"having",
"he",
"he'd",
"he'll",
"he's",
"her",
"here",
"here's",
"hers",
"herself",
"him",
"himself",
"his",
"how",
"how's",
"i",
"i'd",
"i'll",
"i'm",
"i've",
"if",
"in",
"into",
"is",
"isn't",
"it",
"it's",
"its",
"itself",
"let's",
"me",
"more",
"most",
"mustn't",
"my",
"myself",
"no",
"nor",
"not",
"of",
"off",
"on",
"once",
"only",
"or",
"other",
"ought",
"our",
"ours",
"ourselves",
"out",
"over",
"own",
"same",
"shan't",
"she",
"she'd",
"she'll",
"she's",
"should",
"shouldn't",
"so",
"some",
"such",
"than",
"that",
"that's",
"the",
"their",
"theirs",
"them",
"themselves",
"then",
"there",
"there's",
"these",
"they",
"they'd",
"they'll",
"they're",
"they've",
"this",
"those",
"through",
"to",
"too",
"under",
"until",
"up",
"very",
"was",
"wasn't",
"we",
"we'd",
"we'll",
"we're",
"we've",
"were",
"weren't",
"what",
"what's",
"when",
"when's",
"where",
"where's",
"which",
"while",
"who",
"who's",
"whom",
"why",
"why's",
"with",
"won't",
"would",
"wouldn't",
"you",
"you'd",
"you'll",
"you're",
"you've",
"your",
"yours",
"yourself",
"yourselves",
"ing",
]
.into_iter()
.map(String::from)
.collect()
});
static PROGRAMMING_STOP_WORDS: Lazy<HashSet<String>> = Lazy::new(|| {
vec![
"func",
"type",
"struct",
"interface",
"chan",
"map",
"go",
"defer",
"var",
"let",
"const",
"return",
"if",
"else",
"for",
"while",
"switch",
"case",
"break",
"continue",
"default",
"try",
"catch",
"finally",
"throw",
"new",
"super",
"extends",
"implements",
"function",
"class",
"method",
"this",
"public",
"private",
"protected",
"static",
"final",
"async",
"await",
"string",
"int",
"bool",
"float",
"void",
"null",
"nil",
"class",
"enum",
"impl",
"fn",
"mod",
]
.into_iter()
.map(String::from)
.collect()
});
static SPECIAL_CASE_WORDS: Lazy<HashSet<String>> = Lazy::new(|| {
vec![
"oauth",
"oauth2",
"ipv4",
"ipv6",
"ipv",
"graphql",
"postgresql",
"mysql",
"mongodb",
"javascript",
"typescript",
"nodejs",
"reactjs",
"vuejs",
"angularjs",
"github",
"gitlab",
"bitbucket",
"kubernetes",
"docker",
"webpack",
"rollup",
"vite",
"eslint",
"prettier",
"axios",
"fetch",
"grpc",
"http2",
"whitelist",
"blacklist",
"allowlist",
"blocklist",
"denylist",
]
.into_iter()
.map(String::from)
.collect()
});
#[inline]
fn is_uppercase(c: char) -> bool {
c.is_ascii_uppercase()
}
#[inline]
fn is_lowercase(c: char) -> bool {
c.is_ascii_lowercase()
}
#[inline]
fn is_number(c: char) -> bool {
c.is_ascii_digit()
}
pub fn is_special_case(word: &str) -> bool {
let lowercase = word.to_lowercase();
if SPECIAL_CASE_WORDS.contains(&lowercase) {
return true;
}
let special_terms = DYNAMIC_SPECIAL_TERMS.lock().unwrap();
if special_terms.contains(&lowercase) {
if std::env::var("DEBUG").unwrap_or_default() == "1" {
println!("DEBUG: Found dynamic special term: {lowercase}");
}
return true;
}
false
}
pub fn split_camel_case(input: &str) -> Vec<String> {
let _debug_mode = std::env::var("DEBUG").unwrap_or_default() == "1";
if input.is_empty() {
return vec![];
}
if is_special_case(input) {
return vec![input.to_lowercase()];
}
let lowercase = input.to_lowercase();
if lowercase.starts_with("oauth2") {
let remaining = &input[6..]; if !remaining.is_empty() {
let mut result = vec!["oauth2".to_string()];
result.extend(split_camel_case(remaining));
return result;
}
}
let mut special_cases: Vec<&String> = SPECIAL_CASE_WORDS.iter().collect();
special_cases.sort_by_key(|b| std::cmp::Reverse(b.len()));
for special_case in special_cases {
if lowercase.starts_with(special_case) {
let _original_part = &input[0..special_case.len()];
let remaining = &input[special_case.len()..];
if !remaining.is_empty() {
let mut result = vec![special_case.clone()];
result.extend(split_camel_case(remaining));
return result;
}
}
}
if input == lowercase && !input.contains('_') && input.len() > 3 {
let _potential_splits: Vec<String> = Vec::new();
let common_terms = EXCEPTION_TERMS
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>();
for term in common_terms {
if input.contains(term) && term != input {
let parts: Vec<&str> = input.split(term).collect();
if parts.len() > 1 {
let mut result = Vec::new();
for (i, part) in parts.iter().enumerate() {
if !part.is_empty() {
result.push(part.to_string());
}
if i < parts.len() - 1 {
result.push(term.to_string());
}
}
if !result.is_empty() {
return result;
}
}
}
}
}
let chars: Vec<char> = input.chars().collect();
let mut result = Vec::new();
let mut current_word = String::new();
let mut prev_is_lower = false;
let mut prev_is_upper = false;
let mut prev_is_digit = false;
for (i, &c) in chars.iter().enumerate() {
let is_upper = is_uppercase(c);
let is_lower = is_lowercase(c);
let is_digit = is_number(c);
let start_new_word =
!current_word.is_empty() && (
(prev_is_lower && is_upper) ||
(prev_is_digit != is_digit) ||
(prev_is_upper && is_upper && i + 1 < chars.len() && is_lowercase(chars[i + 1]))
);
if start_new_word {
result.push(current_word);
current_word = String::new();
}
current_word.push(c);
prev_is_lower = is_lower;
prev_is_upper = is_upper;
prev_is_digit = is_digit;
}
if !current_word.is_empty() {
result.push(current_word);
}
result.into_iter().map(|word| word.to_lowercase()).collect()
}
pub fn is_english_stop_word(word: &str) -> bool {
if let Ok(num) = word.parse::<u32>() {
if num <= 10 {
return true;
}
}
ENGLISH_STOP_WORDS.contains(word)
}
pub fn is_programming_stop_word(word: &str) -> bool {
PROGRAMMING_STOP_WORDS.contains(word)
}
pub fn is_stop_word(word: &str) -> bool {
is_english_stop_word(word) || is_programming_stop_word(word)
}
pub fn split_compound_word(word: &str, vocab: &HashSet<String>) -> Vec<String> {
if is_special_case(word) {
return vec![word.to_lowercase()];
}
let common_terms = EXCEPTION_TERMS
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>();
if common_terms.contains(&word.to_lowercase().as_str()) {
return vec![word.to_string()];
}
if vocab.contains(&word.to_lowercase()) {
return vec![word.to_string()];
}
let is_valid_word = |w: &str| vocab.contains(&w.to_lowercase());
match decompound(word, &is_valid_word, DecompositionOptions::empty()) {
Ok(parts) if !parts.is_empty() => parts,
_ => vec![word.to_string()],
}
}
pub fn load_vocabulary() -> &'static HashSet<String> {
static VOCABULARY: Lazy<HashSet<String>> = Lazy::new(|| {
vec![
"white",
"black",
"list",
"mail",
"back",
"ground",
"book",
"mark",
"key",
"word",
"pass",
"fire",
"wall",
"firewall",
"water",
"fall",
"data",
"base",
"time",
"stamp",
"air",
"port",
"blue",
"tooth",
"green",
"house",
"red",
"hat",
"yellow",
"pages",
"blue",
"print",
"type",
"script",
"java",
"script",
"note",
"pad",
"web",
"site",
"page",
"view",
"code",
"base",
"name",
"space",
"class",
"room",
"work",
"flow",
"life",
"cycle",
"end",
"point",
"check",
"box",
"drop",
"down",
"pop",
"up",
"side",
"bar",
"tool",
"tip",
"drag",
"drop",
"click",
"stream",
"line",
"dead",
"lock",
"race",
"condition",
"thread",
"safe",
"memory",
"leak",
"stack",
"trace",
"heap",
"dump",
"core",
"file",
"system",
"disk",
"drive",
"hard",
"soft",
"ware",
"firm",
"middle",
"front",
"back",
"end",
"full",
"stack",
"dev",
"ops",
"micro",
"service",
"mono",
"lith",
"container",
"docker",
"pod",
"cloud",
"native",
"server",
"less",
"function",
"as",
"service",
"infra",
"structure",
"platform",
"test",
"driven",
"behavior",
"continuous",
"integration",
"deployment",
"delivery",
"pipeline",
"git",
"hub",
"lab",
"version",
"control",
"branch",
"merge",
"pull",
"request",
"commit",
"push",
"clone",
"fork",
"repository",
"issue",
"bug",
"feature",
"release",
"tag",
"semantic",
"versioning",
"major",
"minor",
"patch",
"alpha",
"beta",
"stable",
"unstable",
"deprecated",
"legacy",
"modern",
"framework",
"library",
"package",
"module",
"component",
"prop",
"state",
"hook",
"effect",
"context",
"provider",
"consumer",
"reducer",
"action",
"store",
"dispatch",
"subscribe",
"publish",
"event",
"handler",
"listener",
"callback",
"promise",
"async",
"await",
"future",
"stream",
"observable",
"reactive",
"functional",
"object",
"oriented",
"procedural",
"declarative",
"imperative",
"mutable",
"immutable",
"pure",
"side",
"effect",
"higher",
"order",
"first",
"class",
"citizen",
"closure",
"scope",
"lexical",
"dynamic",
"static",
"type",
"inference",
"checking",
"compile",
"time",
"run",
"error",
"exception",
"try",
"catch",
"finally",
"throw",
"raise",
"handle",
"logging",
"debug",
"info",
"warn",
"error",
"fatal",
"trace",
"metric",
"monitor",
"alert",
"notification",
"dashboard",
"report",
"analytics",
"insight",
"data",
"science",
"machine",
"learning",
"artificial",
"intelligence",
"neural",
"network",
"deep",
"reinforcement",
"supervised",
"unsupervised",
"classification",
"regression",
"clustering",
"recommendation",
"prediction",
"inference",
"training",
"validation",
"test",
"accuracy",
"precision",
"recall",
"f1",
"score",
"loss",
"function",
"gradient",
"descent",
"back",
"propagation",
"forward",
"pass",
"epoch",
"batch",
"mini",
"over",
"fitting",
"under",
"regularization",
"dropout",
"batch",
"normalization",
"activation",
"sigmoid",
"tanh",
"relu",
"leaky",
"softmax",
"convolution",
"pooling",
"recurrent",
"lstm",
"gru",
"transformer",
"attention",
"encoder",
"decoder",
"embedding",
"token",
"tokenization",
"stemming",
"lemmatization",
"stop",
"word",
"n",
"gram",
"tf",
"idf",
"cosine",
"similarity",
"euclidean",
"distance",
"manhattan",
"jaccard",
"index",
"precision",
"recall",
"relevance",
"ranking",
"page",
"rank",
"search",
"engine",
"crawler",
"indexer",
"query",
"result",
"snippet",
"cache",
"hit",
"miss",
"eviction",
"policy",
"lru",
"fifo",
"lifo",
"priority",
"queue",
"stack",
"heap",
"tree",
"binary",
"balanced",
"avl",
"red",
"black",
"b",
"trie",
"hash",
"map",
"set",
"list",
"linked",
"doubly",
"circular",
"array",
"vector",
"matrix",
"tensor",
"graph",
"directed",
"undirected",
"weighted",
"unweighted",
"adjacency",
"matrix",
"list",
"edge",
"vertex",
"node",
"path",
"cycle",
"traversal",
"breadth",
"first",
"depth",
"topological",
"sort",
"minimum",
"spanning",
"tree",
"shortest",
"path",
"dijkstra",
"bellman",
"ford",
"floyd",
"warshall",
"kruskal",
"prim",
"greedy",
"dynamic",
"programming",
"divide",
"conquer",
"backtracking",
"branch",
"bound",
"heuristic",
"approximation",
"randomized",
"parallel",
"concurrent",
"distributed",
"synchronous",
"asynchronous",
"blocking",
"non",
"mutex",
"semaphore",
"lock",
"atomic",
"volatile",
"transaction",
"acid",
"consistency",
"isolation",
"durability",
"serializable",
"repeatable",
"read",
"committed",
"uncommitted",
"phantom",
"dirty",
"read",
"write",
"skew",
"conflict",
"resolution",
"optimistic",
"pessimistic",
"two",
"phase",
"commit",
"rollback",
"savepoint",
"checkpoint",
"recovery",
"backup",
"restore",
"archive",
"log",
"journal",
"redo",
"undo",
"write",
"ahead",
"logging",
"snapshot",
"isolation",
"level",
"serializable",
"repeatable",
"read",
"committed",
"uncommitted",
"phantom",
"dirty",
"read",
"write",
"skew",
"conflict",
"resolution",
"optimistic",
"pessimistic",
"two",
"phase",
"commit",
"rollback",
"savepoint",
"checkpoint",
"recovery",
"backup",
"restore",
"archive",
"log",
"journal",
"redo",
"undo",
"write",
"ahead",
"logging",
"snapshot",
"isolation",
"level",
]
.into_iter()
.map(String::from)
.collect()
});
&VOCABULARY
}
#[allow(dead_code)]
pub fn tokenize_and_stem(keyword: &str) -> Vec<String> {
let stemmer = get_stemmer();
let vocabulary = load_vocabulary();
let camel_parts = split_camel_case(keyword);
if camel_parts.len() > 1 {
camel_parts
.into_iter()
.filter(|part| !is_stop_word(part))
.map(|part| stemmer.stem(&part).to_string())
.collect()
} else {
let compound_parts = split_compound_word(keyword, vocabulary);
if compound_parts.len() > 1 {
compound_parts
.into_iter()
.filter(|part| !is_stop_word(part))
.map(|part| stemmer.stem(&part).to_string())
.collect()
} else {
vec![stemmer.stem(keyword).to_string()]
}
}
}
pub fn tokenize(text: &str) -> Vec<String> {
let stemmer = get_stemmer();
let vocabulary = load_vocabulary();
let mut negated_terms = HashSet::new();
let mut tokens = Vec::new();
for word in text.split_whitespace() {
let is_negated = word.starts_with('-');
let mut current_token = String::new();
let mut chars = word.chars();
if is_negated {
chars.next();
}
for c in chars {
if c.is_alphanumeric() {
current_token.push(c);
} else if !current_token.is_empty() {
if is_negated {
negated_terms.insert(current_token.to_lowercase());
}
tokens.push(current_token);
current_token = String::new();
}
}
if !current_token.is_empty() {
if is_negated {
negated_terms.insert(current_token.to_lowercase());
}
tokens.push(current_token);
}
}
let mut processed_tokens = HashSet::new();
let mut result = Vec::new();
for token in tokens {
let parts = split_camel_case(&token);
for part in parts {
let lowercase_part = part.to_lowercase();
if is_stop_word(&lowercase_part) {
continue;
}
if negated_terms.contains(&lowercase_part) {
continue;
}
let compound_parts = split_compound_word(&lowercase_part, vocabulary);
for compound_part in compound_parts {
if is_stop_word(&compound_part) {
continue;
}
if negated_terms.contains(&compound_part) {
continue;
}
if is_exception_term(&compound_part)
&& processed_tokens.insert(compound_part.clone())
{
result.push(compound_part.clone());
}
let stemmed_part = stemmer.stem(&compound_part).to_string();
if negated_terms.contains(&stemmed_part) {
continue;
}
if processed_tokens.insert(stemmed_part.clone()) {
result.push(stemmed_part);
}
}
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_split_camel_case() {
assert_eq!(split_camel_case("camelCase"), vec!["camel", "case"]);
assert_eq!(split_camel_case("PascalCase"), vec!["pascal", "case"]);
assert_eq!(
split_camel_case("parseJSONToHTML5"),
vec!["parse", "json", "to", "html", "5"]
);
assert_eq!(split_camel_case("APIDefinition"), vec!["api", "definition"]);
assert_eq!(
split_camel_case("OAuth2Provider"),
vec!["oauth2", "provider"]
);
assert_eq!(split_camel_case("typeIgnore"), vec!["type", "ignore"]);
assert_eq!(
split_camel_case("migrateEndpointMetaByType"),
vec!["migrate", "endpoint", "meta", "by", "type"]
);
}
#[test]
fn test_stop_words() {
assert!(is_programming_stop_word("func"));
assert!(is_programming_stop_word("type"));
assert!(is_programming_stop_word("struct"));
assert!(!is_programming_stop_word("migrate"));
assert!(!is_programming_stop_word("endpoint"));
}
#[test]
fn test_tokenize() {
let tokens = tokenize("func (a *APIDefinition) MigrateEndpointMeta()");
assert!(tokens.contains(&"api".to_string()));
assert!(tokens.contains(&"definit".to_string())); assert!(tokens.contains(&"migrat".to_string())); assert!(
tokens.contains(&"endpoint".to_string())
|| (tokens.contains(&"end".to_string()) && tokens.contains(&"point".to_string()))
);
assert!(tokens.contains(&"meta".to_string()));
let tokens = tokenize("func ParseJSONToHTML5()");
assert!(tokens.contains(&"pars".to_string())); assert!(tokens.contains(&"json".to_string()));
assert!(tokens.contains(&"html".to_string()));
let tokens = tokenize("typeIgnore typeWhitelist");
assert!(tokens.contains(&"ignor".to_string()));
let tokens = tokenize("whitelist blackmail firewall");
assert!(tokens.contains(&"whitelist".to_string()));
assert!(tokens.contains(&"black".to_string()));
assert!(tokens.contains(&"mail".to_string()));
assert!(tokens.contains(&"firewall".to_string()));
let tokens = tokenize("enableFirewallWhitelist");
assert!(tokens.contains(&"enabl".to_string())); assert!(tokens.contains(&"firewall".to_string())); assert!(tokens.contains(&"whitelist".to_string()));
}
#[test]
fn test_compound_word_splitting() {
let vocab = HashSet::from([
"white".to_string(),
"list".to_string(),
"black".to_string(),
"mail".to_string(),
]);
let parts = split_compound_word("whitelist", &vocab);
assert_eq!(parts, vec!["whitelist".to_string()]);
let parts = split_compound_word("blackmail", &vocab);
assert_eq!(parts, vec!["black".to_string(), "mail".to_string()]);
let parts = split_compound_word("computer", &vocab);
assert_eq!(parts, vec!["computer".to_string()]);
}
#[test]
fn test_tokenize_with_compound_words() {
let tokens = tokenize("whitelist blackmail firewall");
assert!(tokens.contains(&"whitelist".to_string()));
assert!(tokens.contains(&"black".to_string()));
assert!(tokens.contains(&"mail".to_string()));
assert!(tokens.contains(&"firewall".to_string()));
}
}