use sinter_core::{
Edge, Evidence, FieldBinding, Node, NodeId, Reference, Relation, Span, SymbolKind,
};
use streaming_iterator::StreamingIterator;
use tree_sitter::{Node as TsNode, Parser, Query, QueryCursor};
const DOC_CAP_BYTES: usize = 4096;
use sinter_core::FileFacts;
use crate::language::LanguageSpec;
#[derive(Debug, thiserror::Error)]
pub enum ExtractError {
#[error("bad grammar or query for {language}: {message}")]
Query {
language: &'static str,
message: String,
},
#[error("parser returned no tree for {0}")]
Parse(String),
}
pub struct Extractor {
spec: &'static LanguageSpec,
parser: Parser,
query: Query,
inline: Option<(Parser, Query)>,
}
struct RawEntry {
start: usize,
end: usize,
name: String,
kind: Option<SymbolKind>,
qualifier: Option<String>,
signature: String,
doc: Option<String>,
}
struct RawRef {
start: usize,
end: usize,
name: String,
path: Option<String>,
alias: Option<String>,
relation: Relation,
}
struct RawLocal {
start: usize,
end: usize,
name: String,
type_name: Option<String>,
}
#[derive(Default)]
struct Collected {
refs: Vec<RawRef>,
locals: Vec<RawLocal>,
fields: Vec<(usize, usize, String, String)>,
embeds: Vec<(usize, usize, String)>,
trait_impls: Vec<(usize, usize, String)>,
alias_spans: Vec<(usize, usize)>,
def_name_spans: Vec<(usize, usize)>,
docs: Vec<(usize, usize, String)>,
}
impl Extractor {
pub fn new(spec: &'static LanguageSpec) -> Result<Self, ExtractError> {
let language = (spec.grammar)();
let mut parser = Parser::new();
parser
.set_language(&language)
.map_err(|e| ExtractError::Query {
language: spec.name,
message: e.to_string(),
})?;
let query = Query::new(&language, spec.query_source).map_err(|e| ExtractError::Query {
language: spec.name,
message: e.to_string(),
})?;
let inline = spec
.inline
.map(|i| {
let language = (i.grammar)();
let mut parser = Parser::new();
parser
.set_language(&language)
.map_err(|e| (spec.name, e.to_string()))?;
let query = Query::new(&language, i.query_source)
.map_err(|e| (spec.name, e.to_string()))?;
Ok((parser, query))
})
.transpose()
.map_err(
|(language, message): (&'static str, String)| ExtractError::Query {
language,
message,
},
)?;
Ok(Self {
spec,
parser,
query,
inline,
})
}
pub fn extract(&mut self, file: &str, source: &str) -> Result<FileFacts, ExtractError> {
let tree = self
.parser
.parse(source, None)
.ok_or_else(|| ExtractError::Parse(file.to_string()))?;
let root = tree.root_node();
let mut entries = Vec::new();
let mut collected = Collected::default();
collect(
&self.query,
self.spec,
root,
source,
&mut entries,
&mut collected,
);
if let (Some((parser, query)), Some(ispec)) = (&mut self.inline, self.spec.inline) {
let ranges = container_ranges(root, ispec.container_kinds);
if !ranges.is_empty() {
parser.set_included_ranges(&ranges).map_err(|e| {
ExtractError::Parse(format!("{} (inline ranges: {e})", self.spec.name))
})?;
let inline_tree = parser
.parse(source, None)
.ok_or_else(|| ExtractError::Parse(format!("{} (inline)", self.spec.name)))?;
collect(
query,
self.spec,
inline_tree.root_node(),
source,
&mut entries,
&mut collected,
);
}
}
entries.sort_by_key(|e| (e.start, usize::MAX - e.end));
let mut doc_set = vec![false; entries.len()];
for (d_start, d_end, text) in &collected.docs {
let owner = entries
.iter_mut()
.enumerate()
.filter(|(_, e)| e.kind.is_some() && e.start <= *d_start && *d_end <= e.end)
.min_by_key(|(_, e)| e.end - e.start);
if let Some((i, entry)) = owner {
let cleaned: Vec<&str> = text.lines().map(str::trim).collect();
let trimmed = cleaned.join("\n");
let trimmed = trimmed.trim_matches('\n');
if trimmed.is_empty() {
continue;
}
entry.doc = match entry.doc.take().filter(|_| doc_set[i]) {
Some(prev) if prev.len() >= DOC_CAP_BYTES => Some(prev),
Some(prev) => Some(format!("{prev}\n\n{trimmed}")),
None => Some(trimmed.to_string()),
};
doc_set[i] = true;
}
}
entries.dedup_by(|b, a| {
let same = a.start == b.start && a.end == b.end && a.name == b.name;
if same && a.kind == Some(SymbolKind::Variable) && b.kind.is_some() {
a.kind = b.kind;
}
same
});
let file_id = NodeId::new(file);
let mut nodes = vec![Node {
id: file_id.clone(),
kind: SymbolKind::File,
name: file.to_string(),
file: file.to_string(),
span: Span {
start: 0,
end: source.len().max(1) as u64,
},
signature: String::new(),
doc: None,
}];
let mut contains = Vec::new();
let mut stack: Vec<(usize, String, Option<NodeId>)> = Vec::new();
let mut def_spans: Vec<(usize, usize, NodeId, SymbolKind)> = Vec::new();
for entry in &entries {
while stack.last().is_some_and(|(end, _, _)| *end <= entry.start) {
stack.pop();
}
let mut path: Vec<&str> = stack.iter().map(|(_, name, _)| name.as_str()).collect();
if let Some(q) = &entry.qualifier {
path.push(q);
}
path.push(&entry.name);
let qualified = path.join("::");
let scope_segment = entry
.qualifier
.as_ref()
.map_or(entry.name.clone(), |q| format!("{q}::{}", entry.name));
let id = if let Some(kind) = entry.kind {
let id = NodeId::new(format!("{file}#{qualified}@{}", entry.start));
let parent = stack
.iter()
.rev()
.find_map(|(_, _, id)| id.clone())
.unwrap_or_else(|| file_id.clone());
nodes.push(Node {
id: id.clone(),
kind,
name: entry.name.clone(),
file: file.to_string(),
span: Span {
start: entry.start as u64,
end: entry.end as u64,
},
signature: entry.signature.clone(),
doc: entry.doc.clone(),
});
contains.push(Edge {
src: parent,
dst: id.clone(),
relation: Relation::Contains,
evidence: Evidence::Structural,
confidence: Evidence::Structural.confidence(),
site: None,
});
def_spans.push((entry.start, entry.end, id.clone(), kind));
Some(id)
} else {
None
};
stack.push((entry.end, scope_segment, id));
}
let mut refs = collected.refs;
refs.retain(|r| !collected.def_name_spans.contains(&(r.start, r.end)));
if self.spec.name == "rust" {
refs.retain(|r| {
r.path.is_some()
|| !RUST_PRELUDE.contains(&r.name.as_str())
|| entries.iter().any(|e| e.name == r.name)
});
}
refs.sort_by_key(|r| (r.start, r.end, r.path.is_none()));
refs.dedup_by_key(|r| (r.start, r.end));
let references = refs
.into_iter()
.map(|r| {
let enclosing = def_spans
.iter()
.filter(|(s, e, _, _)| *s <= r.start && r.end <= *e)
.min_by_key(|(s, e, _, _)| e - s)
.map(|(_, _, id, _)| id.clone());
Reference {
file: file.to_string(),
name: r.name,
path: r.path,
relation: r.relation,
span: Span {
start: r.start as u64,
end: r.end as u64,
},
enclosing,
alias: r.alias,
}
})
.collect();
let alias_spans = collected.alias_spans;
let embeds = collected
.embeds
.iter()
.filter_map(|(start, end, type_name)| {
let owner = def_spans
.iter()
.filter(|(s, e, _, _)| s <= start && end <= e)
.min_by_key(|(s, e, _, _)| e - s)
.map(|(_, _, id, _)| id.clone())?;
Some(sinter_core::Embed {
owner,
type_name: type_name.clone(),
})
})
.collect();
let mut raw_locals = collected
.locals
.into_iter()
.filter(|l| !alias_spans.contains(&(l.start, l.end)))
.collect::<Vec<_>>();
raw_locals.sort_by_key(|l| (l.start, l.end, l.name.clone(), l.type_name.is_none()));
raw_locals.dedup_by(|b, a| a.start == b.start && a.end == b.end && a.name == b.name);
let locals = raw_locals
.into_iter()
.map(|l| {
let scope_end = def_spans
.iter()
.filter(|(s, e, _, _)| *s <= l.start && l.end <= *e)
.min_by_key(|(s, e, _, _)| e - s)
.map_or(source.len() as u64, |(_, e, _, _)| *e as u64);
sinter_core::LocalBinding {
file: file.to_string(),
name: l.name,
span: Span {
start: l.start as u64,
end: l.end as u64,
},
scope_end,
type_name: l.type_name,
}
})
.collect();
let fields = collected
.fields
.iter()
.filter_map(|(start, end, name, type_name)| {
let owner = def_spans
.iter()
.filter(|(s, e, _, kind)| {
*s <= *start
&& *end <= *e
&& matches!(
kind,
SymbolKind::Struct
| SymbolKind::Class
| SymbolKind::Trait
| SymbolKind::Interface
)
})
.min_by_key(|(s, e, _, _)| e - s)
.map(|(_, _, id, _)| id.clone())?;
Some(FieldBinding {
owner,
name: name.clone(),
type_name: type_name.clone(),
})
})
.collect();
let trait_impls = collected
.trait_impls
.iter()
.map(|(start, end, trait_name)| sinter_core::TraitImpl {
file: file.to_string(),
trait_name: trait_name.clone(),
span: Span {
start: *start as u64,
end: *end as u64,
},
})
.collect();
let scopes = crate::scope::node_scopes(self.spec.name, source, &nodes);
let body_terms = crate::body_terms::body_terms(source, &nodes, &scopes);
Ok(FileFacts {
file: file.to_string(),
content_hash: blake3::hash(source.as_bytes()).to_hex().to_string(),
has_syntax_errors: root.has_error(),
nodes,
contains,
references,
locals,
fields,
embeds,
trait_impls,
scopes,
body_terms,
})
}
}
const RUST_PRELUDE: &[&str] = &[
"Ok", "Err", "Some", "None", "Vec", "String", "Box", "Option", "Result", "Self", "bool",
"char", "str", "u8", "u16", "u32", "u64", "u128", "usize", "i8", "i16", "i32", "i64", "i128",
"isize", "f32", "f64",
];
fn container_ranges(root: TsNode<'_>, kinds: &[&str]) -> Vec<tree_sitter::Range> {
let mut out = Vec::new();
let mut stack = vec![root];
while let Some(node) = stack.pop() {
if kinds.contains(&node.kind()) {
out.push(node.range());
} else {
for i in (0..node.child_count()).rev() {
stack.extend(node.child(i));
}
}
}
out.sort_by_key(|r| r.start_byte);
out
}
fn collect(
query: &Query,
spec: &LanguageSpec,
root: TsNode<'_>,
source: &str,
entries: &mut Vec<RawEntry>,
out: &mut Collected,
) {
{
let mut cursor = QueryCursor::new();
let mut matches = cursor.matches(query, root, source.as_bytes());
while let Some(m) = matches.next() {
let mut def: Option<(TsNode, SymbolKind)> = None;
let mut scope: Option<TsNode> = None;
let mut name: Option<TsNode> = None;
let mut qualifier: Option<TsNode> = None;
let mut reference: Option<(TsNode, Relation)> = None;
let mut refpath: Option<TsNode> = None;
let mut import_path: Option<TsNode> = None;
let mut import_module: Option<TsNode> = None;
let mut import_name: Option<TsNode> = None;
let mut import_alias: Option<TsNode> = None;
let mut import_star = false;
let mut match_locals: Vec<TsNode> = Vec::new();
let mut local_type: Option<TsNode> = None;
let mut field_name: Option<TsNode> = None;
let mut field_type: Option<TsNode> = None;
let mut trait_name: Option<TsNode> = None;
let mut trait_impl: Option<TsNode> = None;
for cap in m.captures {
let cap_name = &query.capture_names()[cap.index as usize];
if let Some(kind_str) = cap_name.strip_prefix("def.") {
if let Some(kind) = SymbolKind::from_str_opt(kind_str) {
def = Some((cap.node, kind));
}
} else if let Some(rel) = cap_name.strip_prefix("ref.") {
let relation = match rel {
"use" => Relation::Uses,
_ => Relation::Calls,
};
reference = Some((cap.node, relation));
} else {
match *cap_name {
"scope" => scope = Some(cap.node),
"name" => name = Some(cap.node),
"qualifier" => qualifier = Some(cap.node),
"refpath" => refpath = Some(cap.node),
"import" => import_path = Some(cap.node),
"import.module" => import_module = Some(cap.node),
"import.name" => import_name = Some(cap.node),
"import.alias" => import_alias = Some(cap.node),
"import.star" => import_star = true,
"local" => match_locals.push(cap.node),
"local.type" => local_type = Some(cap.node),
"field.name" => field_name = Some(cap.node),
"field.type" => field_type = Some(cap.node),
"trait" => trait_name = Some(cap.node),
"trait.impl" => trait_impl = Some(cap.node),
"doc" => out.docs.push((
cap.node.start_byte(),
cap.node.end_byte(),
text(cap.node, source).to_string(),
)),
"embed" => out.embeds.push((
cap.node.start_byte(),
cap.node.end_byte(),
text(cap.node, source).to_string(),
)),
_ => {}
}
}
}
let sep = spec.path_separators.first().copied().unwrap_or(".");
if let (Some(t), Some(block)) = (trait_name, trait_impl) {
out.trait_impls.push((
block.start_byte(),
block.end_byte(),
text(t, source).to_string(),
));
}
if let (Some(name), Some(ty)) = (field_name, field_type) {
out.fields.push((
name.start_byte(),
ty.end_byte(),
text(name, source).to_string(),
text(ty, source).to_string(),
));
}
if let Some(a) = import_alias {
out.alias_spans.push((a.start_byte(), a.end_byte()));
}
let alias = import_alias.map(|a| text(a, source).to_string());
for l in &match_locals {
out.locals.push(RawLocal {
start: l.start_byte(),
end: l.end_byte(),
name: text(*l, source).to_string(),
type_name: local_type.map(|t| text(t, source).to_string()),
});
}
if let Some(path_node) = import_path {
out.refs.push(RawRef {
start: path_node.start_byte(),
end: path_node.end_byte(),
name: text(path_node, source)
.trim_matches(['"', '\'', '`'])
.to_string(),
path: None,
alias: alias.or_else(|| import_star.then(|| "*".to_string())),
relation: Relation::Imports,
});
} else if let (Some(module), Some(item)) = (import_module, import_name) {
let module_text = text(module, source).trim_matches(['"', '\'', '`']);
out.refs.push(RawRef {
start: module.start_byte().min(item.start_byte()),
end: item.end_byte().max(module.end_byte()),
name: format!("{module_text}{sep}{}", text(item, source)),
path: None,
alias,
relation: Relation::Imports,
});
} else if let (Some(module), true) = (import_module, import_star) {
let module_text = text(module, source).trim_matches(['"', '\'', '`']);
out.refs.push(RawRef {
start: module.start_byte(),
end: module.end_byte(),
name: format!("{module_text}{sep}*"),
path: None,
alias: Some("*".to_string()),
relation: Relation::Imports,
});
}
if let Some((node, relation)) = reference {
out.refs.push(RawRef {
start: node.start_byte(),
end: node.end_byte(),
name: text(node, source).to_string(),
path: refpath.map(|p| qualified_path(p, node, source)),
alias: None,
relation,
});
}
let container = def.map(|(n, _)| n).or(scope);
if let (Some(container), Some(name_node)) = (container, name) {
if def.is_some() {
out.def_name_spans
.push((name_node.start_byte(), name_node.end_byte()));
}
entries.push(RawEntry {
start: container.start_byte(),
end: container.end_byte(),
name: text(name_node, source).to_string(),
kind: def.map(|(_, k)| k),
qualifier: qualifier.map(|q| text(q, source).to_string()),
signature: signature(container, source),
doc: doc_comment(container, source, spec.comment_kinds, spec.doc_skip_kinds),
});
}
}
}
}
fn qualified_path(refpath: TsNode<'_>, name: TsNode<'_>, source: &str) -> String {
let receiver = refpath
.named_child(0)
.filter(|c| refpath.id() != name.id() && c.end_byte() <= name.start_byte());
let Some(receiver) = receiver else {
return text(refpath, source).to_string();
};
let sep = source[receiver.end_byte()..name.start_byte()].trim();
format!("{}{sep}{}", text(receiver, source), text(name, source))
}
fn text<'a>(node: TsNode<'_>, source: &'a str) -> &'a str {
&source[node.start_byte()..node.end_byte()]
}
fn signature(node: TsNode<'_>, source: &str) -> String {
let t = text(node, source);
let first_line = t.lines().next().unwrap_or(t);
let head = if first_line.trim_end().ends_with(':') || first_line.contains('{') {
first_line.split('{').next().unwrap_or(first_line)
} else {
let up_to_brace = t.split('{').next().unwrap_or(t);
if up_to_brace.len() == t.len() {
first_line
} else {
up_to_brace
}
};
head.split_whitespace().collect::<Vec<_>>().join(" ")
}
fn doc_comment(
node: TsNode<'_>,
source: &str,
comment_kinds: &[&str],
skip_kinds: &[&str],
) -> Option<String> {
let comments = preceding_comments(node, comment_kinds, skip_kinds).or_else(|| {
node.parent()
.and_then(|p| preceding_comments(p, comment_kinds, skip_kinds))
})?;
let mut lines = Vec::new();
for c in comments {
for line in text(c, source).lines() {
let mut l = line.trim();
for marker in ["///", "//!", "//", "/**", "/*", "*/", "--"] {
if let Some(stripped) = l.strip_prefix(marker) {
l = stripped;
break;
}
}
if let Some(rest) = l.strip_prefix('*')
&& !l.starts_with("**")
{
l = rest;
}
l = l.strip_suffix("*/").unwrap_or(l);
lines.push(l.trim());
}
}
while lines.first().is_some_and(|l| l.is_empty()) {
lines.remove(0);
}
while lines.last().is_some_and(|l| l.is_empty()) {
lines.pop();
}
if lines.is_empty() {
None
} else {
Some(lines.join("\n"))
}
}
fn preceding_comments<'t>(
node: TsNode<'t>,
comment_kinds: &[&str],
skip_kinds: &[&str],
) -> Option<Vec<TsNode<'t>>> {
let mut comments = Vec::new();
let mut cur = node.prev_named_sibling();
let mut skips = 0;
while let Some(sib) = cur {
if !comment_kinds.contains(&sib.kind()) {
if skips < 2 && comments.is_empty() && skip_kinds.contains(&sib.kind()) {
skips += 1;
cur = sib.prev_named_sibling();
continue;
}
break;
}
comments.push(sib);
cur = sib.prev_named_sibling();
}
comments.reverse();
if comments.is_empty() {
None
} else {
Some(comments)
}
}