use std::sync::{Arc, OnceLock, RwLock};
use ahash::AHashMap;
use streaming_iterator::StreamingIterator;
use tree_sitter::{Node, Query, QueryCursor, Tree};
use crate::lang::{self, LangError, LangId};
fn vendored_locals(lang: LangId) -> Option<&'static str> {
match lang {
"python" => Some(include_str!("../queries/python-locals.scm")),
"typescript" => Some(include_str!("../queries/typescript-locals.scm")),
"tsx" => Some(include_str!("../queries/tsx-locals.scm")),
"go" => Some(include_str!("../queries/go-locals.scm")),
_ => None,
}
}
type CachedQuery = Option<Arc<Query>>;
static LOCALS_QUERIES: OnceLock<RwLock<AHashMap<LangId, CachedQuery>>> = OnceLock::new();
pub fn locals_query(lang: LangId) -> Result<CachedQuery, LangError> {
let lock = LOCALS_QUERIES.get_or_init(|| RwLock::new(AHashMap::new()));
if let Some(slot) = lock.read().expect("locals query pool poisoned").get(&lang) {
return Ok(slot.as_ref().map(Arc::clone));
}
let source: Option<&'static str> =
vendored_locals(lang).or_else(|| tree_sitter_language_pack::get_locals_query(lang));
let cached = match source {
Some(src) => {
let ts_lang = lang::language(lang)?;
let query = Query::new(&ts_lang, src).map_err(|e| LangError::QueryCompile {
lang,
kind: "locals",
msg: format!("{e}"),
})?;
Some(Arc::new(query))
}
None => None,
};
lock.write()
.expect("locals query pool poisoned")
.insert(lang, cached.as_ref().map(Arc::clone));
Ok(cached)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct ResolvedSpan {
use_end: u32,
def_start: u32,
def_end: u32,
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct LocalBindings {
ref_to_def: AHashMap<u32, ResolvedSpan>,
}
impl LocalBindings {
pub fn resolved_def(&self, ref_start_byte: u32) -> Option<u32> {
self.ref_to_def.get(&ref_start_byte).map(|span| span.def_start)
}
pub fn is_local(&self, ref_start_byte: u32) -> bool {
self.ref_to_def.contains_key(&ref_start_byte)
}
pub fn edges(&self) -> impl Iterator<Item = (u32, u32)> + '_ {
self.ref_to_def
.iter()
.map(|(&use_start, span)| (use_start, span.def_start))
}
pub fn edges_spanned(&self) -> impl Iterator<Item = (u32, u32, u32, u32)> + '_ {
self.ref_to_def
.iter()
.map(|(&use_start, span)| (use_start, span.use_end, span.def_start, span.def_end))
}
pub fn is_empty(&self) -> bool {
self.ref_to_def.is_empty()
}
pub fn len(&self) -> usize {
self.ref_to_def.len()
}
}
pub fn resolve_locals(lang: LangId, tree: &Tree, source: &[u8]) -> Result<LocalBindings, LangError> {
let Some(query) = locals_query(lang)? else {
return Ok(LocalBindings::default());
};
Ok(build_bindings(&query, tree.root_node(), source))
}
fn build_bindings(query: &Query, root: Node, source: &[u8]) -> LocalBindings {
let names = query.capture_names();
debug_assert!(source.len() <= u32::MAX as usize, "source exceeds u32 byte range");
let src_len = source.len() as u32;
let mut scopes: Vec<(u32, u32)> = vec![(0, src_len)];
let mut defs: Vec<(&[u8], u32, u32)> = Vec::new();
let mut refs: Vec<(&[u8], u32, u32)> = Vec::new();
let mut cursor = QueryCursor::new();
let mut iter = cursor.matches(query, root, source);
while let Some(m) = iter.next() {
for cap in m.captures {
let cname = names[cap.index as usize];
let node = cap.node;
let start = node.start_byte() as u32;
let end = node.end_byte() as u32;
if cname == "local.scope" {
scopes.push((start, end));
} else if cname.starts_with("local.definition") {
let name = &source[node.start_byte()..node.end_byte()];
defs.push((name, start, end));
} else if cname.starts_with("local.reference") {
let name = &source[node.start_byte()..node.end_byte()];
refs.push((name, start, end));
}
}
}
LocalBindings {
ref_to_def: resolve_bindings(&scopes, &defs, &refs),
}
}
fn resolve_bindings(
scopes: &[(u32, u32)],
defs: &[(&[u8], u32, u32)],
refs: &[(&[u8], u32, u32)],
) -> AHashMap<u32, ResolvedSpan> {
let mut defs_by_scope: AHashMap<usize, Vec<usize>> = AHashMap::new();
for (di, &(_, ds, de)) in defs.iter().enumerate() {
let owner = innermost_scope(scopes, ds, de);
defs_by_scope.entry(owner).or_default().push(di);
}
let mut scope_order: Vec<usize> = (0..scopes.len()).collect();
scope_order.sort_by_key(|&i| scopes[i].1.saturating_sub(scopes[i].0));
let mut ref_to_def: AHashMap<u32, ResolvedSpan> = AHashMap::new();
for &(rname, rs, re) in refs {
for &sc in &scope_order {
let (s, e) = scopes[sc];
if s <= rs
&& re <= e
&& let Some(dis) = defs_by_scope.get(&sc)
&& let Some(&di) = dis.iter().find(|&&di| defs[di].0 == rname)
{
let (_, def_start, def_end) = defs[di];
ref_to_def.insert(
rs,
ResolvedSpan {
use_end: re,
def_start,
def_end,
},
);
break;
}
}
}
ref_to_def
}
fn innermost_scope(scopes: &[(u32, u32)], start: u32, end: u32) -> usize {
debug_assert!(!scopes.is_empty(), "scopes must contain at least the root scope");
let mut best = 0usize;
let mut best_area = scopes[0].1 - scopes[0].0;
for (i, &(s, e)) in scopes.iter().enumerate() {
if s <= start && end <= e {
let area = e - s;
if area < best_area {
best = i;
best_area = area;
}
}
}
best
}
#[cfg(test)]
mod tests {
use super::*;
type NamedSpan = (&'static [u8], u32, u32);
type Fixture = (Vec<(u32, u32)>, Vec<NamedSpan>, Vec<NamedSpan>);
fn fixture() -> Fixture {
let scopes = vec![(0u32, 200u32), (10, 100), (50, 90)];
let defs: Vec<(&[u8], u32, u32)> = vec![
(b"a".as_slice(), 20, 21),
(b"x".as_slice(), 30, 31),
(b"inner".as_slice(), 45, 50),
(b"b".as_slice(), 60, 61),
];
let refs: Vec<(&[u8], u32, u32)> = vec![
(b"x".as_slice(), 70, 71),
(b"b".as_slice(), 74, 75),
(b"inner".as_slice(), 95, 100),
(b"a".as_slice(), 96, 97),
];
(scopes, defs, refs)
}
#[test]
fn reference_binds_to_nearest_enclosing_definition() {
let (scopes, defs, refs) = fixture();
let map = resolve_bindings(&scopes, &defs, &refs);
let x = map.get(&70).expect("x should bind");
assert_eq!(
(x.def_start, x.def_end),
(30, 31),
"x should bind to outer let x@30..31"
);
assert_eq!(x.use_end, 71, "x ref end must be the identifier end, not the start");
assert_eq!(
map.get(&74).map(|s| s.def_start),
Some(60),
"b should bind to inner param b@60"
);
let inner = map.get(&95).expect("inner should bind");
assert_eq!(
(inner.def_start, inner.def_end),
(45, 50),
"inner should bind to function inner@45..50"
);
assert_eq!(inner.use_end, 100, "inner ref end must be the identifier end");
assert_eq!(
map.get(&96).map(|s| s.def_start),
Some(20),
"a should bind to outer param a@20"
);
}
#[test]
fn shadowing_prefers_inner_scope() {
let scopes = vec![(0u32, 200u32), (10, 100), (50, 90)];
let defs: Vec<(&[u8], u32, u32)> = vec![(b"v".as_slice(), 30, 31), (b"v".as_slice(), 55, 56)];
let refs: Vec<(&[u8], u32, u32)> = vec![(b"v".as_slice(), 70, 71)];
let map = resolve_bindings(&scopes, &defs, &refs);
assert_eq!(
map.get(&70).map(|s| s.def_start),
Some(55),
"inner v@55 must shadow outer v@30"
);
}
#[test]
fn unbound_reference_is_absent() {
let (scopes, defs, _) = fixture();
let refs: Vec<(&[u8], u32, u32)> = vec![(b"fetch".as_slice(), 70, 75)];
let map = resolve_bindings(&scopes, &defs, &refs);
assert!(map.get(&70).is_none(), "free name `fetch` must stay unresolved");
assert!(!LocalBindings { ref_to_def: map }.is_local(70));
}
#[test]
fn module_level_binding_resolves_to_root() {
let scopes = vec![(0u32, 200u32), (10, 100)];
let defs: Vec<(&[u8], u32, u32)> = vec![(b"helper".as_slice(), 5, 11)];
let refs: Vec<(&[u8], u32, u32)> = vec![(b"helper".as_slice(), 50, 56)];
let map = resolve_bindings(&scopes, &defs, &refs);
assert_eq!(
map.get(&50).map(|s| s.def_start),
Some(5),
"root-level helper must still bind"
);
}
fn try_parse(lang: LangId, bytes: &[u8]) -> Option<Tree> {
use crate::lang::{self, ParseOutcome};
if !matches!(locals_query(lang), Ok(Some(_))) {
return None;
}
match lang::with_parser(lang, |p| lang::parse_with_default_timeout(p, bytes)) {
Ok(ParseOutcome::Ok(tree)) => Some(tree),
_ => None,
}
}
fn spanned_for(bindings: &LocalBindings, use_start: u32) -> Option<(u32, u32, u32, u32)> {
bindings.edges_spanned().find(|&(us, ..)| us == use_start)
}
#[test]
fn resolve_locals_binds_python_local() {
let lang = "python";
let src = "def outer(a):\n x = 1\n return x + a\n";
let bytes = src.as_bytes();
let Some(tree) = try_parse(lang, bytes) else {
eprintln!("skip {lang}: grammar or locals query unavailable in this environment");
return;
};
let bindings = resolve_locals(lang, &tree, bytes).expect("resolve_locals must not error");
let x_def = src.find("x = 1").expect("x def present");
let a_def = src.find("(a)").expect("a param present") + "(".len();
let x_use = src.rfind("x + a").expect("x use present");
let a_use = src.rfind("x + a").expect("a use present") + "x + ".len();
assert_eq!(
bindings.resolved_def(x_use as u32),
Some(x_def as u32),
"python: `x` use must bind to local `x = 1`"
);
assert_eq!(
bindings.resolved_def(a_use as u32),
Some(a_def as u32),
"python: `a` use must bind to param `a`"
);
let (us, ue, ds, de) = spanned_for(&bindings, x_use as u32).expect("python: x edge spanned");
assert_eq!(
ue - us,
"x".len() as u32,
"python: x ref end must be the identifier end"
);
assert_eq!(
de - ds,
"x".len() as u32,
"python: x def end must be the identifier end"
);
}
#[test]
fn resolve_locals_binds_go_local() {
let lang = "go";
let src = "package main\n\nfunc outer(a int) int {\n\tx := 1\n\treturn x + a\n}\n";
let bytes = src.as_bytes();
let Some(tree) = try_parse(lang, bytes) else {
eprintln!("skip {lang}: grammar or locals query unavailable in this environment");
return;
};
let bindings = resolve_locals(lang, &tree, bytes).expect("resolve_locals must not error");
let x_def = src.find("x :=").expect("x def present");
let a_def = src.find("a int").expect("a param present");
let x_use = src.rfind("x + a").expect("x use present");
let a_use = src.rfind("x + a").expect("a use present") + "x + ".len();
assert_eq!(
bindings.resolved_def(x_use as u32),
Some(x_def as u32),
"go: `x` use must bind to local `x := 1`"
);
assert_eq!(
bindings.resolved_def(a_use as u32),
Some(a_def as u32),
"go: `a` use must bind to param `a`"
);
let (us, ue, ds, de) = spanned_for(&bindings, x_use as u32).expect("go: x edge spanned");
assert_eq!(ue - us, "x".len() as u32, "go: x ref end must be the identifier end");
assert_eq!(de - ds, "x".len() as u32, "go: x def end must be the identifier end");
}
#[test]
fn resolve_locals_binds_typescript_local() {
let lang = "typescript";
let src = "function outer(a: number): number {\n const x = 1;\n return x + a;\n}\n";
let bytes = src.as_bytes();
let Some(tree) = try_parse(lang, bytes) else {
eprintln!("skip {lang}: grammar or locals query unavailable in this environment");
return;
};
let bindings = resolve_locals(lang, &tree, bytes).expect("resolve_locals must not error");
let x_def = src.find("x = 1").expect("x def present");
let a_def = src.find("a: number").expect("a param present");
let x_use = src.rfind("x + a").expect("x use present");
let a_use = src.rfind("x + a").expect("a use present") + "x + ".len();
assert_eq!(
bindings.resolved_def(x_use as u32),
Some(x_def as u32),
"typescript: `x` use must bind to local `const x`"
);
assert_eq!(
bindings.resolved_def(a_use as u32),
Some(a_def as u32),
"typescript: `a` use must bind to param `a`"
);
let (us, ue, ds, de) = spanned_for(&bindings, x_use as u32).expect("typescript: x edge spanned");
assert_eq!(
ue - us,
"x".len() as u32,
"typescript: x ref end must be the identifier end"
);
assert_eq!(
de - ds,
"x".len() as u32,
"typescript: x def end must be the identifier end"
);
}
#[test]
fn resolve_locals_binds_tsx_local() {
let lang = "tsx";
let src = "function outer(a: number): number {\n const x = 1;\n return x + a;\n}\n";
let bytes = src.as_bytes();
let Some(tree) = try_parse(lang, bytes) else {
eprintln!("skip {lang}: grammar or locals query unavailable in this environment");
return;
};
let bindings = resolve_locals(lang, &tree, bytes).expect("resolve_locals must not error");
let x_def = src.find("x = 1").expect("x def present");
let a_def = src.find("a: number").expect("a param present");
let x_use = src.rfind("x + a").expect("x use present");
let a_use = src.rfind("x + a").expect("a use present") + "x + ".len();
assert_eq!(
bindings.resolved_def(x_use as u32),
Some(x_def as u32),
"tsx: `x` use must bind to local `const x`"
);
assert_eq!(
bindings.resolved_def(a_use as u32),
Some(a_def as u32),
"tsx: `a` use must bind to param `a`"
);
let (us, ue, ds, de) = spanned_for(&bindings, x_use as u32).expect("tsx: x edge spanned");
assert_eq!(ue - us, "x".len() as u32, "tsx: x ref end must be the identifier end");
assert_eq!(de - ds, "x".len() as u32, "tsx: x def end must be the identifier end");
}
#[test]
fn resolve_locals_binds_real_javascript() {
use crate::lang::{self, ParseOutcome};
let lang = "javascript";
let Ok(Some(_)) = locals_query(lang) else {
return; };
let src = "function outer(a) {\n let x = 1;\n return function () { return x + a; };\n}\n";
let bytes = src.as_bytes();
let tree = match lang::with_parser(lang, |p| lang::parse_with_default_timeout(p, bytes)) {
Ok(ParseOutcome::Ok(t)) => t,
_ => return, };
let bindings = resolve_locals(lang, &tree, bytes).expect("resolve_locals must not error");
let x_def = src.find("let x").unwrap() + "let ".len();
let a_def = src.find("(a)").unwrap() + "(".len();
let x_use = src.rfind("x + a").unwrap();
let a_use = src.rfind("x + a").unwrap() + "x + ".len();
assert_eq!(
bindings.resolved_def(x_use as u32),
Some(x_def as u32),
"inner `x` use must bind to outer `let x`"
);
assert_eq!(
bindings.resolved_def(a_use as u32),
Some(a_def as u32),
"inner `a` use must bind to param `a`"
);
let (us, ue, ds, de) = spanned_for(&bindings, x_use as u32).expect("javascript: x edge spanned");
assert_eq!(
ue - us,
"x".len() as u32,
"javascript: x ref end must be the identifier end"
);
assert_eq!(
de - ds,
"x".len() as u32,
"javascript: x def end must be the identifier end"
);
}
}