Improve Rust analysis output for analyze tool (#5072)

crates/goose-mcp/src/developer/analyze/formatter.rs

@@ -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

@@ -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

@@ -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

@@ -380,11 +380,19 @@ impl ElementExtractor {
                             ast_recursion_limit,
                         );
                         if let Some(method_name) = method_name {
-                            (
+                            // Use language-specific handler to find receiver type, or fall back to text
-                                ReferenceType::MethodDefinition,
+                            let type_name = Self::find_receiver_type(&node, source, language)
-                                method_name,
+                                .or_else(|| Some(text.to_string()));
-                                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

@@ -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

@@ -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");
+}