use std::path::Path;
use proc_macro2::TokenStream;
use quote::ToTokens;
use syn::visit::Visit;
use syn::visit_mut::VisitMut;
use crate::error::SourceResult;
use crate::ops;
use crate::visitor::{IdentifierCollector, ImportCollector};
#[derive(Debug, Clone)]
pub struct RustAST {
file: syn::File,
source: Option<String>,
}
impl RustAST {
pub fn parse(source: &str) -> SourceResult<Self> {
let file = syn::parse_file(source)?;
Ok(Self {
file,
source: Some(source.to_string()),
})
}
pub fn from_file(path: &Path) -> SourceResult<Self> {
let source = std::fs::read_to_string(path)?;
Self::parse(&source)
}
pub fn file(&self) -> &syn::File {
&self.file
}
pub fn file_mut(&mut self) -> &mut syn::File {
&mut self.file
}
pub fn source(&self) -> Option<&str> {
self.source.as_deref()
}
pub fn to_string_pretty(&self) -> String {
prettyplease::unparse(&self.file)
}
pub fn to_token_stream(&self) -> TokenStream {
self.file.to_token_stream()
}
pub fn collect_imports(&self) -> Vec<&syn::ItemUse> {
let mut collector = ImportCollector::new();
collector.visit_file(&self.file);
collector.imports
}
pub fn collect_used_identifiers(&self) -> std::collections::HashSet<String> {
let mut collector = IdentifierCollector::new();
collector.visit_file(&self.file);
collector.identifiers
}
pub fn find_unused_imports(&self) -> Vec<UnusedImport> {
ops::RemoveUnusedImports::detect(self)
}
pub fn remove_unused_imports(&mut self) -> Vec<UnusedImport> {
ops::RemoveUnusedImports::apply(self)
}
pub fn visit_mut<V: VisitMut>(&mut self, visitor: &mut V) {
visitor.visit_file_mut(&mut self.file);
}
pub fn visit<'a, V: Visit<'a>>(&'a self, visitor: &mut V) {
visitor.visit_file(&self.file);
}
pub fn items(&self) -> &[syn::Item] {
&self.file.items
}
pub fn items_mut(&mut self) -> &mut Vec<syn::Item> {
&mut self.file.items
}
pub fn filter_items<F>(&self, predicate: F) -> Vec<&syn::Item>
where
F: Fn(&syn::Item) -> bool,
{
self.file.items.iter().filter(|i| predicate(i)).collect()
}
pub fn functions(&self) -> Vec<&syn::ItemFn> {
self.file
.items
.iter()
.filter_map(|item| {
if let syn::Item::Fn(f) = item {
Some(f)
} else {
None
}
})
.collect()
}
pub fn structs(&self) -> Vec<&syn::ItemStruct> {
self.file
.items
.iter()
.filter_map(|item| {
if let syn::Item::Struct(s) = item {
Some(s)
} else {
None
}
})
.collect()
}
pub fn impls(&self) -> Vec<&syn::ItemImpl> {
self.file
.items
.iter()
.filter_map(|item| {
if let syn::Item::Impl(i) = item {
Some(i)
} else {
None
}
})
.collect()
}
}
impl std::fmt::Display for RustAST {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.to_string_pretty())
}
}
#[derive(Debug, Clone)]
pub struct UnusedImport {
pub path: String,
pub name: String,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_simple() {
let ast = RustAST::parse("fn main() {}").unwrap();
assert_eq!(ast.functions().len(), 1);
}
#[test]
fn test_parse_with_imports() {
let ast = RustAST::parse("use std::io;\nuse std::fs;\nfn main() {}").unwrap();
assert_eq!(ast.collect_imports().len(), 2);
}
#[test]
fn test_to_string() {
let ast = RustAST::parse("fn main() {}").unwrap();
let output = ast.to_string();
assert!(output.contains("fn main"));
}
#[test]
fn test_collect_identifiers() {
let ast = RustAST::parse(
r#"
use std::io;
fn main() {
let x = io::stdin();
println!("{}", x);
}
"#,
)
.unwrap();
let idents = ast.collect_used_identifiers();
assert!(idents.contains("io"));
assert!(idents.contains("x"));
assert!(idents.contains("println"));
}
}