use tree_sitter::{Language, Node, Parser};
#[derive(Debug, Clone, serde::Serialize)]
pub struct CodeScope {
pub kind: String,
pub name: String,
pub start_line: usize, pub end_line: usize, }
pub fn get_language(file_ext: &str) -> Option<Language> {
match file_ext {
"rs" => Some(tree_sitter_rust::LANGUAGE.into()),
"ts" | "tsx" => Some(tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into()),
"js" | "jsx" => Some(tree_sitter_javascript::LANGUAGE.into()),
"py" => Some(tree_sitter_python::LANGUAGE.into()),
"go" => Some(tree_sitter_go::LANGUAGE.into()),
"java" => Some(tree_sitter_java::LANGUAGE.into()),
"scala" | "sc" => Some(tree_sitter_scala::LANGUAGE.into()),
_ => None,
}
}
pub fn find_enclosing_scope(source: &str, file_ext: &str, line: usize) -> Option<CodeScope> {
let language = get_language(file_ext)?;
let mut parser = Parser::new();
parser.set_language(&language).ok()?;
let tree = parser.parse(source, None)?;
let target_line = line.checked_sub(1)?; let scope_node_types = scope_types_for_ext(file_ext);
let mut best: Option<(Node, &str)> = None;
walk_tree(tree.root_node(), target_line, &scope_node_types, &mut best);
let (node, kind) = best?;
let name = extract_name(node, source);
Some(CodeScope {
kind: kind.to_string(),
name: name.unwrap_or_else(|| "<anonymous>".to_string()),
start_line: node.start_position().row + 1,
end_line: node.end_position().row + 1,
})
}
fn walk_tree<'a>(
node: Node<'a>,
target_line: usize,
scope_types: &[(&str, &'a str)],
best: &mut Option<(Node<'a>, &'a str)>,
) {
let start = node.start_position().row;
let end = node.end_position().row;
if target_line < start || target_line > end {
return;
}
for (node_type, kind) in scope_types {
if node.kind() == *node_type
&& best.as_ref().is_none_or(|(b, _)| {
let b_range = b.end_position().row - b.start_position().row;
let n_range = end - start;
n_range < b_range
})
{
*best = Some((node, kind));
}
}
for i in 0..node.child_count() {
if let Some(child) = node.child(i) {
walk_tree(child, target_line, scope_types, best);
}
}
}
fn scope_types_for_ext(ext: &str) -> Vec<(&'static str, &'static str)> {
match ext {
"rs" => vec![
("function_item", "function"),
("impl_item", "impl"),
("struct_item", "struct"),
("enum_item", "enum"),
("trait_item", "trait"),
("mod_item", "module"),
],
"ts" | "tsx" | "js" | "jsx" => vec![
("function_declaration", "function"),
("method_definition", "method"),
("arrow_function", "function"),
("class_declaration", "class"),
("interface_declaration", "interface"),
],
"py" => vec![
("function_definition", "function"),
("class_definition", "class"),
],
"go" => vec![
("function_declaration", "function"),
("method_declaration", "method"),
("type_declaration", "type"),
],
"java" => vec![
("method_declaration", "method"),
("constructor_declaration", "constructor"),
("class_declaration", "class"),
("interface_declaration", "interface"),
("enum_declaration", "enum"),
],
"scala" | "sc" => vec![
("function_definition", "function"),
("class_definition", "class"),
("object_definition", "object"),
("trait_definition", "trait"),
],
_ => vec![],
}
}
fn extract_name(node: Node, source: &str) -> Option<String> {
for i in 0..node.child_count() {
if let Some(child) = node.child(i) {
let field_name = node.field_name_for_child(i as u32);
if matches!(field_name, Some("name")) {
let text = &source[child.byte_range()];
return Some(text.to_string());
}
if child.kind() == "identifier" || child.kind() == "type_identifier" {
let text = &source[child.byte_range()];
return Some(text.to_string());
}
}
}
None
}
pub fn fallback_scope(source: &str, line: usize, context_lines: usize) -> CodeScope {
let total_lines = source.lines().count();
let start = line.saturating_sub(context_lines).max(1);
let end = (line + context_lines).min(total_lines);
CodeScope {
kind: "region".to_string(),
name: format!("lines {start}-{end}"),
start_line: start,
end_line: end,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn get_language_rust() {
assert!(get_language("rs").is_some());
}
#[test]
fn get_language_typescript() {
assert!(get_language("ts").is_some());
assert!(get_language("tsx").is_some());
}
#[test]
fn get_language_javascript() {
assert!(get_language("js").is_some());
assert!(get_language("jsx").is_some());
}
#[test]
fn get_language_python() {
assert!(get_language("py").is_some());
}
#[test]
fn get_language_go() {
assert!(get_language("go").is_some());
}
#[test]
fn get_language_java() {
assert!(get_language("java").is_some());
}
#[test]
fn get_language_scala() {
assert!(get_language("scala").is_some());
assert!(get_language("sc").is_some());
}
#[test]
fn get_language_unknown_returns_none() {
assert!(get_language("txt").is_none());
assert!(get_language("md").is_none());
}
#[test]
fn find_scope_unknown_language_returns_none() {
let src = "some text content\nline two\n";
let scope = find_enclosing_scope(src, "txt", 1);
assert!(scope.is_none());
}
#[test]
fn fallback_scope_near_start() {
let src = "line1\nline2\nline3\nline4\nline5\n";
let scope = fallback_scope(src, 1, 10);
assert_eq!(scope.start_line, 1);
}
#[test]
fn fallback_scope_near_end() {
let src = "line1\nline2\nline3\nline4\nline5\n";
let scope = fallback_scope(src, 5, 10);
assert_eq!(scope.end_line, 5);
}
}