use std::collections::HashSet;
use crate::engine::text;
use super::{MAX_ENTITIES_PER_MEMORY, MAX_ENTITY_CHARS, clipped_chars};
pub(super) struct EntityCandidate {
pub(super) kind: &'static str,
pub(super) value: String,
pub(super) origin: &'static str,
}
pub(super) fn extract_entities(
statement: &str,
keywords: &[String],
evidence_texts: &[&str],
) -> Vec<EntityCandidate> {
let mut out = Vec::new();
let mut seen = HashSet::new();
collect_entities_from_text(statement, "statement", &mut out, &mut seen);
for keyword in keywords {
collect_entities_from_text(keyword, "keyword", &mut out, &mut seen);
}
for evidence in evidence_texts {
collect_entities_from_text(evidence, "evidence", &mut out, &mut seen);
}
out.truncate(MAX_ENTITIES_PER_MEMORY);
out
}
fn collect_entities_from_text(
text: &str,
origin: &'static str,
out: &mut Vec<EntityCandidate>,
seen: &mut HashSet<String>,
) {
if out.len() >= MAX_ENTITIES_PER_MEMORY {
return;
}
let sanitized = text::sanitize(text);
push_entity_candidates(&sanitized, origin, out, seen);
}
fn push_entity_candidates(
text: &str,
origin: &'static str,
out: &mut Vec<EntityCandidate>,
seen: &mut HashSet<String>,
) {
for span in extract_backtick_spans(text) {
classify_and_push(&span, origin, out, seen);
}
for token in tokenize_entity_candidates(text) {
classify_and_push(&token, origin, out, seen);
if out.len() >= MAX_ENTITIES_PER_MEMORY {
return;
}
}
let mut words = text.split_whitespace().peekable();
while let Some(word) = words.next() {
let lower = word.to_ascii_lowercase();
if matches!(
lower.as_str(),
"fn" | "struct" | "enum" | "trait" | "mod" | "type" | "const" | "static" | "impl"
) && let Some(name) = words.peek()
{
let cleaned = trim_entity_token(name);
if is_symbol_name(&cleaned) {
push_candidate("symbol", &cleaned, origin, out, seen);
}
}
}
}
fn extract_backtick_spans(text: &str) -> Vec<String> {
let mut spans = Vec::new();
let mut rest = text;
while let Some(start) = rest.find('`') {
rest = &rest[start + 1..];
if let Some(end) = rest.find('`') {
let span = rest[..end].trim();
if !span.is_empty() && !span.contains('\n') {
spans.push(span.to_string());
}
rest = &rest[end + 1..];
} else {
break;
}
}
spans
}
fn tokenize_entity_candidates(text: &str) -> Vec<String> {
text.split(|character: char| {
character.is_whitespace()
|| matches!(
character,
',' | ';' | '(' | ')' | '[' | ']' | '{' | '}' | '"' | '\''
)
})
.map(trim_entity_token)
.filter(|token| token.len() > 1 && token.len() <= MAX_ENTITY_CHARS)
.collect()
}
fn trim_entity_token(token: &str) -> String {
token
.trim_matches(|character: char| {
matches!(
character,
'.' | ','
| ';'
| ':'
| '!'
| '?'
| '"'
| '\''
| '`'
| '('
| ')'
| '['
| ']'
| '{'
| '}'
| '<'
| '>'
)
})
.to_string()
}
fn classify_and_push(
raw: &str,
origin: &'static str,
out: &mut Vec<EntityCandidate>,
seen: &mut HashSet<String>,
) {
if out.len() >= MAX_ENTITIES_PER_MEMORY {
return;
}
let token = trim_entity_token(raw);
if token.len() < 2 || token.len() > MAX_ENTITY_CHARS {
return;
}
if let Some(kind) = classify_entity(&token) {
push_candidate(kind, &token, origin, out, seen);
}
}
fn classify_entity(token: &str) -> Option<&'static str> {
if is_path_entity(token) {
return Some("path");
}
if is_command_entity(token) {
return Some("command");
}
if is_crate_entity(token) {
return Some("crate");
}
if is_symbol_entity(token) {
return Some("symbol");
}
if is_concept_entity(token) {
return Some("concept");
}
None
}
fn is_path_entity(token: &str) -> bool {
if token.contains("://") {
return false;
}
let lowered = token.to_ascii_lowercase();
if lowered.starts_with("./") || lowered.starts_with("../") || lowered.starts_with("~/") {
return token.contains('/') || token.contains('\\');
}
if token.starts_with('/') && token.contains('/') {
return true;
}
if token.contains('/') {
let segments: Vec<&str> = token.split('/').filter(|part| !part.is_empty()).collect();
if segments.len() >= 2 {
return true;
}
if let Some(last) = segments.last()
&& last.contains('.')
&& last.rsplit_once('.').is_some_and(|(_, ext)| {
(1..=6).contains(&ext.len()) && ext.chars().all(|c| c.is_ascii_alphanumeric())
})
{
return true;
}
}
if token.contains('.')
&& !token.starts_with('.')
&& token.rsplit_once('.').is_some_and(|(stem, ext)| {
!stem.is_empty()
&& (1..=8).contains(&ext.len())
&& ext.chars().all(|c| c.is_ascii_alphanumeric())
&& stem
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
})
{
return true;
}
token.contains('\\')
}
fn is_command_entity(token: &str) -> bool {
let cleaned = token.trim_start_matches('$').trim();
let first = cleaned.split_whitespace().next().unwrap_or("");
matches!(
first,
"cargo"
| "npm"
| "npx"
| "pnpm"
| "yarn"
| "bun"
| "git"
| "go"
| "python"
| "python3"
| "pip"
| "make"
| "cmake"
| "docker"
| "kubectl"
| "rg"
| "grep"
| "sed"
| "awk"
| "curl"
| "wget"
| "rustc"
| "clippy"
| "rustfmt"
| "goosedump"
| "pi"
) || cleaned.starts_with("cargo ")
|| cleaned.starts_with("npm ")
|| cleaned.starts_with("git ")
|| cleaned.starts_with("docker ")
}
fn is_crate_entity(token: &str) -> bool {
if token.starts_with("crate::") {
return true;
}
if let Some(rest) = token.strip_prefix("use ") {
let name = rest.split("::").next().unwrap_or("").trim();
return !name.is_empty()
&& name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-');
}
token.contains('-')
&& token
.chars()
.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_')
&& token.matches('-').count() >= 1
&& token.len() >= 3
}
fn is_symbol_entity(token: &str) -> bool {
if token.contains("::") {
let parts: Vec<&str> = token.split("::").collect();
return parts.len() >= 2
&& parts.iter().all(|part| {
!part.is_empty() && part.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
});
}
is_symbol_name(token)
&& (token.contains('_')
|| token.chars().any(|c| c.is_ascii_uppercase())
|| token.ends_with('!'))
}
fn is_symbol_name(token: &str) -> bool {
let trimmed = token.trim_end_matches('!');
if trimmed.is_empty() || trimmed.len() > 80 {
return false;
}
let mut chars = trimmed.chars();
let Some(first) = chars.next() else {
return false;
};
if !(first.is_ascii_alphabetic() || first == '_') {
return false;
}
chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}
fn is_concept_entity(token: &str) -> bool {
let lower = token.to_ascii_lowercase();
matches!(
lower.as_str(),
"fts5"
| "bm25"
| "sqlite"
| "wal"
| "jsonl"
| "gguf"
| "compaction"
| "tombstone"
| "provenance"
| "supersession"
)
}
fn push_candidate(
kind: &'static str,
value: &str,
origin: &'static str,
out: &mut Vec<EntityCandidate>,
seen: &mut HashSet<String>,
) {
let value = clipped_chars(value.trim(), MAX_ENTITY_CHARS);
if value.len() < 2 {
return;
}
let key = format!("{kind}\0{}\0{origin}", normalize_entity_value(&value));
if !seen.insert(key) {
return;
}
out.push(EntityCandidate {
kind,
value,
origin,
});
}
pub(super) fn normalize_entity_value(value: &str) -> String {
value.trim().to_ascii_lowercase()
}