From a1fb9585d34c9ff60dae057468da58fa14be4554 Mon Sep 17 00:00:00 2001 From: Lum1104 Date: Wed, 15 Apr 2026 18:38:26 +0800 Subject: [PATCH] feat: add RustExtractor for tree-sitter Rust structural analysis Co-Authored-By: Claude Opus 4.6 (1M context) --- .../__tests__/rust-extractor.test.ts | 755 ++++++++++++++++++ .../src/plugins/extractors/rust-extractor.ts | 509 ++++++++++++ 2 files changed, 1264 insertions(+) create mode 100644 understand-anything-plugin/packages/core/src/plugins/extractors/__tests__/rust-extractor.test.ts create mode 100644 understand-anything-plugin/packages/core/src/plugins/extractors/rust-extractor.ts diff --git a/understand-anything-plugin/packages/core/src/plugins/extractors/__tests__/rust-extractor.test.ts b/understand-anything-plugin/packages/core/src/plugins/extractors/__tests__/rust-extractor.test.ts new file mode 100644 index 0000000..0687960 --- /dev/null +++ b/understand-anything-plugin/packages/core/src/plugins/extractors/__tests__/rust-extractor.test.ts @@ -0,0 +1,755 @@ +import { describe, it, expect, beforeAll } from "vitest"; +import { createRequire } from "node:module"; +import { RustExtractor } from "../rust-extractor.js"; + +const require = createRequire(import.meta.url); + +// Load tree-sitter + Rust grammar once +let Parser: any; +let Language: any; +let rustLang: any; + +beforeAll(async () => { + const mod = await import("web-tree-sitter"); + Parser = mod.Parser; + Language = mod.Language; + await Parser.init(); + const wasmPath = require.resolve( + "tree-sitter-rust/tree-sitter-rust.wasm", + ); + rustLang = await Language.load(wasmPath); +}); + +function parse(code: string) { + const parser = new Parser(); + parser.setLanguage(rustLang); + const tree = parser.parse(code); + const root = tree.rootNode; + return { tree, parser, root }; +} + +describe("RustExtractor", () => { + const extractor = new RustExtractor(); + + it("has correct languageIds", () => { + expect(extractor.languageIds).toEqual(["rust"]); + }); + + // ---- Functions ---- + + describe("extractStructure - functions", () => { + it("extracts top-level functions with params and return types", () => { + const { tree, parser, root } = parse(` +pub fn check_port(port: u16) -> bool { + port > 0 +} + +fn helper(name: String, count: i32) -> String { + name.repeat(count) +} +`); + const result = extractor.extractStructure(root); + + expect(result.functions).toHaveLength(2); + + expect(result.functions[0].name).toBe("check_port"); + expect(result.functions[0].params).toEqual(["port"]); + expect(result.functions[0].returnType).toBe("bool"); + + expect(result.functions[1].name).toBe("helper"); + expect(result.functions[1].params).toEqual(["name", "count"]); + expect(result.functions[1].returnType).toBe("String"); + + tree.delete(); + parser.delete(); + }); + + it("extracts functions with no params and no return type", () => { + const { tree, parser, root } = parse(` +fn noop() { +} +`); + const result = extractor.extractStructure(root); + + expect(result.functions).toHaveLength(1); + expect(result.functions[0].name).toBe("noop"); + expect(result.functions[0].params).toEqual([]); + expect(result.functions[0].returnType).toBeUndefined(); + + tree.delete(); + parser.delete(); + }); + + it("reports correct line ranges for multi-line functions", () => { + const { tree, parser, root } = parse(` +fn multiline( + a: i32, + b: i32, +) -> i32 { + let result = a + b; + result +} +`); + const result = extractor.extractStructure(root); + + expect(result.functions).toHaveLength(1); + expect(result.functions[0].lineRange[0]).toBe(2); + expect(result.functions[0].lineRange[1]).toBe(8); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Methods (impl blocks) ---- + + describe("extractStructure - impl blocks and methods", () => { + it("extracts methods from impl blocks and links them to structs", () => { + const { tree, parser, root } = parse(` +pub struct Config { + name: String, + port: u16, +} + +impl Config { + pub fn new(name: String, port: u16) -> Self { + Config { name, port } + } + + fn validate(&self) -> bool { + true + } +} +`); + const result = extractor.extractStructure(root); + + // Methods appear in functions list + const fnNames = result.functions.map((f) => f.name); + expect(fnNames).toContain("new"); + expect(fnNames).toContain("validate"); + + // new() skips self parameter, captures real params + const newFn = result.functions.find((f) => f.name === "new"); + expect(newFn?.params).toEqual(["name", "port"]); + expect(newFn?.returnType).toBe("Self"); + + // validate() has &self — should have no params + const validateFn = result.functions.find((f) => f.name === "validate"); + expect(validateFn?.params).toEqual([]); + expect(validateFn?.returnType).toBe("bool"); + + // Methods linked to struct + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Config"); + expect(result.classes[0].methods).toContain("new"); + expect(result.classes[0].methods).toContain("validate"); + + tree.delete(); + parser.delete(); + }); + + it("handles impl blocks for enums", () => { + const { tree, parser, root } = parse(` +enum Status { + Active, + Inactive, +} + +impl Status { + fn is_active(&self) -> bool { + true + } +} +`); + const result = extractor.extractStructure(root); + + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Status"); + expect(result.classes[0].methods).toContain("is_active"); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Structs ---- + + describe("extractStructure - structs", () => { + it("extracts struct with fields", () => { + const { tree, parser, root } = parse(` +pub struct Config { + name: String, + port: u16, +} +`); + const result = extractor.extractStructure(root); + + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Config"); + expect(result.classes[0].properties).toEqual(["name", "port"]); + expect(result.classes[0].methods).toEqual([]); + expect(result.classes[0].lineRange[0]).toBe(2); + + tree.delete(); + parser.delete(); + }); + + it("extracts empty struct (unit struct)", () => { + const { tree, parser, root } = parse(` +struct Empty; +`); + const result = extractor.extractStructure(root); + + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Empty"); + expect(result.classes[0].properties).toEqual([]); + expect(result.classes[0].methods).toEqual([]); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Enums ---- + + describe("extractStructure - enums", () => { + it("extracts enum with variants as properties", () => { + const { tree, parser, root } = parse(` +enum Status { + Active, + Inactive, +} +`); + const result = extractor.extractStructure(root); + + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Status"); + expect(result.classes[0].properties).toEqual(["Active", "Inactive"]); + expect(result.classes[0].methods).toEqual([]); + + tree.delete(); + parser.delete(); + }); + + it("extracts pub enum", () => { + const { tree, parser, root } = parse(` +pub enum Direction { + North, + South, + East, + West, +} +`); + const result = extractor.extractStructure(root); + + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Direction"); + expect(result.classes[0].properties).toEqual(["North", "South", "East", "West"]); + + // pub enum should be exported + const exportNames = result.exports.map((e) => e.name); + expect(exportNames).toContain("Direction"); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Traits ---- + + describe("extractStructure - traits", () => { + it("extracts trait with method signatures", () => { + const { tree, parser, root } = parse(` +pub trait Validator { + fn validate(&self) -> bool; + fn name(&self) -> &str; +} +`); + const result = extractor.extractStructure(root); + + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Validator"); + expect(result.classes[0].methods).toEqual(["validate", "name"]); + expect(result.classes[0].properties).toEqual([]); + + tree.delete(); + parser.delete(); + }); + + it("extracts empty trait", () => { + const { tree, parser, root } = parse(` +trait Marker {} +`); + const result = extractor.extractStructure(root); + + expect(result.classes).toHaveLength(1); + expect(result.classes[0].name).toBe("Marker"); + expect(result.classes[0].methods).toEqual([]); + + tree.delete(); + parser.delete(); + }); + + it("exports pub traits", () => { + const { tree, parser, root } = parse(` +pub trait Serializable { + fn serialize(&self) -> String; +} + +trait Internal { + fn process(&self); +} +`); + const result = extractor.extractStructure(root); + + const exportNames = result.exports.map((e) => e.name); + expect(exportNames).toContain("Serializable"); + expect(exportNames).not.toContain("Internal"); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Imports ---- + + describe("extractStructure - imports", () => { + it("extracts scoped identifier imports", () => { + const { tree, parser, root } = parse(` +use std::collections::HashMap; +`); + const result = extractor.extractStructure(root); + + expect(result.imports).toHaveLength(1); + expect(result.imports[0].source).toBe("std::collections"); + expect(result.imports[0].specifiers).toEqual(["HashMap"]); + + tree.delete(); + parser.delete(); + }); + + it("extracts scoped use list imports", () => { + const { tree, parser, root } = parse(` +use std::io::{self, Read, Write}; +`); + const result = extractor.extractStructure(root); + + expect(result.imports).toHaveLength(1); + expect(result.imports[0].source).toBe("std::io"); + expect(result.imports[0].specifiers).toEqual(["self", "Read", "Write"]); + + tree.delete(); + parser.delete(); + }); + + it("extracts wildcard imports", () => { + const { tree, parser, root } = parse(` +use std::prelude::*; +`); + const result = extractor.extractStructure(root); + + expect(result.imports).toHaveLength(1); + expect(result.imports[0].source).toBe("std::prelude"); + expect(result.imports[0].specifiers).toEqual(["*"]); + + tree.delete(); + parser.delete(); + }); + + it("extracts simple identifier imports", () => { + const { tree, parser, root } = parse(` +use foo; +`); + const result = extractor.extractStructure(root); + + expect(result.imports).toHaveLength(1); + expect(result.imports[0].source).toBe("foo"); + expect(result.imports[0].specifiers).toEqual(["foo"]); + + tree.delete(); + parser.delete(); + }); + + it("extracts crate-relative imports", () => { + const { tree, parser, root } = parse(` +use crate::config::Settings; +`); + const result = extractor.extractStructure(root); + + expect(result.imports).toHaveLength(1); + expect(result.imports[0].source).toBe("crate::config"); + expect(result.imports[0].specifiers).toEqual(["Settings"]); + + tree.delete(); + parser.delete(); + }); + + it("reports correct import line numbers", () => { + const { tree, parser, root } = parse(` +use std::collections::HashMap; +use std::io::{self, Read}; +`); + const result = extractor.extractStructure(root); + + expect(result.imports).toHaveLength(2); + expect(result.imports[0].lineNumber).toBe(2); + expect(result.imports[1].lineNumber).toBe(3); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Exports ---- + + describe("extractStructure - exports", () => { + it("exports pub items and not private items", () => { + const { tree, parser, root } = parse(` +pub struct Config { + name: String, +} + +struct Internal { + value: i32, +} + +pub fn check_port(port: u16) -> bool { + port > 0 +} + +fn helper() {} +`); + const result = extractor.extractStructure(root); + + const exportNames = result.exports.map((e) => e.name); + expect(exportNames).toContain("Config"); + expect(exportNames).toContain("check_port"); + expect(exportNames).not.toContain("Internal"); + expect(exportNames).not.toContain("helper"); + + tree.delete(); + parser.delete(); + }); + + it("exports pub methods inside impl blocks", () => { + const { tree, parser, root } = parse(` +pub struct Config { + name: String, +} + +impl Config { + pub fn new(name: String) -> Self { + Config { name } + } + + fn validate(&self) -> bool { + true + } +} +`); + const result = extractor.extractStructure(root); + + const exportNames = result.exports.map((e) => e.name); + expect(exportNames).toContain("Config"); + expect(exportNames).toContain("new"); + expect(exportNames).not.toContain("validate"); + + tree.delete(); + parser.delete(); + }); + + it("reports correct export line numbers", () => { + const { tree, parser, root } = parse(` +pub struct Config { + name: String, +} + +pub fn check_port(port: u16) -> bool { + port > 0 +} +`); + const result = extractor.extractStructure(root); + + const configExport = result.exports.find((e) => e.name === "Config"); + expect(configExport?.lineNumber).toBe(2); + + const checkPortExport = result.exports.find((e) => e.name === "check_port"); + expect(checkPortExport?.lineNumber).toBe(6); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Call Graph ---- + + describe("extractCallGraph", () => { + it("extracts simple function calls", () => { + const { tree, parser, root } = parse(` +fn validate(port: u16) -> bool { + check_port(port) +} + +fn main() { + validate(8080); +} +`); + const result = extractor.extractCallGraph(root); + + const validateCalls = result.filter((e) => e.caller === "validate"); + expect(validateCalls.some((e) => e.callee === "check_port")).toBe(true); + + const mainCalls = result.filter((e) => e.caller === "main"); + expect(mainCalls.some((e) => e.callee === "validate")).toBe(true); + + tree.delete(); + parser.delete(); + }); + + it("extracts method calls (field_expression)", () => { + const { tree, parser, root } = parse(` +fn process(&self) { + self.validate(); +} +`); + const result = extractor.extractCallGraph(root); + + expect(result).toHaveLength(1); + expect(result[0].caller).toBe("process"); + expect(result[0].callee).toBe("self.validate"); + + tree.delete(); + parser.delete(); + }); + + it("extracts scoped calls (e.g., Vec::new)", () => { + const { tree, parser, root } = parse(` +fn create() { + Vec::new(); +} +`); + const result = extractor.extractCallGraph(root); + + expect(result).toHaveLength(1); + expect(result[0].caller).toBe("create"); + expect(result[0].callee).toBe("Vec::new"); + + tree.delete(); + parser.delete(); + }); + + it("tracks correct caller for methods in impl blocks", () => { + const { tree, parser, root } = parse(` +impl Config { + fn validate(&self) -> bool { + check_port(self.port) + } +} +`); + const result = extractor.extractCallGraph(root); + + expect(result).toHaveLength(1); + expect(result[0].caller).toBe("validate"); + expect(result[0].callee).toBe("check_port"); + + tree.delete(); + parser.delete(); + }); + + it("reports correct line numbers for calls", () => { + const { tree, parser, root } = parse(` +fn main() { + foo(); + bar(); +} +`); + const result = extractor.extractCallGraph(root); + + expect(result).toHaveLength(2); + expect(result[0].lineNumber).toBe(3); + expect(result[1].lineNumber).toBe(4); + + tree.delete(); + parser.delete(); + }); + + it("ignores top-level calls (no caller)", () => { + const { tree, parser, root } = parse(` +static X: i32 = compute(); +`); + const result = extractor.extractCallGraph(root); + + // Top-level calls have no enclosing function, so they are skipped + expect(result).toHaveLength(0); + + tree.delete(); + parser.delete(); + }); + }); + + // ---- Comprehensive ---- + + describe("comprehensive Rust file", () => { + it("handles the full test scenario from the spec", () => { + const { tree, parser, root } = parse(`use std::collections::HashMap; +use std::io::{self, Read}; + +pub struct Config { + name: String, + port: u16, +} + +impl Config { + pub fn new(name: String, port: u16) -> Self { + Config { name, port } + } + + fn validate(&self) -> bool { + check_port(self.port) + } +} + +pub fn check_port(port: u16) -> bool { + port > 0 +} + +enum Status { + Active, + Inactive, +} +`); + const result = extractor.extractStructure(root); + + // Functions: new, validate, check_port + expect(result.functions).toHaveLength(3); + const fnNames = result.functions.map((f) => f.name).sort(); + expect(fnNames).toEqual(["check_port", "new", "validate"]); + + // new() params + const newFn = result.functions.find((f) => f.name === "new"); + expect(newFn?.params).toEqual(["name", "port"]); + expect(newFn?.returnType).toBe("Self"); + + // validate() has &self — no visible params + const validateFn = result.functions.find((f) => f.name === "validate"); + expect(validateFn?.params).toEqual([]); + expect(validateFn?.returnType).toBe("bool"); + + // check_port() + const checkPortFn = result.functions.find((f) => f.name === "check_port"); + expect(checkPortFn?.params).toEqual(["port"]); + expect(checkPortFn?.returnType).toBe("bool"); + + // Classes: Config (struct) and Status (enum) + expect(result.classes).toHaveLength(2); + + const configClass = result.classes.find((c) => c.name === "Config"); + expect(configClass).toBeDefined(); + expect(configClass!.properties).toEqual(["name", "port"]); + expect(configClass!.methods).toContain("new"); + expect(configClass!.methods).toContain("validate"); + + const statusClass = result.classes.find((c) => c.name === "Status"); + expect(statusClass).toBeDefined(); + expect(statusClass!.properties).toEqual(["Active", "Inactive"]); + expect(statusClass!.methods).toEqual([]); + + // Imports: 2 + expect(result.imports).toHaveLength(2); + expect(result.imports[0].source).toBe("std::collections"); + expect(result.imports[0].specifiers).toEqual(["HashMap"]); + expect(result.imports[1].source).toBe("std::io"); + expect(result.imports[1].specifiers).toEqual(["self", "Read"]); + + // Exports: Config, new, check_port (those with pub) + const exportNames = result.exports.map((e) => e.name).sort(); + expect(exportNames).toEqual(["Config", "check_port", "new"]); + + // Call graph: validate -> check_port + const calls = extractor.extractCallGraph(root); + const validateCalls = calls.filter((e) => e.caller === "validate"); + expect(validateCalls).toHaveLength(1); + expect(validateCalls[0].callee).toBe("check_port"); + + tree.delete(); + parser.delete(); + }); + + it("handles a realistic Rust module with traits and multiple impls", () => { + const { tree, parser, root } = parse(`use std::fmt; +use std::io::{self, Write}; + +pub trait Displayable { + fn display(&self) -> String; +} + +pub struct Server { + host: String, + port: u16, +} + +impl Server { + pub fn new(host: String, port: u16) -> Self { + Server { host, port } + } + + pub fn start(&self) { + listen(self.port); + } +} + +impl Displayable for Server { + fn display(&self) -> String { + format_server(self.host.clone(), self.port) + } +} + +fn listen(port: u16) { + println!("Listening on port {}", port); +} + +fn format_server(host: String, port: u16) -> String { + format!("{}:{}", host, port) +} +`); + const result = extractor.extractStructure(root); + + // Functions: new, start, display, listen, format_server + expect(result.functions).toHaveLength(5); + + // Trait: Displayable + const trait_ = result.classes.find((c) => c.name === "Displayable"); + expect(trait_).toBeDefined(); + expect(trait_!.methods).toEqual(["display"]); + + // Struct: Server + const server = result.classes.find((c) => c.name === "Server"); + expect(server).toBeDefined(); + expect(server!.properties).toEqual(["host", "port"]); + // Methods from both impl blocks (Server and Displayable for Server) + expect(server!.methods).toContain("new"); + expect(server!.methods).toContain("start"); + expect(server!.methods).toContain("display"); + + // Exports: pub items + const exportNames = result.exports.map((e) => e.name).sort(); + expect(exportNames).toContain("Displayable"); + expect(exportNames).toContain("Server"); + expect(exportNames).toContain("new"); + expect(exportNames).toContain("start"); + expect(exportNames).not.toContain("listen"); + expect(exportNames).not.toContain("format_server"); + + // Call graph + const calls = extractor.extractCallGraph(root); + const startCalls = calls.filter((e) => e.caller === "start"); + expect(startCalls.some((e) => e.callee === "listen")).toBe(true); + + const displayCalls = calls.filter((e) => e.caller === "display"); + expect(displayCalls.some((e) => e.callee === "format_server")).toBe(true); + + tree.delete(); + parser.delete(); + }); + }); +}); diff --git a/understand-anything-plugin/packages/core/src/plugins/extractors/rust-extractor.ts b/understand-anything-plugin/packages/core/src/plugins/extractors/rust-extractor.ts new file mode 100644 index 0000000..98ab38f --- /dev/null +++ b/understand-anything-plugin/packages/core/src/plugins/extractors/rust-extractor.ts @@ -0,0 +1,509 @@ +import type { StructuralAnalysis, CallGraphEntry } from "../../types.js"; +import type { LanguageExtractor, TreeSitterNode } from "./types.js"; +import { findChild, findChildren } from "./base-extractor.js"; + +/** + * Extract parameter names from a Rust `parameters` node. + * + * Each `parameter` child has a `pattern` field (identifier) and a `type` field. + * `self_parameter` nodes (&self, &mut self, self) are skipped since they are + * implicit receivers, not user-facing parameters. + */ +function extractParams(paramsNode: TreeSitterNode | null): string[] { + if (!paramsNode) return []; + const params: string[] = []; + + for (let i = 0; i < paramsNode.childCount; i++) { + const child = paramsNode.child(i); + if (!child) continue; + + if (child.type === "parameter") { + const pattern = child.childForFieldName("pattern"); + if (pattern) { + params.push(pattern.text); + } + } + // Skip self_parameter — it's the receiver, not a real parameter + } + + return params; +} + +/** + * Extract the return type from a function_item node. + * + * In tree-sitter-rust, the return type is accessed via the `return_type` named + * field on function_item. The field value is the type node itself (e.g. + * primitive_type "bool", type_identifier "Self"). + */ +function extractReturnType(node: TreeSitterNode): string | undefined { + const returnType = node.childForFieldName("return_type"); + if (returnType) { + return returnType.text; + } + return undefined; +} + +/** + * Check if a node has a `visibility_modifier` child whose text starts with "pub". + * Covers `pub`, `pub(crate)`, `pub(super)`, etc. + */ +function isPublic(node: TreeSitterNode): boolean { + const visMod = findChild(node, "visibility_modifier"); + return visMod !== null && visMod.text.startsWith("pub"); +} + +/** + * Recursively extract the path portion of a scoped_identifier. + * + * `scoped_identifier` nests: `std::collections::HashMap` is + * scoped_identifier(path: scoped_identifier(path: identifier "std", name: identifier "collections"), name: identifier "HashMap") + * + * This function collects all path segments into a flat "a::b::c" string, + * excluding the final `name` (which is the imported specifier). + */ +function extractScopedPath(node: TreeSitterNode): { path: string; name: string } { + if (node.type === "scoped_identifier") { + const pathNode = node.childForFieldName("path"); + const nameNode = node.childForFieldName("name"); + const name = nameNode ? nameNode.text : ""; + const path = pathNode ? pathNode.text : ""; + return { path, name }; + } + // Bare identifier: `use foo;` + return { path: "", name: node.text }; +} + +/** + * Rust extractor for tree-sitter structural analysis and call graph extraction. + * + * Handles functions, structs, enums, traits, impl blocks, use declarations, + * visibility-based exports, and call graphs for Rust source code. + * + * Rust-specific mapping decisions: + * - Structs, enums, and traits are mapped to the `classes` array. + * - Methods inside `impl` blocks are stored as functions and also listed + * in the corresponding struct/enum's `methods` array. + * - Trait method signatures (function_signature_item) are listed in the + * trait's `methods` array. + * - Exports are determined by the `pub` visibility modifier. + * - Enum variants are extracted as `properties` of the enum class entry. + */ +export class RustExtractor implements LanguageExtractor { + readonly languageIds = ["rust"]; + + extractStructure(rootNode: TreeSitterNode): StructuralAnalysis { + const functions: StructuralAnalysis["functions"] = []; + const classes: StructuralAnalysis["classes"] = []; + const imports: StructuralAnalysis["imports"] = []; + const exports: StructuralAnalysis["exports"] = []; + + // Track methods per impl type so we can attach them to structs/enums + const methodsByType = new Map(); + + for (let i = 0; i < rootNode.childCount; i++) { + const node = rootNode.child(i); + if (!node) continue; + + switch (node.type) { + case "function_item": + this.extractFunction(node, functions, exports); + break; + + case "struct_item": + this.extractStruct(node, classes, exports); + break; + + case "enum_item": + this.extractEnum(node, classes, exports); + break; + + case "trait_item": + this.extractTrait(node, classes, exports); + break; + + case "impl_item": + this.extractImpl(node, functions, exports, methodsByType); + break; + + case "use_declaration": + this.extractUseDeclaration(node, imports); + break; + } + } + + // Attach collected methods to their corresponding structs/enums/traits + for (const cls of classes) { + const methods = methodsByType.get(cls.name); + if (methods) { + cls.methods.push(...methods); + } + } + + return { functions, classes, imports, exports }; + } + + extractCallGraph(rootNode: TreeSitterNode): CallGraphEntry[] { + const entries: CallGraphEntry[] = []; + const functionStack: string[] = []; + + const walkForCalls = (node: TreeSitterNode) => { + let pushedName = false; + + // Track entering function_item declarations + if (node.type === "function_item") { + const nameNode = node.childForFieldName("name"); + if (nameNode) { + functionStack.push(nameNode.text); + pushedName = true; + } + } + + // Extract call expressions + if (node.type === "call_expression") { + if (functionStack.length > 0) { + const callee = this.extractCalleeName(node); + if (callee) { + entries.push({ + caller: functionStack[functionStack.length - 1], + callee, + lineNumber: node.startPosition.row + 1, + }); + } + } + } + + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i); + if (child) walkForCalls(child); + } + + if (pushedName) { + functionStack.pop(); + } + }; + + walkForCalls(rootNode); + + return entries; + } + + // ---- Private helpers ---- + + /** + * Extract the callee name from a call_expression. + * + * Handles: + * - Plain function call: `check_port(x)` -> "check_port" + * - Method call via field_expression: `self.validate()` -> "self.validate" + * - Scoped call: `Vec::new()` -> "Vec::new" + */ + private extractCalleeName(callNode: TreeSitterNode): string | null { + const funcNode = callNode.child(0); + if (!funcNode) return null; + + if (funcNode.type === "identifier") { + return funcNode.text; + } + + if (funcNode.type === "field_expression") { + // e.g., self.validate or obj.method + const field = funcNode.childForFieldName("field"); + const value = funcNode.childForFieldName("value"); + if (field && value) { + return value.text + "." + field.text; + } + } + + if (funcNode.type === "scoped_identifier") { + // e.g., Vec::new + return funcNode.text; + } + + // Fallback: use the full text of the function child + return funcNode.text; + } + + private extractFunction( + node: TreeSitterNode, + functions: StructuralAnalysis["functions"], + exports: StructuralAnalysis["exports"], + ): void { + const nameNode = node.childForFieldName("name"); + if (!nameNode) return; + + const paramsNode = node.childForFieldName("parameters"); + const params = extractParams(paramsNode ?? null); + const returnType = extractReturnType(node); + + functions.push({ + name: nameNode.text, + lineRange: [ + node.startPosition.row + 1, + node.endPosition.row + 1, + ], + params, + returnType, + }); + + if (isPublic(node)) { + exports.push({ + name: nameNode.text, + lineNumber: node.startPosition.row + 1, + }); + } + } + + private extractStruct( + node: TreeSitterNode, + classes: StructuralAnalysis["classes"], + exports: StructuralAnalysis["exports"], + ): void { + const nameNode = node.childForFieldName("name"); + if (!nameNode) return; + + const properties: string[] = []; + const body = node.childForFieldName("body"); + if (body && body.type === "field_declaration_list") { + const fields = findChildren(body, "field_declaration"); + for (const field of fields) { + const fieldName = findChild(field, "field_identifier"); + if (fieldName) { + properties.push(fieldName.text); + } + } + } + + classes.push({ + name: nameNode.text, + lineRange: [ + node.startPosition.row + 1, + node.endPosition.row + 1, + ], + methods: [], // Methods are attached later from methodsByType + properties, + }); + + if (isPublic(node)) { + exports.push({ + name: nameNode.text, + lineNumber: node.startPosition.row + 1, + }); + } + } + + private extractEnum( + node: TreeSitterNode, + classes: StructuralAnalysis["classes"], + exports: StructuralAnalysis["exports"], + ): void { + const nameNode = node.childForFieldName("name"); + if (!nameNode) return; + + const properties: string[] = []; + const body = node.childForFieldName("body"); + if (body && body.type === "enum_variant_list") { + const variants = findChildren(body, "enum_variant"); + for (const variant of variants) { + const variantName = variant.childForFieldName("name"); + if (variantName) { + properties.push(variantName.text); + } + } + } + + classes.push({ + name: nameNode.text, + lineRange: [ + node.startPosition.row + 1, + node.endPosition.row + 1, + ], + methods: [], // Methods are attached later if there's an impl block + properties, + }); + + if (isPublic(node)) { + exports.push({ + name: nameNode.text, + lineNumber: node.startPosition.row + 1, + }); + } + } + + private extractTrait( + node: TreeSitterNode, + classes: StructuralAnalysis["classes"], + exports: StructuralAnalysis["exports"], + ): void { + const nameNode = node.childForFieldName("name"); + if (!nameNode) return; + + const methods: string[] = []; + const body = findChild(node, "declaration_list"); + if (body) { + // Trait bodies contain function_signature_item for method declarations + const sigs = findChildren(body, "function_signature_item"); + for (const sig of sigs) { + const sigName = findChild(sig, "identifier"); + if (sigName) { + methods.push(sigName.text); + } + } + // Also handle default method implementations (function_item) + const fns = findChildren(body, "function_item"); + for (const fn of fns) { + const fnName = fn.childForFieldName("name"); + if (fnName) { + methods.push(fnName.text); + } + } + } + + classes.push({ + name: nameNode.text, + lineRange: [ + node.startPosition.row + 1, + node.endPosition.row + 1, + ], + methods, + properties: [], + }); + + if (isPublic(node)) { + exports.push({ + name: nameNode.text, + lineNumber: node.startPosition.row + 1, + }); + } + } + + private extractImpl( + node: TreeSitterNode, + functions: StructuralAnalysis["functions"], + exports: StructuralAnalysis["exports"], + methodsByType: Map, + ): void { + const typeNode = node.childForFieldName("type"); + const typeName = typeNode ? typeNode.text : null; + + const body = node.childForFieldName("body"); + if (!body) return; + + const fns = findChildren(body, "function_item"); + for (const fn of fns) { + const nameNode = fn.childForFieldName("name"); + if (!nameNode) continue; + + const paramsNode = fn.childForFieldName("parameters"); + const params = extractParams(paramsNode ?? null); + const returnType = extractReturnType(fn); + + functions.push({ + name: nameNode.text, + lineRange: [ + fn.startPosition.row + 1, + fn.endPosition.row + 1, + ], + params, + returnType, + }); + + // Track method association with the impl type + if (typeName) { + if (!methodsByType.has(typeName)) { + methodsByType.set(typeName, []); + } + methodsByType.get(typeName)!.push(nameNode.text); + } + + // pub methods inside impl blocks are exports + if (isPublic(fn)) { + exports.push({ + name: nameNode.text, + lineNumber: fn.startPosition.row + 1, + }); + } + } + } + + private extractUseDeclaration( + node: TreeSitterNode, + imports: StructuralAnalysis["imports"], + ): void { + const argument = node.childForFieldName("argument"); + if (!argument) return; + + switch (argument.type) { + case "identifier": + // `use foo;` + imports.push({ + source: argument.text, + specifiers: [argument.text], + lineNumber: node.startPosition.row + 1, + }); + break; + + case "scoped_identifier": { + // `use std::collections::HashMap;` + const { path, name } = extractScopedPath(argument); + imports.push({ + source: path, + specifiers: [name], + lineNumber: node.startPosition.row + 1, + }); + break; + } + + case "scoped_use_list": { + // `use std::io::{self, Read, Write};` + const pathNode = argument.childForFieldName("path"); + const listNode = argument.childForFieldName("list"); + const source = pathNode ? pathNode.text : ""; + const specifiers: string[] = []; + + if (listNode) { + for (let j = 0; j < listNode.childCount; j++) { + const ch = listNode.child(j); + if (!ch) continue; + if (ch.type === "self" || ch.type === "identifier") { + specifiers.push(ch.text); + } else if (ch.type === "scoped_identifier") { + // Nested scoped identifier inside a use list + specifiers.push(ch.text); + } + } + } + + imports.push({ + source, + specifiers, + lineNumber: node.startPosition.row + 1, + }); + break; + } + + case "use_wildcard": { + // `use std::prelude::*;` + // The path is the scoped_identifier child + const scopedId = findChild(argument, "scoped_identifier"); + const source = scopedId ? scopedId.text : ""; + imports.push({ + source, + specifiers: ["*"], + lineNumber: node.startPosition.row + 1, + }); + break; + } + + default: { + // Fallback for any unhandled pattern + imports.push({ + source: argument.text, + specifiers: [argument.text], + lineNumber: node.startPosition.row + 1, + }); + break; + } + } + } +}