use anyhow::Result;
use async_trait::async_trait;
use std::path::Path;
use syn::{visit::Visit, File as SynFile, ItemEnum, ItemFn, ItemImpl, ItemStruct, ItemTrait};
use super::LanguageStrategy;
use crate::ast::core::{
AstDag, AstKind, ClassKind, FunctionKind, Language, NodeFlags, UnifiedAstNode,
};
pub struct RustStrategy {
}
impl Default for RustStrategy {
fn default() -> Self {
Self::new()
}
}
impl RustStrategy {
#[must_use]
pub fn new() -> Self {
Self {}
}
fn parse_syn_file(&self, content: &str) -> Result<SynFile> {
syn::parse_file(content).map_err(|e| anyhow::anyhow!("Rust parse error: {e}"))
}
fn convert_to_dag(&self, syn_file: &SynFile) -> AstDag {
let mut dag = AstDag::new();
let mut visitor = RustAstVisitor::new(&mut dag);
visitor.visit_file(syn_file);
dag
}
}
#[async_trait]
impl LanguageStrategy for RustStrategy {
fn language(&self) -> Language {
Language::Rust
}
fn can_parse(&self, path: &Path) -> bool {
path.extension()
.and_then(|ext| ext.to_str())
.is_some_and(|ext| ext == "rs")
}
async fn parse_file(&self, _path: &Path, content: &str) -> Result<AstDag> {
let syn_file = self.parse_syn_file(content)?;
Ok(self.convert_to_dag(&syn_file))
}
fn extract_imports(&self, ast: &AstDag) -> Vec<String> {
let mut imports = Vec::new();
for i in 0..ast.nodes.len() {
if let Some(node) = ast.nodes.get(i as u32) {
if node.flags.has(NodeFlags::IMPORT) {
imports.push(format!("import_{i}"));
}
}
}
imports
}
fn extract_functions(&self, ast: &AstDag) -> Vec<UnifiedAstNode> {
let mut functions = Vec::new();
for i in 0..ast.nodes.len() {
if let Some(node) = ast.nodes.get(i as u32) {
if matches!(node.kind, AstKind::Function(_)) {
functions.push(node.clone());
}
}
}
functions
}
fn extract_types(&self, ast: &AstDag) -> Vec<UnifiedAstNode> {
let mut types = Vec::new();
for i in 0..ast.nodes.len() {
if let Some(node) = ast.nodes.get(i as u32) {
if matches!(node.kind, AstKind::Class(_) | AstKind::Type(_)) {
types.push(node.clone());
}
}
}
types
}
fn calculate_complexity(&self, ast: &AstDag) -> (u32, u32) {
let mut cyclomatic = 1;
let mut cognitive = 0;
for i in 0..ast.nodes.len() {
if let Some(node) = ast.nodes.get(i as u32) {
if node.flags.has(NodeFlags::CONTROL_FLOW) {
cyclomatic += 1;
cognitive += 1;
}
}
}
(cyclomatic, cognitive)
}
}
struct RustAstVisitor<'a> {
dag: &'a mut AstDag,
current_parent: Option<u32>,
}
impl<'a> RustAstVisitor<'a> {
fn new(dag: &'a mut AstDag) -> Self {
Self {
dag,
current_parent: None,
}
}
#[allow(dead_code)]
fn add_node(&mut self, kind: AstKind) -> u32 {
let mut node = UnifiedAstNode::new(kind, Language::Rust);
if let Some(parent) = self.current_parent {
node.parent = parent;
}
self.dag.add_node(node)
}
}
impl Visit<'_> for RustAstVisitor<'_> {
fn visit_item_fn(&mut self, node: &ItemFn) {
let mut ast_node =
UnifiedAstNode::new(AstKind::Function(FunctionKind::Regular), Language::Rust);
if node.sig.asyncness.is_some() {
ast_node.flags.set(NodeFlags::ASYNC);
}
let key = self.dag.add_node(ast_node);
let old_parent = self.current_parent;
self.current_parent = Some(key);
syn::visit::visit_item_fn(self, node);
self.current_parent = old_parent;
}
fn visit_item_struct(&mut self, node: &ItemStruct) {
let ast_node = UnifiedAstNode::new(AstKind::Class(ClassKind::Struct), Language::Rust);
let key = self.dag.add_node(ast_node);
let old_parent = self.current_parent;
self.current_parent = Some(key);
syn::visit::visit_item_struct(self, node);
self.current_parent = old_parent;
}
fn visit_item_enum(&mut self, node: &ItemEnum) {
let ast_node = UnifiedAstNode::new(AstKind::Class(ClassKind::Enum), Language::Rust);
let key = self.dag.add_node(ast_node);
let old_parent = self.current_parent;
self.current_parent = Some(key);
syn::visit::visit_item_enum(self, node);
self.current_parent = old_parent;
}
fn visit_item_trait(&mut self, node: &ItemTrait) {
let ast_node = UnifiedAstNode::new(AstKind::Class(ClassKind::Trait), Language::Rust);
let key = self.dag.add_node(ast_node);
let old_parent = self.current_parent;
self.current_parent = Some(key);
syn::visit::visit_item_trait(self, node);
self.current_parent = old_parent;
}
fn visit_item_impl(&mut self, node: &ItemImpl) {
let ast_node = UnifiedAstNode::new(AstKind::Class(ClassKind::Regular), Language::Rust);
let key = self.dag.add_node(ast_node);
let old_parent = self.current_parent;
self.current_parent = Some(key);
syn::visit::visit_item_impl(self, node);
self.current_parent = old_parent;
}
}