use tree_sitter::Node;
pub const ELEMENT_QUERY: &str = r"
(function_declaration
name: (identifier) @func_name) @function
(method_declaration
name: (field_identifier) @method_name) @function
(type_spec
name: (type_identifier) @type_name
type: (struct_type)) @class
(type_spec
name: (type_identifier) @type_name
type: (interface_type)) @class
";
pub const CALL_QUERY: &str = r"
(call_expression
function: (identifier) @call)
(call_expression
function: (selector_expression field: (field_identifier) @call))
";
pub const REFERENCE_QUERY: &str = r"
(type_identifier) @type_ref
";
pub const IMPORT_QUERY: &str = r"
(import_declaration) @import_path
";
pub const DEFUSE_QUERY: &str = r"
(short_var_declaration left: (expression_list (identifier) @write.short))
(assignment_statement left: (expression_list (identifier) @write.assign))
(var_declaration (var_spec (identifier) @write.var))
(inc_statement (identifier) @writeread.inc)
(dec_statement (identifier) @writeread.dec)
(identifier) @read.usage
";
#[must_use]
pub fn extract_function_name(node: &Node, source: &str, _lang: &str) -> Option<String> {
if node.kind() != "function_declaration" && node.kind() != "method_declaration" {
return None;
}
node.child_by_field_name("name").and_then(|n| {
let end = n.end_byte();
if end <= source.len() {
Some(source[n.start_byte()..end].to_string())
} else {
None
}
})
}
#[must_use]
pub fn find_receiver_type(node: &Node, source: &str) -> Option<String> {
if node.kind() != "method_declaration" {
return None;
}
let receiver = node.child_by_field_name("receiver")?;
for i in 0..receiver.named_child_count() {
if let Some(param) = receiver.named_child(i as u32)
&& param.kind() == "parameter_declaration"
{
if let Some(type_node) = param.child_by_field_name("type") {
match type_node.kind() {
"type_identifier" => {
let end = type_node.end_byte();
if end <= source.len() {
return Some(source[type_node.start_byte()..end].to_string());
}
}
"pointer_type" => {
if let Some(inner) = (0..type_node.named_child_count())
.filter_map(|j| type_node.named_child(j as u32))
.find(|n| n.kind() == "type_identifier")
{
let end = inner.end_byte();
if end <= source.len() {
return Some(source[inner.start_byte()..end].to_string());
}
}
}
_ => {}
}
}
}
}
None
}
#[must_use]
pub fn find_method_for_receiver(
node: &Node,
source: &str,
_depth: Option<usize>,
) -> Option<String> {
if node.kind() != "method_declaration" && node.kind() != "function_declaration" {
return None;
}
node.child_by_field_name("name").and_then(|n| {
let start = n.start_byte();
let end = n.end_byte();
if end <= source.len() {
Some(source[start..end].to_string())
} else {
None
}
})
}
#[must_use]
pub fn extract_inheritance(node: &Node, source: &str) -> Vec<String> {
let mut inherits = Vec::new();
if let Some(type_field) = node.child_by_field_name("type") {
match type_field.kind() {
"struct_type" => {
for i in 0..type_field.named_child_count() {
if let Some(field_list) = type_field.named_child(i as u32)
&& field_list.kind() == "field_declaration_list"
{
for j in 0..field_list.named_child_count() {
if let Some(field) = field_list.named_child(j as u32)
&& field.kind() == "field_declaration"
&& field.child_by_field_name("name").is_none()
{
if let Some(type_node) = field.child_by_field_name("type") {
let text =
&source[type_node.start_byte()..type_node.end_byte()];
inherits.push(text.to_string());
}
}
}
}
}
}
"interface_type" => {
for i in 0..type_field.named_child_count() {
if let Some(elem) = type_field.named_child(i as u32)
&& elem.kind() == "type_elem"
{
let text = &source[elem.start_byte()..elem.end_byte()];
inherits.push(text.to_string());
}
}
}
_ => {}
}
}
inherits
}
#[cfg(all(test, feature = "lang-go"))]
mod tests {
use super::*;
use crate::DefUseKind;
use crate::parser::SemanticExtractor;
use tree_sitter::Parser;
fn parse_go(source: &str) -> tree_sitter::Tree {
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_go::LANGUAGE.into())
.expect("failed to set Go language");
parser.parse(source, None).expect("failed to parse source")
}
#[test]
fn test_extract_inheritance_struct_no_embeds() {
let source = "package p\ntype Foo struct { x int }";
let tree = parse_go(source);
let root = tree.root_node();
let type_spec = (0..root.named_child_count())
.filter_map(|i| root.named_child(i as u32))
.find_map(|n| {
if n.kind() == "type_declaration" {
(0..n.named_child_count())
.filter_map(|j| n.named_child(j as u32))
.find(|c| c.kind() == "type_spec")
} else {
None
}
})
.expect("expected type_spec node");
let result = extract_inheritance(&type_spec, source);
assert!(
result.is_empty(),
"expected no inherited types, got {:?}",
result
);
}
#[test]
fn test_find_method_for_receiver_wrong_kind() {
let source = "package p\ntype Bar struct {}";
let tree = parse_go(source);
let root = tree.root_node();
let node = root.named_child(0).expect("expected child");
let result = find_method_for_receiver(&node, source, None);
assert_eq!(result, None);
}
#[test]
fn test_extract_function_name() {
let source = "package p\nfunc Foo() {}";
let tree = parse_go(source);
let root = tree.root_node();
let func_node = (0..root.named_child_count())
.filter_map(|i| root.named_child(i as u32))
.find(|n| n.kind() == "function_declaration")
.expect("expected function_declaration");
let result = extract_function_name(&func_node, source, "go");
assert_eq!(result, Some("Foo".to_string()));
}
#[test]
fn test_extract_method_name() {
let source = "package p\nfunc (r *Receiver) Bar() {}";
let tree = parse_go(source);
let root = tree.root_node();
let method_node = (0..root.named_child_count())
.filter_map(|i| root.named_child(i as u32))
.find(|n| n.kind() == "method_declaration")
.expect("expected method_declaration");
let result = extract_function_name(&method_node, source, "go");
assert_eq!(result, Some("Bar".to_string()));
}
#[test]
fn test_extract_function_name_wrong_kind() {
let source = "package p\ntype Baz struct {}";
let tree = parse_go(source);
let root = tree.root_node();
let node = root.named_child(0).expect("expected child");
let result = extract_function_name(&node, source, "go");
assert_eq!(result, None);
}
#[test]
fn test_find_receiver_type() {
let source = "package p\nfunc (r Receiver) Foo() {}";
let tree = parse_go(source);
let root = tree.root_node();
let method_node = (0..root.named_child_count())
.filter_map(|i| root.named_child(i as u32))
.find(|n| n.kind() == "method_declaration")
.expect("expected method_declaration");
let result = find_receiver_type(&method_node, source);
assert_eq!(result, Some("Receiver".to_string()));
}
#[test]
fn test_find_receiver_type_pointer() {
let source = "package p\nfunc (r *Receiver) Foo() {}";
let tree = parse_go(source);
let root = tree.root_node();
let method_node = (0..root.named_child_count())
.filter_map(|i| root.named_child(i as u32))
.find(|n| n.kind() == "method_declaration")
.expect("expected method_declaration");
let result = find_receiver_type(&method_node, source);
assert_eq!(result, Some("Receiver".to_string()));
}
#[test]
fn test_defuse_query_write_site() {
let src = "package p\nfunc main() { x := 1 }\n";
let sites = SemanticExtractor::extract_def_use_for_file(src, "go", "x", "test.go", None);
assert!(!sites.is_empty(), "defuse sites should not be empty");
let has_write = sites.iter().any(|s| matches!(s.kind, DefUseKind::Write));
assert!(has_write, "should contain a Write DefUseSite");
}
#[test]
fn test_defuse_go_short_var_decl() {
let src = "package p\nfunc main() { x := 42 }\n";
let sites = SemanticExtractor::extract_def_use_for_file(src, "go", "x", "test.go", None);
assert!(
!sites.is_empty(),
"short var decl should produce defuse sites"
);
let has_write = sites.iter().any(|s| matches!(s.kind, DefUseKind::Write));
assert!(has_write, "short var decl should be Write");
}
#[test]
fn test_defuse_go_multi_return() {
let src = "package p\nfunc main() { a, b := f() }\nfunc f() (int, int) { return 1, 2 }\n";
let sites_a = SemanticExtractor::extract_def_use_for_file(src, "go", "a", "test.go", None);
let sites_b = SemanticExtractor::extract_def_use_for_file(src, "go", "b", "test.go", None);
assert!(
!sites_a.is_empty(),
"multi-return a should produce defuse sites"
);
assert!(
!sites_b.is_empty(),
"multi-return b should produce defuse sites"
);
let a_write = sites_a.iter().any(|s| matches!(s.kind, DefUseKind::Write));
let b_write = sites_b.iter().any(|s| matches!(s.kind, DefUseKind::Write));
assert!(a_write, "multi-return a should be Write");
assert!(b_write, "multi-return b should be Write");
}
#[test]
fn test_defuse_go_blank_identifier() {
let src =
"package p\nfunc main() { _, err := f() }\nfunc f() (int, error) { return 1, nil }\n";
let sites = SemanticExtractor::extract_def_use_for_file(src, "go", "_", "test.go", None);
if !sites.is_empty() {
let has_write = sites.iter().any(|s| matches!(s.kind, DefUseKind::Write));
assert!(has_write, "blank identifier if captured should be Write");
}
}
}