#[cfg(feature = "code-graph")]
use std::path::Path;
#[cfg(feature = "code-graph")]
use tree_sitter::{Node, Parser};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Symbol {
pub name: String,
pub kind: SymbolKind,
pub file_path: String,
pub start_line: usize,
pub end_line: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SymbolKind {
Function,
Struct,
Enum,
Trait,
Impl,
Import,
}
impl std::fmt::Display for SymbolKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SymbolKind::Function => write!(f, "function"),
SymbolKind::Struct => write!(f, "struct"),
SymbolKind::Enum => write!(f, "enum"),
SymbolKind::Trait => write!(f, "trait"),
SymbolKind::Impl => write!(f, "impl"),
SymbolKind::Import => write!(f, "import"),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CallEdge {
pub caller: String,
pub callee: String,
pub file_path: String,
pub line: usize,
}
#[cfg(feature = "code-graph")]
pub struct SymbolExtractor {
parser: Parser,
}
#[cfg(feature = "code-graph")]
impl SymbolExtractor {
pub fn new() -> Result<Self, anyhow::Error> {
let mut parser = Parser::new();
let language = tree_sitter_rust::LANGUAGE;
parser.set_language(&language.into())?;
Ok(Self { parser })
}
pub fn extract(
&mut self,
file_path: &Path,
source: &str,
) -> Result<(Vec<Symbol>, Vec<CallEdge>), anyhow::Error> {
let tree = self
.parser
.parse(source, None)
.ok_or_else(|| anyhow::anyhow!("Failed to parse {}", file_path.display()))?;
let file_path_str = file_path.to_string_lossy().to_string();
let mut symbols = Vec::new();
let mut call_edges = Vec::new();
self.extract_from_node(
tree.root_node(),
source,
&file_path_str,
None, &mut symbols,
&mut call_edges,
);
Ok((symbols, call_edges))
}
fn extract_from_node(
&self,
node: Node,
source: &str,
file_path: &str,
current_function: Option<&str>,
symbols: &mut Vec<Symbol>,
call_edges: &mut Vec<CallEdge>,
) {
let kind = node.kind();
match kind {
"function_item" => {
if let Some(name) = self.extract_function_name(node, source) {
let start_line = node.start_position().row;
let end_line = node.end_position().row;
symbols.push(Symbol {
name: name.clone(),
kind: SymbolKind::Function,
file_path: file_path.to_string(),
start_line,
end_line,
});
for child in node.children(&mut node.walk()) {
self.extract_from_node(
child,
source,
file_path,
Some(&name),
symbols,
call_edges,
);
}
}
}
"struct_item" => {
if let Some(name) = self.extract_type_name(node, source) {
symbols.push(Symbol {
name,
kind: SymbolKind::Struct,
file_path: file_path.to_string(),
start_line: node.start_position().row,
end_line: node.end_position().row,
});
}
}
"enum_item" => {
if let Some(name) = self.extract_type_name(node, source) {
symbols.push(Symbol {
name,
kind: SymbolKind::Enum,
file_path: file_path.to_string(),
start_line: node.start_position().row,
end_line: node.end_position().row,
});
}
}
"trait_item" => {
if let Some(name) = self.extract_type_name(node, source) {
symbols.push(Symbol {
name,
kind: SymbolKind::Trait,
file_path: file_path.to_string(),
start_line: node.start_position().row,
end_line: node.end_position().row,
});
}
}
"impl_item" => {
let impl_name = self.extract_impl_name(node, source);
symbols.push(Symbol {
name: impl_name,
kind: SymbolKind::Impl,
file_path: file_path.to_string(),
start_line: node.start_position().row,
end_line: node.end_position().row,
});
for child in node.children(&mut node.walk()) {
self.extract_from_node(
child,
source,
file_path,
current_function,
symbols,
call_edges,
);
}
}
"use_declaration" => {
let import_path = self.extract_use_path(node, source);
symbols.push(Symbol {
name: import_path,
kind: SymbolKind::Import,
file_path: file_path.to_string(),
start_line: node.start_position().row,
end_line: node.end_position().row,
});
}
"call_expression" => {
if let (Some(caller), Some(callee)) =
(current_function, self.extract_callee(node, source))
{
call_edges.push(CallEdge {
caller: caller.to_string(),
callee,
file_path: file_path.to_string(),
line: node.start_position().row,
});
}
for child in node.children(&mut node.walk()) {
self.extract_from_node(
child,
source,
file_path,
current_function,
symbols,
call_edges,
);
}
}
_ => {
for child in node.children(&mut node.walk()) {
self.extract_from_node(
child,
source,
file_path,
current_function,
symbols,
call_edges,
);
}
}
}
}
fn extract_function_name(&self, node: Node, source: &str) -> Option<String> {
node.child_by_field_name("name")
.map(|n| n.utf8_text(source.as_bytes()).unwrap_or("").to_string())
}
fn extract_type_name(&self, node: Node, source: &str) -> Option<String> {
node.child_by_field_name("name")
.map(|n| n.utf8_text(source.as_bytes()).unwrap_or("").to_string())
}
fn extract_impl_name(&self, node: Node, source: &str) -> String {
let mut result = String::from("impl");
if let Some(type_node) = node.child_by_field_name("type") {
let type_name = type_node.utf8_text(source.as_bytes()).unwrap_or("");
result.push(' ');
result.push_str(type_name);
}
result
}
fn extract_use_path(&self, node: Node, source: &str) -> String {
node.child_by_field_name("argument")
.map(|n| n.utf8_text(source.as_bytes()).unwrap_or("").to_string())
.unwrap_or_else(|| "use".to_string())
}
fn extract_callee(&self, node: Node, source: &str) -> Option<String> {
let raw = node
.child_by_field_name("function")
.map(|n| n.utf8_text(source.as_bytes()).unwrap_or("").to_string())?;
let callee = match raw.rsplit_once('.') {
Some((_, method)) => method.to_string(),
None => raw,
};
if matches!(callee.as_str(), "Some" | "Ok" | "Err" | "None") {
return None;
}
Some(callee)
}
}
#[cfg(feature = "code-graph")]
pub(crate) fn extract_and_store(store: &super::db::Store, key: &str, body: &str) {
let mut extractor = match SymbolExtractor::new() {
Ok(e) => e,
Err(e) => {
tracing::debug!("code-graph: extractor init failed: {e}");
return;
}
};
let path = std::path::Path::new(key);
match extractor.extract(path, body) {
Ok((symbols, call_edges)) => {
for sym in symbols.iter().filter(|s| s.kind != SymbolKind::Import) {
if let Err(e) = store.insert_symbol(
&sym.name,
&sym.kind.to_string(),
key,
sym.start_line,
sym.end_line,
) {
tracing::debug!("code-graph: insert_symbol {key} {}: {e}", sym.name);
}
}
for edge in call_edges {
if let Err(e) = store.insert_call_edge(&edge.caller, &edge.callee, key, edge.line) {
tracing::debug!("code-graph: insert_call_edge {key} {}: {e}", edge.caller);
}
}
for sym in symbols.iter().filter(|s| s.kind == SymbolKind::Import) {
if let Err(e) = store.insert_import(&sym.name, key, sym.start_line) {
tracing::debug!("code-graph: insert_import {key} {}: {e}", sym.name);
}
}
}
Err(e) => tracing::debug!("code-graph: failed to extract symbols from {key}: {e}"),
}
}