use super::engine::{ScanRefusal, list_files_matching, read_scannable_reason};
use super::glob::FileFilter;
use crate::parse::grammar::LanguageId;
use ast_grep_core::Language;
use ast_grep_core::matcher::{Pattern, PatternBuilder, PatternError};
use ast_grep_core::meta_var::MetaVariable;
use ast_grep_core::tree_sitter::{LanguageExt, StrDoc, TSLanguage};
use rayon::prelude::*;
use std::borrow::Cow;
use std::path::Path;
use std::time::Instant;
const EXPANDO_CHAR: char = 'µ';
const UNLIMITED_FILE_BUFFER: usize = 10_000;
pub const MAX_PATTERN_BYTES: usize = 4_096;
pub const NO_CAP: usize = 0;
const MAX_CANDIDATE_FILES: usize = 300_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StructuralHit {
pub line: u32,
pub col: u32,
pub start_byte: usize,
pub end_byte: usize,
pub kind: String,
pub text: String,
pub capture: Option<String>,
pub captured: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StructuralFile {
pub rel: String,
pub lang: &'static str,
pub total: usize,
pub hits: Vec<StructuralHit>,
pub capped: bool,
}
#[derive(Debug, Clone)]
pub struct StructuralOptions {
pub limit: usize,
pub offset: usize,
pub per_file_cap: usize,
pub max_match_chars: usize,
pub deadline: Option<Instant>,
}
impl Default for StructuralOptions {
fn default() -> Self {
Self {
limit: 100,
offset: 0,
per_file_cap: 20,
max_match_chars: 200,
deadline: None,
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct StructuralStats {
pub candidates: usize,
pub scanned: usize,
pub skipped_deadline: usize,
pub unreadable: usize,
pub files_matched: usize,
pub total_matches: usize,
pub files_capped: usize,
pub millis: u128,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct StructuralOutput {
pub files: Vec<StructuralFile>,
pub returned: usize,
pub has_more: bool,
pub stopped_by_deadline: bool,
pub stats: StructuralStats,
}
impl StructuralOutput {
pub fn total_matches(&self) -> usize {
self.stats.total_matches
}
}
pub struct CompiledPattern {
lang: StructuralLang,
pattern: Pattern,
name: &'static str,
}
impl std::fmt::Debug for CompiledPattern {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CompiledPattern")
.field("lang", &self.name)
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StructuralError {
UnsupportedLanguage(String),
GrammarUnavailable(&'static str),
InvalidPattern {
lang: &'static str,
detail: String,
},
UnusablePattern(String),
}
impl std::fmt::Display for StructuralError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::UnsupportedLanguage(name) => write!(
f,
"structural search does not support language '{name}' — pass one LeIndex \
has a tree-sitter grammar for (e.g. rust, python, go, typescript)"
),
Self::GrammarUnavailable(name) => write!(
f,
"language '{name}' is parsed as a document flavour and has no tree-sitter \
grammar, so structural search cannot match it"
),
Self::InvalidPattern { lang, detail } => {
write!(f, "pattern is not valid {lang} code: {detail}")
}
Self::UnusablePattern(detail) => write!(f, "unusable pattern: {detail}"),
}
}
}
impl std::error::Error for StructuralError {}
#[derive(Clone)]
struct StructuralLang {
id: LanguageId,
grammar: TSLanguage,
}
impl StructuralLang {
fn new(id: LanguageId) -> Result<Self, StructuralError> {
if is_document_flavour(id) {
return Err(StructuralError::GrammarUnavailable(
id.config().name.as_str(),
));
}
let grammar = id
.from_cache()
.map_err(|_| StructuralError::GrammarUnavailable(id.config().name.as_str()))?;
Ok(Self { id, grammar })
}
}
impl Language for StructuralLang {
fn pre_process_pattern<'q>(&self, query: &'q str) -> Cow<'q, str> {
if query.contains('$') {
Cow::Owned(query.replace('$', "\u{b5}"))
} else {
Cow::Borrowed(query)
}
}
fn expando_char(&self) -> char {
EXPANDO_CHAR
}
fn kind_to_id(&self, kind: &str) -> u16 {
self.grammar.id_for_node_kind(kind, true)
}
fn field_to_id(&self, field: &str) -> Option<u16> {
self.grammar.field_id_for_name(field).map(|f| f.get())
}
fn build_pattern(&self, builder: &PatternBuilder<'_>) -> Result<Pattern, PatternError> {
builder.build(|src| StrDoc::try_new(src, self.clone()))
}
}
impl LanguageExt for StructuralLang {
fn get_ts_language(&self) -> TSLanguage {
self.grammar.clone()
}
}
fn is_document_flavour(id: LanguageId) -> bool {
matches!(
id,
LanguageId::Markdown | LanguageId::Rst | LanguageId::Adoc | LanguageId::PlainText
)
}
const STRUCTURAL_LANGUAGES: [LanguageId; 37] = [
LanguageId::Python,
LanguageId::JavaScript,
LanguageId::TypeScript,
LanguageId::Go,
LanguageId::Rust,
LanguageId::Java,
LanguageId::Cpp,
LanguageId::CSharp,
LanguageId::Ruby,
LanguageId::Php,
LanguageId::Lua,
LanguageId::Scala,
LanguageId::C,
LanguageId::Bash,
LanguageId::Json,
LanguageId::Swift,
LanguageId::Kotlin,
LanguageId::Dart,
LanguageId::Html,
LanguageId::Css,
LanguageId::Scss,
LanguageId::Yaml,
LanguageId::Cmake,
LanguageId::Elixir,
LanguageId::Erlang,
LanguageId::Haskell,
LanguageId::Perl,
LanguageId::R,
LanguageId::Zig,
LanguageId::Graphql,
LanguageId::Hcl,
LanguageId::Make,
LanguageId::Elisp,
LanguageId::Julia,
LanguageId::D,
LanguageId::Glsl,
LanguageId::EmbeddedTemplate,
];
pub fn resolve_language(language: &str) -> Option<LanguageId> {
let key = language.trim();
if key.is_empty() {
return None;
}
if let Some(id) = LanguageId::from_extension(key) {
return Some(id);
}
STRUCTURAL_LANGUAGES.into_iter().find(|id| {
id.config().name.eq_ignore_ascii_case(key)
|| id
.config()
.extensions
.iter()
.any(|ext| ext.eq_ignore_ascii_case(key))
})
}
pub fn language_for_path(path: &Path) -> Option<LanguageId> {
match path.file_name().and_then(|name| name.to_str()) {
Some("Makefile") | Some("makefile") | Some("GNUmakefile") => {
return Some(LanguageId::Make);
}
Some("CMakeLists.txt") => return Some(LanguageId::Cmake),
_ => {}
}
let ext = path.extension()?.to_str()?;
LanguageId::from_extension(ext).filter(|id| !is_document_flavour(*id))
}
impl CompiledPattern {
pub fn new(language: &str, pattern: &str) -> Result<Self, StructuralError> {
let id = resolve_language(language)
.ok_or_else(|| StructuralError::UnsupportedLanguage(language.to_string()))?;
Self::for_language(id, pattern)
}
pub fn for_language(id: LanguageId, pattern: &str) -> Result<Self, StructuralError> {
let name = id.config().name.as_str();
let lang = StructuralLang::new(id)?;
if pattern.trim().is_empty() {
return Err(StructuralError::UnusablePattern(
"pattern must not be empty".to_string(),
));
}
if pattern.len() > MAX_PATTERN_BYTES {
return Err(StructuralError::UnusablePattern(format!(
"pattern is {} bytes; the maximum is {MAX_PATTERN_BYTES}",
pattern.len()
)));
}
let pattern = Pattern::try_new(pattern, lang.clone()).map_err(|e| {
StructuralError::InvalidPattern {
lang: name,
detail: format!("{e:?}"),
}
})?;
Ok(Self {
lang,
pattern,
name,
})
}
pub fn language(&self) -> &'static str {
self.name
}
pub fn language_id(&self) -> LanguageId {
self.lang.id
}
fn scan(&self, rel: &str, source: &str, options: &StructuralOptions) -> Option<StructuralFile> {
let doc = self.lang.ast_grep(source);
let buffer_cap = if options.per_file_cap == NO_CAP {
UNLIMITED_FILE_BUFFER
} else {
options.per_file_cap
};
let mut hits = Vec::new();
let mut total = 0usize;
for matched in doc.root().find_all(&self.pattern) {
total += 1;
if hits.len() >= buffer_cap {
continue;
}
let node = matched.get_node();
let (line, byte_column) = node.start_pos().byte_point();
let range = node.range();
let (capture, captured) = capture_of(&matched, options.max_match_chars);
hits.push(StructuralHit {
line: u32::try_from(line + 1).unwrap_or(u32::MAX),
col: u32::try_from(byte_column + 1).unwrap_or(u32::MAX),
start_byte: range.start,
end_byte: range.end,
kind: node.kind().to_string(),
text: clip(&node.text(), options.max_match_chars),
capture,
captured,
});
}
if total == 0 {
return None;
}
Some(StructuralFile {
rel: rel.to_string(),
lang: self.name,
total,
hits,
capped: total > buffer_cap,
})
}
}
fn capture_of(
matched: &ast_grep_core::NodeMatch<'_, StrDoc<StructuralLang>>,
max_chars: usize,
) -> (Option<String>, Option<String>) {
let env = matched.get_env();
let Some(name) = env.get_matched_variables().find_map(|var| match var {
MetaVariable::Capture(name, _) | MetaVariable::MultiCapture(name) => Some(name),
MetaVariable::Multiple | MetaVariable::Dropped(_) => None,
}) else {
return (None, None);
};
let captured = match env.get_match(&name) {
Some(node) => Some(clip(&node.text(), max_chars)),
None => {
let multi = env.get_multiple_matches(&name);
if multi.is_empty() {
None
} else {
let named: Vec<String> = multi
.iter()
.filter(|node| node.is_named())
.map(|node| node.text().to_string())
.collect();
let texts = if named.is_empty() {
multi
.iter()
.map(|node| node.text().to_string())
.collect::<Vec<_>>()
} else {
named
};
Some(clip(&texts.join(", "), max_chars))
}
}
};
(Some(name), captured)
}
fn clip(text: &str, max_chars: usize) -> String {
if text.chars().count() <= max_chars {
return text.to_string();
}
let mut out: String = text.chars().take(max_chars).collect();
out.push('…');
out
}
pub fn search(
root: &Path,
pattern: &CompiledPattern,
filter: &FileFilter,
options: &StructuralOptions,
) -> StructuralOutput {
let started = Instant::now();
let mut output = StructuralOutput::default();
let (rels, truncation) =
list_files_matching(root, MAX_CANDIDATE_FILES, options.deadline, |rel| {
filter.allows(rel) && language_for_path(Path::new(rel)) == Some(pattern.language_id())
});
output.stats.candidates = rels.len();
let deadline_skips = std::sync::atomic::AtomicUsize::new(0);
let unreadable = std::sync::atomic::AtomicUsize::new(0);
let policy_refusals = std::sync::atomic::AtomicUsize::new(0);
let deadline = options.deadline;
let scanned: Vec<Option<StructuralFile>> = rels
.par_iter()
.map(|rel| {
if deadline.is_some_and(|d| Instant::now() >= d) {
deadline_skips.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
return None;
}
let source = match read_scannable_reason(&root.join(rel)) {
Ok(source) => source,
Err(ScanRefusal::Unreadable) => {
unreadable.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
return None;
}
Err(ScanRefusal::NotScannable) => {
policy_refusals.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
return None;
}
};
if deadline.is_some_and(|d| Instant::now() >= d) {
deadline_skips.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
return None;
}
pattern.scan(rel, &source, options)
})
.collect();
let skipped_deadline = deadline_skips.into_inner();
let unreadable = unreadable.into_inner();
output.stats.scanned = rels
.len()
.saturating_sub(skipped_deadline + unreadable + policy_refusals.into_inner());
output.stats.skipped_deadline = skipped_deadline;
output.stats.unreadable = unreadable;
let mut files: Vec<StructuralFile> = scanned.into_iter().flatten().collect();
files.sort_by(|a, b| a.rel.cmp(&b.rel));
let mut paged: usize = 0;
let mut skipped = options.offset;
for file in files {
output.stats.total_matches += file.total;
output.stats.files_matched += 1;
if file.capped {
output.stats.files_capped += 1;
}
let available = file.hits.len();
let drop = skipped.min(available);
skipped -= drop;
paged += available - drop;
let room = options.limit.saturating_sub(output.returned);
let keep = (available - drop).min(room);
if keep == 0 {
continue;
}
let hits = file.hits[drop..drop + keep].to_vec();
output.returned += hits.len();
output.files.push(StructuralFile { hits, ..file });
}
output.stopped_by_deadline = output.stats.skipped_deadline > 0
|| truncation.deadline
|| deadline.is_some_and(|d| Instant::now() >= d);
output.has_more = output.returned < paged
|| output.stopped_by_deadline
|| truncation.truncated()
|| output.stats.unreadable > 0;
output.stats.millis = started.elapsed().as_millis();
output
}