"""Rust language plugin with advanced features."""
import logging
import re
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional
from ...plugin_base import (
IndexShard,
Reference,
SearchOpts,
SearchResult,
SymbolDef,
)
from ...storage.sqlite_store import SQLiteStore
from ..generic_treesitter_plugin import GenericTreeSitterPlugin
from .cargo_integration import CargoIntegration
from .module_resolver import RustModuleResolver
from .trait_analyzer import RustTraitAnalyzer
logger = logging.getLogger(__name__)
class SpecializedPluginBase(GenericTreeSitterPlugin):
"""Base class for specialized language plugins."""
def __init__(
self,
language_config: Dict[str, Any],
sqlite_store: Optional[SQLiteStore] = None,
enable_semantic: bool = True,
) -> None:
super().__init__(language_config, sqlite_store, enable_semantic)
self._specialized_features = {}
def add_feature(self, name: str, feature: Any) -> None:
"""Add a specialized feature to this plugin."""
self._specialized_features[name] = feature
def get_feature(self, name: str) -> Any:
"""Get a specialized feature by name."""
return self._specialized_features.get(name)
class RustPlugin(SpecializedPluginBase):
"""Rust language plugin with module resolution, trait analysis, and Cargo integration."""
# Rust-specific tree-sitter query for comprehensive symbol extraction
RUST_QUERY = """
; Functions
(function_item
name: (identifier) @function.name) @function
; Methods in impl blocks
(impl_item
body: (declaration_list
(function_item
name: (identifier) @method.name))) @method
; Structs
(struct_item
name: (type_identifier) @struct.name) @struct
; Enums
(enum_item
name: (type_identifier) @enum.name) @enum
; Traits
(trait_item
name: (type_identifier) @trait.name) @trait
; Type aliases
(type_item
name: (type_identifier) @type_alias.name) @type_alias
; Constants
(const_item
name: (identifier) @constant.name) @constant
; Statics
(static_item
name: (identifier) @static.name) @static
; Modules
(mod_item
name: (identifier) @module.name) @module
; Use declarations
(use_declaration
argument: (use_tree) @use) @use_decl
; Macros
(macro_definition
name: (identifier) @macro.name) @macro
; Impl blocks
(impl_item
trait: (type_identifier)? @impl.trait
type: (_) @impl.type) @impl
"""
def __init__(
self, sqlite_store: Optional[SQLiteStore] = None, enable_semantic: bool = True
) -> None:
"""Initialize Rust plugin with specialized features."""
config = {
"code": "rust",
"name": "Rust",
"extensions": [".rs"],
"symbols": [
"function_item",
"struct_item",
"enum_item",
"trait_item",
"impl_item",
"mod_item",
"type_alias",
"const_item",
"static_item",
"macro_definition",
],
"query": self.RUST_QUERY,
}
super().__init__(config, sqlite_store, enable_semantic)
# Initialize specialized features
workspace_root = Path.cwd()
self.add_feature("module_resolver", RustModuleResolver(workspace_root))
self.add_feature("trait_analyzer", RustTraitAnalyzer())
self.add_feature("cargo_integration", CargoIntegration(workspace_root))
# Cache for macro expansions
self._macro_cache = {}
# Additional indexing for Rust-specific features
self._trait_impls = {} # trait -> list of implementing types
self._module_tree = {} # module hierarchy
logger.info("Initialized Rust plugin with specialized features")
def indexFile(self, path: str | Path, content: str) -> IndexShard:
"""Index a Rust file with enhanced analysis."""
# First, do the base indexing
shard = super().indexFile(path, content)
# Add Rust-specific analysis
if isinstance(path, str):
path = Path(path)
# Analyze traits and implementations
trait_analyzer = self.get_feature("trait_analyzer")
analysis = trait_analyzer.analyze_file(content)
# Store trait implementations
for impl_info in analysis["implementations"]:
if impl_info.trait_name:
if impl_info.trait_name not in self._trait_impls:
self._trait_impls[impl_info.trait_name] = []
self._trait_impls[impl_info.trait_name].append(
{
"type": impl_info.target_type,
"file": str(path),
"line": impl_info.line_number,
}
)
# Analyze module structure
module_resolver = self.get_feature("module_resolver")
mod_decls = module_resolver.find_mod_declarations(content)
for mod_name, line_num in mod_decls:
self._module_tree[mod_name] = {
"file": str(path),
"line": line_num,
"parent": self._get_parent_module(path),
}
# Check if this is a Cargo.toml file
if path.name == "Cargo.toml":
cargo = self.get_feature("cargo_integration")
crate_info = cargo.parse_cargo_toml(path)
if crate_info:
logger.info(f"Indexed Cargo.toml for crate: {crate_info.name}")
# Add enhanced symbols with Rust-specific metadata
enhanced_symbols = []
for symbol in shard["symbols"]:
enhanced_symbol = symbol.copy()
# Add trait bounds for generic functions
if symbol["kind"] in ["function", "method"]:
bounds = self._extract_trait_bounds(content, symbol["line"])
if bounds:
enhanced_symbol["trait_bounds"] = bounds
# Add lifetime information
if symbol["kind"] in ["function", "struct", "impl"]:
lifetimes = self._extract_lifetimes(content, symbol["line"])
if lifetimes:
enhanced_symbol["lifetimes"] = lifetimes
enhanced_symbols.append(enhanced_symbol)
shard["symbols"] = enhanced_symbols
shard["rust_analysis"] = analysis
return shard
def getDefinition(self, symbol: str) -> SymbolDef | None:
"""Get definition with module resolution support."""
# First try the base implementation
definition = super().getDefinition(symbol)
if definition:
return definition
# Try module resolution if symbol contains ::
if "::" in symbol:
module_resolver = self.get_feature("module_resolver")
parts = symbol.split("::")
# Try to resolve the module path
for ext in self.file_extensions:
for path in Path(".").rglob(f"*{ext}"):
resolved = module_resolver.resolve_module_path(path, "::".join(parts[:-1]))
if resolved and resolved.exists():
try:
content = resolved.read_text(encoding="utf-8")
# Look for the last part of the symbol
if self.parser:
tree = self.parser.parse(content.encode("utf-8"))
symbols = self._extract_symbols(tree, content)
for sym in symbols:
if sym["symbol"] == parts[-1]:
return SymbolDef(
symbol=symbol,
kind=sym["kind"],
language="rust",
signature=sym["signature"],
doc=None,
defined_in=str(resolved),
line=sym["line"],
span=(
sym["line"],
sym.get("end_line", sym["line"]),
),
)
except Exception:
continue
return None
def findReferences(self, symbol: str) -> list[Reference]:
"""Find references with use statement analysis."""
refs = super().findReferences(symbol)
# Add references from use statements
module_resolver = self.get_feature("module_resolver")
for ext in self.file_extensions:
for path in Path(".").rglob(f"*{ext}"):
try:
content = path.read_text(encoding="utf-8")
use_stmts = module_resolver.find_use_statements(content)
for use_path, line_num in use_stmts:
# Check if symbol is imported
if symbol in use_path:
refs.append(Reference(file=str(path), line=line_num))
except Exception:
continue
return refs
def search(self, query: str, opts: SearchOpts | None = None) -> Iterable[SearchResult]:
"""Enhanced search with Rust-specific features."""
# Get base results
results = list(super().search(query, opts))
# Check if searching for trait implementations
if query.startswith("impl:"):
trait_name = query[5:].strip()
impl_results = self._search_trait_implementations(trait_name)
results.extend(impl_results)
# Check if searching for crate dependencies
elif query.startswith("dep:"):
dep_name = query[4:].strip()
dep_results = self._search_dependencies(dep_name)
results.extend(dep_results)
# Check if searching for macros
elif query.startswith("macro:"):
macro_name = query[6:].strip()
macro_results = self._search_macros(macro_name)
results.extend(macro_results)
return results
def _search_trait_implementations(self, trait_name: str) -> List[SearchResult]:
"""Search for implementations of a specific trait."""
results = []
if trait_name in self._trait_impls:
for impl_info in self._trait_impls[trait_name]:
results.append(
SearchResult(
file=impl_info["file"],
line=impl_info["line"],
column=0,
text=f"impl {trait_name} for {impl_info['type']}",
score=1.0,
)
)
return results
def _search_dependencies(self, dep_name: str) -> List[SearchResult]:
"""Search for crate dependencies."""
results = []
cargo = self.get_feature("cargo_integration")
for cargo_path in cargo.find_cargo_tomls():
crate_info = cargo.parse_cargo_toml(cargo_path)
if crate_info and dep_name in crate_info.dependencies:
results.append(
SearchResult(
file=str(cargo_path),
line=1,
column=0,
text=f"dependency: {dep_name} = {crate_info.dependencies[dep_name]}",
score=1.0,
)
)
return results
def _search_macros(self, macro_name: str) -> List[SearchResult]:
"""Search for macro definitions and usage."""
results = []
pattern = re.compile(rf"\b{re.escape(macro_name)}!\s*[\(\[{{]")
for ext in self.file_extensions:
for path in Path(".").rglob(f"*{ext}"):
try:
content = path.read_text(encoding="utf-8")
lines = content.split("\n")
for i, line in enumerate(lines):
if pattern.search(line):
results.append(
SearchResult(
file=str(path),
line=i + 1,
column=(line.index(macro_name) if macro_name in line else 0),
text=line.strip(),
score=0.9,
)
)
except Exception:
continue
return results
def _extract_trait_bounds(self, content: str, line_num: int) -> List[str]:
"""Extract trait bounds from a function or impl block."""
lines = content.split("\n")
if line_num - 1 >= len(lines):
return []
bounds = []
# Look for where clause
where_pattern = re.compile(r"where\s+(.*?)(?:\{|;)", re.DOTALL)
text = "\n".join(lines[line_num - 1 : min(line_num + 10, len(lines))])
match = where_pattern.search(text)
if match:
where_clause = match.group(1)
# Parse individual bounds
bound_parts = re.split(r",\s*(?![^<>]*>)", where_clause)
for part in bound_parts:
if ":" in part:
bounds.append(part.strip())
return bounds
def _extract_lifetimes(self, content: str, line_num: int) -> List[str]:
"""Extract lifetime parameters from a definition."""
lines = content.split("\n")
if line_num - 1 >= len(lines):
return []
line = lines[line_num - 1]
lifetime_pattern = re.compile(r"'(\w+)(?:\s*[:,])")
lifetimes = []
for match in lifetime_pattern.finditer(line):
lifetime = match.group(1)
if lifetime not in ["static"]: # Exclude 'static
lifetimes.append(f"'{lifetime}")
return lifetimes
def _get_parent_module(self, file_path: Path) -> Optional[str]:
"""Get the parent module for a file."""
# Check if this is in a module directory
if file_path.parent.name != "src" and (file_path.parent / "mod.rs").exists():
return file_path.parent.name
# Check if this is a submodule
if file_path.stem != "mod" and file_path.stem != "lib" and file_path.stem != "main":
return None
return None
def get_crate_info(self, crate_name: str) -> Optional[Dict[str, Any]]:
"""Get information about a specific crate."""
cargo = self.get_feature("cargo_integration")
crate_info = cargo.find_crate_by_name(crate_name)
if crate_info:
return {
"name": crate_info.name,
"version": crate_info.version,
"path": str(crate_info.path),
"dependencies": crate_info.dependencies,
"features": crate_info.features,
}
return None
def analyze_trait_hierarchy(self) -> Dict[str, List[str]]:
"""Get the complete trait hierarchy."""
trait_analyzer = self.get_feature("trait_analyzer")
all_traits = []
# Collect all traits from indexed files
for ext in self.file_extensions:
for path in Path(".").rglob(f"*{ext}"):
try:
content = path.read_text(encoding="utf-8")
analysis = trait_analyzer.analyze_file(content)
all_traits.extend(analysis["traits"])
except Exception:
continue
return trait_analyzer.get_trait_hierarchy(all_traits)
