tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Improve Rust analysis output for analyze tool (#5072)

Browse Source
This commit is contained in:
tlongwell-block
2025-10-10 09:28:26 -04:00
committed by GitHub
parent 3e9a799d8f
commit 228325ffef
6 changed files with 412 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-mcp/src/developer/analyze/formatter.rs
+97
View File
@@ -179,9 +179,106 @@ impl Formatter {
output.push('\n');
}
// References (type tracking) - only show if present
if !result.references.is_empty() {
Self::append_references(&mut output, result);
}
output
}
/// Append reference tracking information (method-to-type associations, type usage)
fn append_references(output: &mut String, result: &AnalysisResult) {
use crate::developer::analyze::types::ReferenceType;
// Group references by type
let mut method_defs = Vec::new();
let mut type_inst = Vec::new();
let mut field_types = Vec::new();
let mut var_types = Vec::new();
let mut param_types = Vec::new();
for ref_info in &result.references {
match ref_info.ref_type {
ReferenceType::MethodDefinition => method_defs.push(ref_info),
ReferenceType::TypeInstantiation => type_inst.push(ref_info),
ReferenceType::FieldType => field_types.push(ref_info),
ReferenceType::VariableType => var_types.push(ref_info),
ReferenceType::ParameterType => param_types.push(ref_info),
ReferenceType::Call | ReferenceType::Definition | ReferenceType::Import => {}
}
}
// Only show section if we have non-call references
if method_defs.is_empty()
&& type_inst.is_empty()
&& field_types.is_empty()
&& var_types.is_empty()
&& param_types.is_empty()
{
return;
}
output.push_str("\nR: ");
let mut sections = Vec::new();
// Method definitions (methods associated with types)
if !method_defs.is_empty() {
let mut method_strs: Vec<String> = method_defs
.iter()
.map(|r| {
if let Some(type_name) = &r.associated_type {
format!("{}({})", r.symbol, type_name)
} else {
r.symbol.clone()
}
})
.collect();
method_strs.sort();
method_strs.dedup();
sections.push(format!("methods[{}]", method_strs.join(" ")));
}
// Type instantiations (struct literals)
if !type_inst.is_empty() {
let mut type_names: Vec<String> = type_inst.iter().map(|r| r.symbol.clone()).collect();
type_names.sort();
type_names.dedup();
sections.push(format!("types[{}]", type_names.join(" ")));
}
// Field types (only show unique types, not all occurrences)
if !field_types.is_empty() {
let mut field_type_names: Vec<String> =
field_types.iter().map(|r| r.symbol.clone()).collect();
field_type_names.sort();
field_type_names.dedup();
sections.push(format!("fields[{}]", field_type_names.join(" ")));
}
// Variable types (only show unique types)
if !var_types.is_empty() {
let mut var_type_names: Vec<String> =
var_types.iter().map(|r| r.symbol.clone()).collect();
var_type_names.sort();
var_type_names.dedup();
sections.push(format!("vars[{}]", var_type_names.join(" ")));
}
// Parameter types (only show unique types)
if !param_types.is_empty() {
let mut param_type_names: Vec<String> =
param_types.iter().map(|r| r.symbol.clone()).collect();
param_type_names.sort();
param_type_names.dedup();
sections.push(format!("params[{}]", param_type_names.join(" ")));
}
output.push_str(&sections.join("; "));
output.push('\n');
}
/// Format directory structure with summary
pub fn format_directory_structure(
base_path: &Path,
crates/goose-mcp/src/developer/analyze/languages/mod.rs
+16 -2
View File
@@ -41,6 +41,10 @@ type ExtractFunctionNameHandler = fn(&tree_sitter::Node, &str, &str) -> Option<S
/// Takes: (receiver_node, source, ast_recursion_limit)
type FindMethodForReceiverHandler = fn(&tree_sitter::Node, &str, Option<usize>) -> Option<String>;
/// Handler for finding the receiver type from a receiver node
/// Takes: (receiver_node, source)
type FindReceiverTypeHandler = fn(&tree_sitter::Node, &str) -> Option<String>;
/// Language configuration containing all language-specific information
///
/// This struct serves as a single source of truth for language support.
@@ -61,6 +65,8 @@ pub struct LanguageInfo {
pub extract_function_name_handler: Option<ExtractFunctionNameHandler>,
/// Optional handler for finding method names from receiver nodes
pub find_method_for_receiver_handler: Option<FindMethodForReceiverHandler>,
/// Optional handler for finding receiver type from receiver nodes
pub find_receiver_type_handler: Option<FindReceiverTypeHandler>,
}
/// Get language configuration for a given language
@@ -76,15 +82,17 @@ pub fn get_language_info(language: &str) -> Option<LanguageInfo> {
function_name_kinds: &["identifier", "field_identifier", "property_identifier"],
extract_function_name_handler: None,
find_method_for_receiver_handler: None,
find_receiver_type_handler: None,
}),
"rust" => Some(LanguageInfo {
element_query: rust::ELEMENT_QUERY,
call_query: rust::CALL_QUERY,
reference_query: "",
reference_query: rust::REFERENCE_QUERY,
function_node_kinds: &["function_item", "impl_item"],
function_name_kinds: &["identifier", "field_identifier", "property_identifier"],
extract_function_name_handler: Some(rust::extract_function_name_for_kind),
find_method_for_receiver_handler: None,
find_method_for_receiver_handler: Some(rust::find_method_for_receiver),
find_receiver_type_handler: Some(rust::find_receiver_type),
}),
"javascript" | "typescript" => Some(LanguageInfo {
element_query: javascript::ELEMENT_QUERY,
@@ -98,6 +106,7 @@ pub fn get_language_info(language: &str) -> Option<LanguageInfo> {
function_name_kinds: &["identifier", "field_identifier", "property_identifier"],
extract_function_name_handler: None,
find_method_for_receiver_handler: None,
find_receiver_type_handler: None,
}),
"go" => Some(LanguageInfo {
element_query: go::ELEMENT_QUERY,
@@ -107,6 +116,7 @@ pub fn get_language_info(language: &str) -> Option<LanguageInfo> {
function_name_kinds: &["identifier", "field_identifier", "property_identifier"],
extract_function_name_handler: None,
find_method_for_receiver_handler: Some(go::find_method_for_receiver),
find_receiver_type_handler: None,
}),
"java" => Some(LanguageInfo {
element_query: java::ELEMENT_QUERY,
@@ -116,6 +126,7 @@ pub fn get_language_info(language: &str) -> Option<LanguageInfo> {
function_name_kinds: &["identifier", "field_identifier", "property_identifier"],
extract_function_name_handler: None,
find_method_for_receiver_handler: None,
find_receiver_type_handler: None,
}),
"kotlin" => Some(LanguageInfo {
element_query: kotlin::ELEMENT_QUERY,
@@ -125,6 +136,7 @@ pub fn get_language_info(language: &str) -> Option<LanguageInfo> {
function_name_kinds: &["identifier", "field_identifier", "property_identifier"],
extract_function_name_handler: None,
find_method_for_receiver_handler: None,
find_receiver_type_handler: None,
}),
"swift" => Some(LanguageInfo {
element_query: swift::ELEMENT_QUERY,
@@ -139,6 +151,7 @@ pub fn get_language_info(language: &str) -> Option<LanguageInfo> {
function_name_kinds: &["simple_identifier"],
extract_function_name_handler: Some(swift::extract_function_name_for_kind),
find_method_for_receiver_handler: None,
find_receiver_type_handler: None,
}),
"ruby" => Some(LanguageInfo {
element_query: ruby::ELEMENT_QUERY,
@@ -148,6 +161,7 @@ pub fn get_language_info(language: &str) -> Option<LanguageInfo> {
function_name_kinds: &["identifier", "field_identifier", "property_identifier"],
extract_function_name_handler: None,
find_method_for_receiver_handler: Some(ruby::find_method_for_receiver),
find_receiver_type_handler: None,
}),
_ => None,
}
crates/goose-mcp/src/developer/analyze/languages/rust.rs
+94
View File
@@ -27,6 +27,48 @@ pub const CALL_QUERY: &str = r#"
macro: (identifier) @macro.call)
"#;
/// Tree-sitter query for extracting Rust type references and usage patterns
pub const REFERENCE_QUERY: &str = r#"
; Method receivers - capture self parameters to associate methods with impl types
(self_parameter) @method.receiver
; Struct instantiation - struct literals
(struct_expression
name: (type_identifier) @struct.literal)
; Field type declarations in structs
(field_declaration
type: (type_identifier) @field.type)
; Field with reference type
(field_declaration
type: (reference_type
(type_identifier) @field.type))
; Field with generic type
(field_declaration
type: (generic_type
type: (type_identifier) @field.type))
; Variable type annotations
(let_declaration
type: (type_identifier) @var.type)
; Variable with reference type
(let_declaration
type: (reference_type
(type_identifier) @var.type))
; Function parameter types
(parameter
type: (type_identifier) @param.type)
; Parameter with reference type
(parameter
type: (reference_type
(type_identifier) @param.type))
"#;
/// Extract function name for Rust-specific node kinds
///
/// Rust has special cases like impl_item blocks that should be
@@ -48,3 +90,55 @@ pub fn extract_function_name_for_kind(
}
None
}
/// Find the method name for a method receiver node in Rust
///
/// The receiver_node is a self_parameter. This walks up to find the
/// containing function_item and returns the method name.
pub fn find_method_for_receiver(
receiver_node: &tree_sitter::Node,
source: &str,
_ast_recursion_limit: Option<usize>,
) -> Option<String> {
// Walk up to find the function_item that contains this self_parameter
let mut current = *receiver_node;
while let Some(parent) = current.parent() {
if parent.kind() == "function_item" {
// Found the function, get its name
for i in 0..parent.child_count() {
if let Some(child) = parent.child(i) {
if child.kind() == "identifier" {
return Some(source[child.byte_range()].to_string());
}
}
}
}
current = parent;
}
None
}
/// Find the receiver type for a self parameter in Rust
///
/// In Rust, self parameters are special - they don't explicitly state their type.
/// This function walks up from a self_parameter node to find the impl block
/// and extracts the type being implemented.
pub fn find_receiver_type(node: &tree_sitter::Node, source: &str) -> Option<String> {
// Walk up from self_parameter to find the impl_item
let mut current = *node;
while let Some(parent) = current.parent() {
if parent.kind() == "impl_item" {
// Find the type_identifier in the impl block
for i in 0..parent.child_count() {
if let Some(child) = parent.child(i) {
if child.kind() == "type_identifier" {
return Some(source[child.byte_range()].to_string());
}
}
}
}
current = parent;
}
None
}
crates/goose-mcp/src/developer/analyze/parser.rs
+25 -5
View File
@@ -380,11 +380,19 @@ impl ElementExtractor {
ast_recursion_limit,
);
if let Some(method_name) = method_name {
(
ReferenceType::MethodDefinition,
method_name,
Some(text.to_string()),
)
// Use language-specific handler to find receiver type, or fall back to text
let type_name = Self::find_receiver_type(&node, source, language)
.or_else(|| Some(text.to_string()));
if let Some(type_name) = type_name {
(
ReferenceType::MethodDefinition,
method_name,
Some(type_name),
)
} else {
continue;
}
} else {
continue;
}
@@ -428,6 +436,18 @@ impl ElementExtractor {
.and_then(|handler| handler(receiver_node, source, ast_recursion_limit))
}
fn find_receiver_type(
receiver_node: &tree_sitter::Node,
source: &str,
language: &str,
) -> Option<String> {
use crate::developer::analyze::languages;
languages::get_language_info(language)
.and_then(|info| info.find_receiver_type_handler)
.and_then(|handler| handler(receiver_node, source))
}
fn find_containing_function(
node: &tree_sitter::Node,
source: &str,
crates/goose-mcp/src/developer/analyze/tests/mod.rs
+1
View File
@@ -9,4 +9,5 @@ pub mod integration_tests;
pub mod large_output_tests;
pub mod parser_tests;
pub mod ruby_test;
pub mod rust_test;
pub mod traversal_tests;
crates/goose-mcp/src/developer/analyze/tests/rust_test.rs
+179
View File
@@ -0,0 +1,179 @@
use crate::developer::analyze::graph::CallGraph;
use crate::developer::analyze::parser::{ElementExtractor, ParserManager};
use crate::developer::analyze::types::{AnalysisResult, ReferenceType};
use std::collections::HashSet;
use std::path::PathBuf;
fn parse_and_extract(code: &str) -> AnalysisResult {
let manager = ParserManager::new();
let tree = manager.parse(code, "rust").unwrap();
ElementExtractor::extract_with_depth(&tree, code, "rust", "semantic", None).unwrap()
}
fn build_test_graph(files: Vec<(&str, &str)>) -> CallGraph {
let manager = ParserManager::new();
let results: Vec<_> = files
.iter()
.map(|(path, code)| {
let tree = manager.parse(code, "rust").unwrap();
let result =
ElementExtractor::extract_with_depth(&tree, code, "rust", "semantic", None)
.unwrap();
(PathBuf::from(*path), result)
})
.collect();
CallGraph::build_from_results(&results)
}
#[test]
fn test_rust_self_parameter_type_resolution() {
// Test that self parameters correctly resolve to their impl type
let code = r#"
struct MyStruct {
value: i32,
}
impl MyStruct {
fn method_with_self(&self) -> i32 {
self.value
}
fn method_with_mut_self(&mut self) {
self.value += 1;
}
fn associated_function() -> Self {
MyStruct { value: 0 }
}
}
"#;
let result = parse_and_extract(code);
// Find method references with self parameters
let self_methods: Vec<_> = result
.references
.iter()
.filter(|r| r.ref_type == ReferenceType::MethodDefinition)
.collect();
// Should find both methods with self parameters
assert_eq!(
self_methods.len(),
2,
"Expected 2 methods with self parameters"
);
// Both should be associated with MyStruct
for method_ref in &self_methods {
assert_eq!(
method_ref.associated_type.as_deref(),
Some("MyStruct"),
"Method {} should be associated with MyStruct",
method_ref.symbol
);
}
// Verify the specific methods
let method_names: HashSet<_> = self_methods.iter().map(|r| r.symbol.as_str()).collect();
assert!(method_names.contains("method_with_self"));
assert!(method_names.contains("method_with_mut_self"));
}
#[test]
fn test_rust_struct_and_impl_tracking() {
let code = r#"
struct Config {
host: String,
port: u16,
}
struct Handler {
cfg: Config,
}
impl Handler {
fn new(cfg: Config) -> Self {
Handler { cfg }
}
fn start(&self) -> Result<(), String> {
Ok(())
}
}
fn main() {
let cfg = Config { host: "localhost".to_string(), port: 8080 };
let handler = Handler::new(cfg);
let _ = handler.start();
}
"#;
let result = parse_and_extract(code);
let graph = build_test_graph(vec![("test.rs", code)]);
// Test struct extraction (includes impl blocks)
assert_eq!(result.class_count, 3); // Config, Handler, impl Handler
let struct_names: HashSet<_> = result.classes.iter().map(|c| c.name.as_str()).collect();
assert!(struct_names.contains("Config"));
assert!(struct_names.contains("Handler"));
// Test method extraction
let method_names: HashSet<_> = result.functions.iter().map(|f| f.name.as_str()).collect();
assert!(method_names.contains("new"));
assert!(method_names.contains("start"));
assert!(method_names.contains("main"));
// Test method-to-type associations (only methods with self parameter)
let handler_methods: Vec<_> = result
.references
.iter()
.filter(|r| {
r.ref_type == ReferenceType::MethodDefinition
&& r.associated_type.as_deref() == Some("Handler")
})
.collect();
assert!(
!handler_methods.is_empty(),
"Expected at least 1 method on Handler (start), found {}",
handler_methods.len()
);
// Verify the method is 'start' (new doesn't have self, so it's not tracked)
assert!(
handler_methods.iter().any(|r| r.symbol == "start"),
"Expected to find 'start' method on Handler"
);
// Test field type tracking
let field_type_refs: Vec<_> = result
.references
.iter()
.filter(|r| r.ref_type == ReferenceType::FieldType)
.collect();
assert!(
!field_type_refs.is_empty(),
"Expected to find field type references"
);
// Test struct instantiation
let config_literals: Vec<_> = result
.references
.iter()
.filter(|r| r.symbol == "Config" && r.ref_type == ReferenceType::TypeInstantiation)
.collect();
assert!(
!config_literals.is_empty(),
"Expected to find Config struct literals"
);
// Test call graph integration
let incoming = graph.find_incoming_chains("Handler", 1);
assert!(
!incoming.is_empty(),
"Expected to find incoming references to Handler"
);
let outgoing = graph.find_outgoing_chains("Handler", 1);
assert!(!outgoing.is_empty(), "Expected to find methods on Handler");
}