use std::collections::hash_map::Entry;
use std::path::Path;
use compact_str::CompactString;
use rustc_hash::FxHashMap;
use tree_sitter::{QueryCursor, StreamingIterator, Tree};
use super::{MAX_SYMBOLS_PER_FILE, Symbol, SymbolKind, kind_specificity};
use crate::{LanguageId, ParserPool};
struct RawTag {
name: String,
kind: SymbolKind,
line: usize,
column: usize,
name_byte: usize,
start_byte: usize,
end_byte: usize,
}
pub fn extract_symbols(source: &str, path: &str, pool: &mut ParserPool<'_>) -> Vec<Symbol> {
if u32::try_from(source.len()).is_err() {
return Vec::new();
}
let Some(id) = pool.registry().for_path(Path::new(path)) else {
return Vec::new();
};
if pool
.registry()
.get(id)
.is_none_or(|spec| spec.tags_query.is_none())
{
return Vec::new();
}
let tree = {
let Some(parser) = pool.parser(id) else {
return Vec::new();
};
match parser.parse(source, None) {
Some(tree) => tree,
None => return Vec::new(),
}
};
extract_symbols_from_tree(source, &tree, id, pool)
}
pub(crate) fn extract_symbols_from_tree(
source: &str,
tree: &Tree,
id: LanguageId,
pool: &mut ParserPool<'_>,
) -> Vec<Symbol> {
const _: () = assert!(MAX_SYMBOLS_PER_FILE <= u16::MAX as usize);
let merge_adjacent_same_name_definitions = pool
.registry()
.get(id)
.is_some_and(|spec| spec.merge_adjacent_same_name_definitions);
let Some(query) = pool.tags_query(id) else {
return Vec::new();
};
let capture_names = query.capture_names();
let bytes = source.as_bytes();
let mut raw: Vec<RawTag> = Vec::new();
let mut cursor = QueryCursor::new();
let mut matches = cursor.matches(query, tree.root_node(), bytes);
while let Some(m) = matches.next() {
let mut name_node = None;
let mut definition = None;
for capture in m.captures {
let capture_name = capture_names[capture.index as usize];
if capture_name == "name" {
name_node.get_or_insert(capture.node);
} else if let Some(kind) = SymbolKind::from_capture(capture_name) {
definition.get_or_insert((kind, capture.node));
}
}
let (Some(name_node), Some((kind, def_node))) = (name_node, definition) else {
continue;
};
let Ok(name) = name_node.utf8_text(bytes) else {
continue;
};
let name = name.trim();
if name.is_empty() {
continue;
}
let position = name_node.start_position();
raw.push(RawTag {
name: name.to_string(),
kind,
line: position.row + 1,
column: char_column(source, name_node.start_byte(), position.column),
name_byte: name_node.start_byte(),
start_byte: def_node.start_byte(),
end_byte: def_node.end_byte(),
});
if raw.len() >= MAX_SYMBOLS_PER_FILE {
break;
}
}
let mut raw = collapse_duplicate_tags(raw);
raw.sort_by(compare_raw_tags);
let mut symbols: Vec<Symbol> = Vec::with_capacity(raw.len());
let mut enclosing: Vec<usize> = Vec::new();
for tag in raw {
while enclosing.last().is_some_and(|end| *end <= tag.start_byte) {
enclosing.pop();
}
let depth = enclosing.len();
enclosing.push(tag.end_byte);
if merge_adjacent_same_name_definitions
&& let Some(previous) = symbols.last()
&& previous.name.as_str() == tag.name
&& previous.kind == tag.kind
&& usize::from(previous.depth) == depth
{
continue;
}
debug_assert!(tag.line <= u32::MAX as usize);
debug_assert!(tag.column <= u32::MAX as usize);
debug_assert!(depth <= u16::MAX as usize);
debug_assert!(tag.name_byte <= u32::MAX as usize);
debug_assert!(tag.start_byte <= u32::MAX as usize);
debug_assert!(tag.end_byte <= u32::MAX as usize);
symbols.push(Symbol {
name: CompactString::from(tag.name),
kind: tag.kind,
line: tag.line as u32,
column: tag.column as u32,
depth: depth as u16,
name_start: tag.name_byte as u32,
def_start: tag.start_byte as u32,
def_end: tag.end_byte as u32,
});
}
symbols
}
fn compare_raw_tags(a: &RawTag, b: &RawTag) -> std::cmp::Ordering {
a.start_byte
.cmp(&b.start_byte)
.then(b.end_byte.cmp(&a.end_byte))
.then(a.line.cmp(&b.line))
.then(a.name.cmp(&b.name))
.then(a.name_byte.cmp(&b.name_byte))
}
fn collapse_duplicate_tags(raw: Vec<RawTag>) -> Vec<RawTag> {
let mut best: FxHashMap<usize, RawTag> = FxHashMap::default();
best.reserve(raw.len());
for tag in raw {
match best.entry(tag.name_byte) {
Entry::Vacant(slot) => {
slot.insert(tag);
}
Entry::Occupied(mut slot) => {
let current = slot.get();
let span = tag.end_byte.saturating_sub(tag.start_byte);
let current_span = current.end_byte.saturating_sub(current.start_byte);
let better = (kind_specificity(tag.kind), span)
< (kind_specificity(current.kind), current_span);
if better {
slot.insert(tag);
}
}
}
}
best.into_values().collect()
}
fn char_column(source: &str, start_byte: usize, byte_column: usize) -> usize {
let line_start = start_byte.saturating_sub(byte_column);
if line_start >= source.len() || !source.is_char_boundary(line_start) {
return byte_column;
}
let end = start_byte.min(source.len());
if !source.is_char_boundary(end) {
return byte_column;
}
source[line_start..end].chars().count()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn raw_tag_sort_uses_name_byte_as_the_final_tie_breaker() {
let raw_tag = |name_byte| RawTag {
name: "same".to_owned(),
kind: SymbolKind::Function,
line: 7,
column: 0,
name_byte,
start_byte: 10,
end_byte: 20,
};
let mut tags = vec![raw_tag(19), raw_tag(11), raw_tag(15)];
tags.sort_by(compare_raw_tags);
assert_eq!(
tags.into_iter()
.map(|tag| tag.name_byte)
.collect::<Vec<_>>(),
[11, 15, 19]
);
}
}