use std::collections::HashSet;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExtractedEntity {
pub entity: String,
pub etype: &'static str,
}
const MAX_ENTITIES_PER_CHUNK: usize = 64;
const MIN_ENTITY_LEN: usize = 2;
const TRIM: &[char] = &[
'"', '\'', '`', ',', ';', '.', '(', ')', '[', ']', '{', '}', '?', '!', '<', '>', '*', '|',
'=', '@', '#', '%', '\\',
];
pub fn extract_entities(content: &str, trigger_desc: Option<&str>) -> Vec<ExtractedEntity> {
let mut raw: Vec<(String, &'static str)> = Vec::new();
for text in [Some(content), trigger_desc].into_iter().flatten() {
collect_backtick_spans(text, &mut raw);
for tok in text.split_whitespace() {
if let Some(cl) = classify(tok) {
raw.push(cl);
}
}
}
let mut seen: HashSet<String> = HashSet::new();
let mut out: Vec<ExtractedEntity> = Vec::new();
for (entity, etype) in raw {
if entity.len() >= MIN_ENTITY_LEN && seen.insert(entity.clone()) {
out.push(ExtractedEntity { entity, etype });
if out.len() >= MAX_ENTITIES_PER_CHUNK {
break;
}
}
}
out
}
fn collect_backtick_spans(text: &str, raw: &mut Vec<(String, &'static str)>) {
let mut in_span = false;
let mut buf = String::new();
for ch in text.chars() {
if ch == '`' {
if in_span {
for piece in buf.split_whitespace() {
if let Some(cl) = classify(piece) {
raw.push(cl);
} else if let Some(sym) = as_identifier(piece) {
raw.push((sym, "symbol"));
}
}
buf.clear();
}
in_span = !in_span;
} else if in_span {
buf.push(ch);
}
}
}
fn classify(tok: &str) -> Option<(String, &'static str)> {
let pre = tok.trim_matches(TRIM);
if pre.starts_with("--") && pre.len() >= 4 && pre[2..].chars().all(is_ident_char) {
return Some((pre.to_lowercase(), "flag"));
}
let t = pre;
if t.len() < MIN_ENTITY_LEN {
return None;
}
if let Some((alpha, digits)) = split_alpha_digits(t) {
if (1..=4).contains(&alpha.len()) && digits.len() >= 3 {
return Some((t.to_lowercase(), "error"));
}
}
if t.contains("::") && t.chars().all(is_ident_char) {
return Some((t.to_lowercase(), "path"));
}
if t.contains('/') && t.contains('.') && t.chars().all(is_ident_char) {
return Some((t.trim_end_matches('/').to_lowercase(), "path"));
}
if (t.contains('_') || t.contains('-')) && is_identifier_token(t) {
return Some((t.to_lowercase(), "symbol"));
}
if is_camel_case(t) {
return Some((t.to_lowercase(), "symbol"));
}
None
}
fn as_identifier(tok: &str) -> Option<String> {
let t = tok.trim_matches(TRIM);
if t.len() >= MIN_ENTITY_LEN && is_identifier_token(t) {
Some(t.to_lowercase())
} else {
None
}
}
fn is_ident_char(c: char) -> bool {
c.is_alphanumeric() || matches!(c, '_' | '-' | ':' | '.' | '/')
}
fn is_identifier_token(t: &str) -> bool {
t.chars().all(is_ident_char) && t.chars().any(|c| c.is_alphabetic())
}
fn split_alpha_digits(t: &str) -> Option<(&str, &str)> {
let split = t.find(|c: char| c.is_ascii_digit())?;
let (alpha, digits) = t.split_at(split);
if !alpha.is_empty()
&& alpha.chars().all(|c| c.is_ascii_alphabetic())
&& !digits.is_empty()
&& digits.chars().all(|c| c.is_ascii_digit())
{
Some((alpha, digits))
} else {
None
}
}
fn is_camel_case(t: &str) -> bool {
if !t.chars().all(|c| c.is_alphanumeric()) {
return false;
}
let bytes: Vec<char> = t.chars().collect();
bytes
.windows(2)
.any(|w| w[0].is_lowercase() && w[1].is_uppercase())
}
#[cfg(test)]
mod tests {
use super::*;
fn ents(content: &str) -> Vec<String> {
extract_entities(content, None)
.into_iter()
.map(|e| e.entity)
.collect()
}
#[test]
fn extracts_error_codes() {
assert!(ents("hit error E0277 while building").contains(&"e0277".to_string()));
assert!(!ents("bumped to 1.2.3 in 2026").contains(&"1.2.3".to_string()));
}
#[test]
fn extracts_flags_and_paths() {
let e = ents("run cargo build --release and edit core/src/kb/recall.rs");
assert!(e.contains(&"--release".to_string()));
assert!(e.contains(&"core/src/kb/recall.rs".to_string()));
}
#[test]
fn extracts_symbols_not_plain_words() {
let e = ents("call KnowledgeBase::recall via the snake_case helper get_deps");
assert!(e.contains(&"knowledgebase::recall".to_string()));
assert!(e.contains(&"snake_case".to_string()));
assert!(e.contains(&"get_deps".to_string()));
assert!(!e.contains(&"call".to_string()));
assert!(!e.contains(&"via".to_string()));
}
#[test]
fn backtick_spans_lower_the_bar() {
let e = ents("the `recall` method matters");
assert!(e.contains(&"recall".to_string()));
assert!(!ents("recall the method").contains(&"recall".to_string()));
}
#[test]
fn dedups_and_caps() {
let e = extract_entities("E0277 E0277 E0277", None);
assert_eq!(e.iter().filter(|x| x.entity == "e0277").count(), 1);
}
}