diff --git a/mcp_codebase.py b/mcp_codebase.py index 305d21c..bc99f59 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2505,11 +2505,7 @@ def parse_rust_file(filepath: Path) -> List[Dict]: @lru_cache(maxsize=1000) def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: - """Cached Rust parsing with enhanced metadata for graph relationships. - ALWAYS returns a tuple of (text, metadata_tuple) entries. If the Rust helper - fails or produces no chunks, return a single fallback chunk containing the - whole file text and minimal metadata. - """ + """Enhanced Rust parsing with better metadata extraction.""" filepath = Path(filepath_str) chunks = [] @@ -2518,15 +2514,7 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: rust_helper = Path("./tools/target/release/parse_rust_ast") if not rust_helper.is_file(): logger.warning(f"Rust parser binary not found: {rust_helper}. Using fallback whole-file chunk for {filepath}") - code = filepath.read_text(encoding="utf-8") - fallback_meta = { - "file": str(filepath), - "name": filepath.name, - "type": "file", - "line": 1, - "language": "rust", - } - return ((code, tuple(fallback_meta.items())),) + return create_fallback_chunk(filepath) try: code = filepath.read_text(encoding="utf-8") @@ -2543,61 +2531,41 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: logger.info(f" AST helper found {len(decls)} declarations") for i, d in enumerate(decls): - start_line = max(0, d.get("start_line", 1) - 1) - end_line = d.get("end_line", start_line + 1) - # guard the end_line to be at least start_line+1 - if end_line <= start_line: - end_line = start_line + 1 + # FIX: Ensure valid line numbers + start_line = max(1, d.get("start_line", 1)) + end_line = max(start_line, d.get("end_line", start_line + 1)) - # Make slice safe even if end_line > len(lines) - chunk_code = "".join(lines[start_line:end_line]) + # Ensure we don't exceed file bounds + if start_line > len(lines): + start_line = len(lines) + if end_line > len(lines): + end_line = len(lines) + + chunk_code = "".join(lines[start_line-1:end_line]) metadata = { "file": str(filepath), - "name": d.get("name", "") or filepath.name, + "name": d.get("name", "") or filepath.stem, "type": d.get("type_", "") or "unknown", - "line": start_line + 1, + "line": start_line, "language": "rust", } - # copy other fields defensively - if d.get("visibility"): - metadata["visibility"] = d.get("visibility") - if d.get("is_async"): - metadata["is_async"] = True - if d.get("is_unsafe"): - metadata["is_unsafe"] = True - if d.get("generics"): - metadata["generics"] = d.get("generics") - if d.get("traits"): - metadata["implements_traits"] = d.get("traits") - if d.get("fields"): - metadata["fields"] = d.get("fields") - if d.get("methods"): - metadata["methods"] = d.get("methods") - if d.get("return_type"): - metadata["return_type"] = d.get("return_type") - if d.get("parameters"): - metadata["parameters"] = d.get("parameters") + # Enhanced metadata extraction + enhanced_fields = [ + "visibility", "is_async", "is_unsafe", "generics", "traits", + "fields", "methods", "return_type", "parameters", "docstring", + "imports", "calls" # NEW fields + ] + + for field in enhanced_fields: + if field in d and d[field] is not None: + metadata[field] = d[field] logger.info(f" Declaration {i+1}: {metadata.get('type', 'unknown')} {metadata.get('name', 'unnamed')}") - chunk_text = f"File: {filepath}\nType: {metadata.get('type')}\nName: {metadata.get('name')}\n" - if metadata.get("visibility"): - chunk_text += f"Visibility: {metadata.get('visibility')}\n" - if metadata.get("is_async"): - chunk_text += "Async: yes\n" - if metadata.get("is_unsafe"): - chunk_text += "Unsafe: yes\n" - if metadata.get("generics"): - chunk_text += f"Generics: {', '.join(metadata.get('generics', []))}\n" - if metadata.get("implements_traits"): - chunk_text += f"Implements: {', '.join(metadata.get('implements_traits', []))}\n" - if metadata.get("fields"): - chunk_text += f"Fields: {', '.join(metadata.get('fields', []))}\n" - - chunk_text += f"Code:\n{chunk_code}" - + # Build rich chunk text + chunk_text = build_rust_chunk_text(metadata, chunk_code) chunks.append({ "text": chunk_text, "metadata": metadata @@ -2605,27 +2573,61 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: except Exception as e: logger.warning(f"Rust AST helper failed for {filepath}: {e}") + return create_fallback_chunk(filepath) - # If parser produced chunks, return them in the (text, metadata_tuple) format expected by parse_rust_file() if chunks: logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}") return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) - # Fallback: emit a whole-file chunk to keep counts consistent - logger.warning(f"❌ No chunks created from Rust AST for {filepath}. Emitting fallback whole-file chunk.") + logger.warning(f"❌ No chunks created from Rust AST for {filepath}. Emitting fallback.") + return create_fallback_chunk(filepath) + +def build_rust_chunk_text(metadata: dict, code: str) -> str: + """Build comprehensive Rust chunk text with all metadata.""" + parts = [ + f"File: {metadata.get('file', '')}", + f"Type: {metadata.get('type', 'unknown')}", + f"Name: {metadata.get('name', '')}", + ] + + # Add enhanced metadata + if metadata.get("visibility"): + parts.append(f"Visibility: {metadata['visibility']}") + if metadata.get("is_async"): + parts.append("Async: yes") + if metadata.get("is_unsafe"): + parts.append("Unsafe: yes") + if metadata.get("generics"): + parts.append(f"Generics: {', '.join(metadata['generics'])}") + if metadata.get("traits"): + parts.append(f"Implements: {', '.join(metadata['traits'])}") + if metadata.get("fields"): + parts.append(f"Fields: {', '.join(metadata['fields'])}") + if metadata.get("imports"): + parts.append(f"Imports: {', '.join(metadata['imports'])}") + if metadata.get("calls"): + parts.append(f"Calls: {', '.join(metadata['calls'])}") + + parts.append(f"Code:\n{code}") + return "\n".join(parts) + +def create_fallback_chunk(filepath: Path) -> tuple: + """Create a minimal fallback chunk with proper metadata.""" try: full_code = filepath.read_text(encoding="utf-8") except Exception: full_code = "" + fallback_meta = { "file": str(filepath), - "name": filepath.name, + "name": filepath.stem, # Use stem instead of full name "type": "file", "line": 1, "language": "rust", } - return ((full_code, tuple(fallback_meta.items())),) + chunk_text = f"File: {filepath}\nType: file\nName: {filepath.stem}\nCode:\n{full_code}" + return ((chunk_text, tuple(fallback_meta.items())),) @lru_cache(maxsize=1000) def parse_shell_file_cached(filepath_str: str) -> tuple: @@ -5445,6 +5447,7 @@ def signal_handler(sig, frame): parse_sql_file_cached.cache_clear() parse_go_file_cached.cache_clear() parse_java_file_cached.cache_clear() + parse_rust_file_cached.cache_clear() parse_svelte_file_cached.cache_clear() parse_shell_file_cached.cache_clear() logger.info("✓ Caches cleared") diff --git a/tools/src/main.rs b/tools/src/main.rs index 3e4a9d4..bac5a94 100644 --- a/tools/src/main.rs +++ b/tools/src/main.rs @@ -7,7 +7,7 @@ use std::fs; use std::process; use syn::{ spanned::Spanned, Attribute, FnArg, Item, ItemEnum, ItemFn, ItemImpl, ItemMod, ItemStruct, - ItemTrait, Pat, Type, Visibility, + ItemTrait, Pat, Type, UseTree, UseTreeKind, Visibility, }; #[derive(Serialize)] @@ -26,6 +26,8 @@ struct RustDecl { return_type: Option, parameters: Vec, docstring: Option, + imports: Vec, // Track imports for cross-file relationships + calls: Vec, // Track function calls } fn main() { @@ -53,37 +55,73 @@ fn main() { }; let mut decls = Vec::new(); + let mut imports = Vec::new(); // Collect imports separately - // Process declarations + // Process imports first + for item in &syntax.items { + if let Item::Use(item_use) = item { + let import_paths = extract_use_paths(&item_use.tree); + imports.extend(import_paths.clone()); + + for path in &import_paths { + decls.push(RustDecl { + name: path.clone(), + type_: "import".to_string(), + start_line: get_span_start(&item_use.span()), + end_line: get_span_end(&item_use.span()), + visibility: visibility_to_string(&item_use.vis), + is_async: false, + is_unsafe: false, + generics: vec![], + traits: vec![], + fields: vec![], + methods: vec![], + return_type: None, + parameters: vec![path.clone()], + docstring: extract_docstring(&item_use.attrs), + imports: vec![], + calls: vec![], + }); + } + } + } + + // Process other declarations with import context for item in &syntax.items { match item { Item::Fn(item_fn) => { - if let Some(decl) = parse_function(&item_fn) { + if let Some(mut decl) = parse_function(&item_fn, &content) { + decl.imports = imports.clone(); // Add imports to function context decls.push(decl); } } Item::Struct(item_struct) => { - if let Some(decl) = parse_struct(&item_struct) { + if let Some(mut decl) = parse_struct(&item_struct) { + decl.imports = imports.clone(); decls.push(decl); } } Item::Enum(item_enum) => { - if let Some(decl) = parse_enum(&item_enum) { + if let Some(mut decl) = parse_enum(&item_enum) { + decl.imports = imports.clone(); decls.push(decl); } } Item::Trait(item_trait) => { - if let Some(decl) = parse_trait(&item_trait) { + if let Some(mut decl) = parse_trait(&item_trait) { + decl.imports = imports.clone(); decls.push(decl); } } Item::Impl(item_impl) => { - if let Some(decl) = parse_impl(&item_impl) { + if let Some(mut decl) = parse_impl(&item_impl, &content) { + decl.imports = imports.clone(); decls.push(decl); } } Item::Mod(item_mod) => { - if let Some(decl) = parse_mod(&item_mod) { + if let Some(mut decl) = parse_mod(&item_mod) { + decl.imports = imports.clone(); decls.push(decl); } } @@ -91,43 +129,74 @@ fn main() { } } - for item in &syntax.items { - if let Item::Use(item_use) = item { - let path = use_tree_to_string(&item_use.tree); - - decls.push(RustDecl { - name: path.clone(), - type_: "import".to_string(), - start_line: 0, - end_line: 0, - visibility: "".to_string(), - is_async: false, - is_unsafe: false, - generics: vec![], - traits: vec![], - fields: vec![], - methods: vec![], - return_type: None, - parameters: vec![path], - docstring: None, - }); - } - } - let output = serde_json::to_string(&decls).unwrap_or_else(|_| "[]".to_string()); println!("{}", output); } -fn span_start_end(node: &T) -> (usize, usize) { - // Requires proc-macro2 "span-locations" feature. - let span = node.span(); - let start = span.start().line; - let end = span.end().line; - (start, end) +// Better span handling (requires proc-macro2 "span-locations" feature) +fn get_span_start(span: &proc_macro2::Span) -> usize { + span.start().line } -fn parse_function(item_fn: &ItemFn) -> Option { - let (start_line, end_line) = span_start_end(item_fn); +fn get_span_end(span: &proc_macro2::Span) -> usize { + span.end().line +} + +// Enhanced import extraction +fn extract_use_paths(tree: &UseTree) -> Vec { + let mut paths = Vec::new(); + extract_use_paths_recursive(tree, "", &mut paths); + paths +} + +fn extract_use_paths_recursive(tree: &UseTree, current_path: &str, paths: &mut Vec) { + match tree { + UseTree::Path(path) => { + let new_path = if current_path.is_empty() { + path.ident.to_string() + } else { + format!("{}::{}", current_path, path.ident) + }; + extract_use_paths_recursive(&path.tree, &new_path, paths); + } + UseTree::Name(name) => { + let full_path = if current_path.is_empty() { + name.ident.to_string() + } else { + format!("{}::{}", current_path, name.ident) + }; + paths.push(full_path); + } + UseTree::Rename(rename) => { + let full_path = if current_path.is_empty() { + rename.ident.to_string() + } else { + format!("{}::{}", current_path, rename.ident) + }; + paths.push(full_path); + } + UseTree::Glob(_) => { + if !current_path.is_empty() { + paths.push(format!("{}::*", current_path)); + } else { + // top-level glob — represent as "*" + paths.push("*".to_string()); + } + } + UseTree::Group(group) => { + for item in &group.items { + extract_use_paths_recursive(item, current_path, paths); + } + } + // If the AST contains other nodes, ignore them (keeps function exhaustive) + _ => {} + } +} + +// Enhanced function parsing with call detection +fn parse_function(item_fn: &ItemFn, file_content: &str) -> Option { + let start_line = get_span_start(&item_fn.span()); + let end_line = get_span_end(&item_fn.span()); let mut parameters = Vec::new(); for input in &item_fn.sig.inputs { @@ -149,6 +218,9 @@ fn parse_function(item_fn: &ItemFn) -> Option { syn::ReturnType::Type(_, ty) => Some(type_to_string(&*ty)), }; + // Extract function calls from body + let calls = extract_function_calls(&item_fn.block, file_content); + Some(RustDecl { name: item_fn.sig.ident.to_string(), type_: "function".to_string(), @@ -170,16 +242,100 @@ fn parse_function(item_fn: &ItemFn) -> Option { return_type, parameters, docstring: extract_docstring(&item_fn.attrs), + imports: Vec::new(), // Filled by caller + calls, }) } +// Simple function call extraction using regex on tokenized block +fn extract_function_calls(block: &syn::Block, _file_content: &str) -> Vec { + let mut calls = Vec::new(); + let content = quote::ToTokens::to_token_stream(block).to_string(); + + // Basic regex to find function-like identifiers followed by '(' + let re = Regex::new(r"\b([a-zA-Z_][a-zA-Z0-9_]*)\s*\(").unwrap(); + for cap in re.captures_iter(&content) { + if let Some(name) = cap.get(1) { + let call = name.as_str().to_string(); + // Filter out common keywords and types, and constructors (capitalized) + if !is_rust_keyword(&call) && !call.chars().next().map_or(false, |c| c.is_uppercase()) { + calls.push(call); + } + } + } + + // Deduplicate while preserving order + let mut deduped = Vec::new(); + for c in calls { + if !deduped.contains(&c) { + deduped.push(c); + } + } + deduped +} + +fn is_rust_keyword(word: &str) -> bool { + matches!( + word, + "if" | "else" + | "for" + | "while" + | "loop" + | "match" + | "return" + | "break" + | "continue" + | "let" + | "mut" + | "const" + | "static" + | "fn" + | "struct" + | "enum" + | "trait" + | "impl" + | "mod" + | "use" + | "pub" + | "crate" + | "self" + | "super" + | "async" + | "await" + | "unsafe" + | "dyn" + | "move" + | "where" + | "in" + | "type" + | "as" + | "ref" + | "box" + | "true" + | "false" + | "Some" + | "None" + | "Ok" + | "Err" + ) +} + +// Update other parse functions to include the new fields... fn parse_struct(item_struct: &ItemStruct) -> Option { - let (start_line, end_line) = span_start_end(item_struct); + let start_line = get_span_start(&item_struct.span()); + let end_line = get_span_end(&item_struct.span()); let fields: Vec = item_struct .fields .iter() - .filter_map(|field| field.ident.as_ref().map(|ident| ident.to_string())) + .filter_map(|field| { + // field.ident is Option for named fields; for tuple structs it's None + field.ident.as_ref().map(|ident| { + let field_name = ident.to_string(); + let field_type = type_to_string(&field.ty); + format!("{}: {}", field_name, field_type) + }) + }) .collect(); Some(RustDecl { @@ -202,11 +358,14 @@ fn parse_struct(item_struct: &ItemStruct) -> Option { return_type: None, parameters: Vec::new(), docstring: extract_docstring(&item_struct.attrs), + imports: Vec::new(), // Filled by caller + calls: Vec::new(), }) } fn parse_enum(item_enum: &ItemEnum) -> Option { - let (start_line, end_line) = span_start_end(item_enum); + let start_line = get_span_start(&item_enum.span()); + let end_line = get_span_end(&item_enum.span()); let variants: Vec = item_enum .variants @@ -234,11 +393,14 @@ fn parse_enum(item_enum: &ItemEnum) -> Option { return_type: None, parameters: Vec::new(), docstring: extract_docstring(&item_enum.attrs), + imports: Vec::new(), + calls: Vec::new(), }) } fn parse_trait(item_trait: &ItemTrait) -> Option { - let (start_line, end_line) = span_start_end(item_trait); + let start_line = get_span_start(&item_trait.span()); + let end_line = get_span_end(&item_trait.span()); let methods: Vec = item_trait .items @@ -272,11 +434,14 @@ fn parse_trait(item_trait: &ItemTrait) -> Option { return_type: None, parameters: Vec::new(), docstring: extract_docstring(&item_trait.attrs), + imports: Vec::new(), + calls: Vec::new(), }) } -fn parse_impl(item_impl: &ItemImpl) -> Option { - let (start_line, end_line) = span_start_end(item_impl); +fn parse_impl(item_impl: &ItemImpl, file_content: &str) -> Option { + let start_line = get_span_start(&item_impl.span()); + let end_line = get_span_end(&item_impl.span()); let target = type_to_string(&*item_impl.self_ty); @@ -293,21 +458,33 @@ fn parse_impl(item_impl: &ItemImpl) -> Option { .collect(); // If this impl implements a trait, extract its path - let implements = item_impl + let implements: Vec = item_impl .trait_ .as_ref() .map(|(_, path, _)| path_to_string(path)) .into_iter() - .collect::>(); + .collect(); // Prefer the first trait name (if any) as the identifier; otherwise fall back to the target type let name = if !implements.is_empty() { implements[0].clone() } else { - target + target.clone() }; - // Extract any docstring attached to the impl block + // Extract calls across all method bodies in the impl (basic) + let mut calls = Vec::new(); + for item in &item_impl.items { + if let syn::ImplItem::Fn(method) = item { + let body_calls = extract_function_calls(&method.block, file_content); + for c in body_calls { + if !calls.contains(&c) { + calls.push(c); + } + } + } + } + let docstring = extract_docstring(&item_impl.attrs); Some(RustDecl { @@ -330,27 +507,14 @@ fn parse_impl(item_impl: &ItemImpl) -> Option { return_type: None, parameters: Vec::new(), docstring, + imports: Vec::new(), + calls, }) } -/// Pull a concatenated doc‑string from a list of attributes. -/// Handles both `///` comments (converted to `#[doc = "..."]` by syn) -/// and explicit `#[doc = "..."]` attributes. -fn extract_docstring(attrs: &[Attribute]) -> Option { - // Convert the attribute’s token stream back to source‑like text. - let text = attrs - .iter() - .map(|a| a.to_token_stream().to_string()) - .collect::>() - .join("\n"); - - let re = Regex::new(r#"(?m)^\s*(?:///|#\[\s*doc\s*=\s*")(.+?)(?:")?$"#).ok()?; - re.captures(&text) - .and_then(|c| c.get(1).map(|m| m.as_str().to_string())) -} - fn parse_mod(item_mod: &ItemMod) -> Option { - let (start_line, end_line) = span_start_end(item_mod); + let start_line = get_span_start(&item_mod.span()); + let end_line = get_span_end(&item_mod.span()); Some(RustDecl { name: item_mod.ident.to_string(), @@ -367,6 +531,8 @@ fn parse_mod(item_mod: &ItemMod) -> Option { return_type: None, parameters: Vec::new(), docstring: extract_docstring(&item_mod.attrs), + imports: Vec::new(), + calls: Vec::new(), }) } @@ -375,25 +541,28 @@ fn visibility_to_string(vis: &Visibility) -> String { Visibility::Public(_) => "pub".to_string(), Visibility::Restricted(r) => { // pub(crate) or pub(in ...) - format!("pub({})", quote::ToTokens::to_token_stream(r).to_string()) + quote::ToTokens::to_token_stream(r).to_string() } _ => "private".to_string(), } } fn type_to_string(ty: &Type) -> String { - // Simple string representation for common types match ty { Type::Path(type_path) => type_path .path .segments - .last() + .iter() .map(|seg| seg.ident.to_string()) - .unwrap_or_else(|| "unknown".to_string()), + .collect::>() + .join("::"), Type::Reference(type_ref) => { - // type_ref.elem is Box format!("&{}", type_to_string(&*type_ref.elem)) } + Type::Tuple(tuple) => { + let elems: Vec = tuple.elems.iter().map(type_to_string).collect(); + format!("({})", elems.join(", ")) + } _ => quote::ToTokens::to_token_stream(ty).to_string(), } } @@ -414,19 +583,50 @@ fn generic_param_to_string(param: &syn::GenericParam) -> String { } } -fn use_tree_to_string(tree: &syn::UseTree) -> String { - match tree { - syn::UseTree::Path(p) => { - let base = p.ident.to_string(); - let nested = use_tree_to_string(&*p.tree); - format!("{}::{}", base, nested) - } - syn::UseTree::Name(n) => n.ident.to_string(), - syn::UseTree::Rename(r) => r.ident.to_string(), - syn::UseTree::Glob(_) => "*".to_string(), - syn::UseTree::Group(g) => { - let items: Vec = g.items.iter().map(use_tree_to_string).collect(); - items.join("::") +/// Pull a concatenated doc-string from a list of attributes. +/// Handles both `///` comments (converted to `#[doc = "..."]` by syn) +/// and explicit `#[doc = "..."]` attributes. +fn extract_docstring(attrs: &[Attribute]) -> Option { + // Prefer parsing doc attributes explicitly + let mut lines: Vec = Vec::new(); + + for attr in attrs { + // If attribute path is "doc", try to parse the meta and extract the literal + if attr.path().is_ident("doc") { + if let Ok(syn::Meta::NameValue(nv)) = attr.parse_meta() { + if let syn::Lit::Str(litstr) = nv.lit { + lines.push(litstr.value()); + } + } else { + // Fallback: stringify the tokens and try regex extraction + let tok = attr.to_token_stream().to_string(); + let re = Regex::new(r#"(?m)#\s*\[\s*doc\s*=\s*"(.+?)"\s*\]"#).ok(); + if let Some(re) = re { + if let Some(cap) = re.captures(&tok) { + if let Some(m) = cap.get(1) { + lines.push(m.as_str().to_string()); + } + } + } + } + } else { + // capture leading "///" style comments as attributes may not always be "doc" + let tok = attr.to_token_stream().to_string(); + // crude fallback to capture /// "..." inside tokens + let re_line = Regex::new(r#"(?m)///\s*(.+)$"#).ok(); + if let Some(re_line) = re_line { + for cap in re_line.captures_iter(&tok) { + if let Some(m) = cap.get(1) { + lines.push(m.as_str().to_string()); + } + } + } } } + + if lines.is_empty() { + None + } else { + Some(lines.join("\n")) + } }