ryo_source/

ast.rs

//! Core RustAST implementation.

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};

/// Rust source code with parsed AST.
///
/// This is the main entry point for AST manipulation.
/// It wraps a `syn::File` and provides high-level operations.
#[derive(Debug, Clone)]
pub struct RustAST {
    /// The parsed AST.
    file: syn::File,
    /// Original source (optional, for span information).
    source: Option<String>,
}

impl RustAST {
    /// Parse Rust source code into an AST.
    pub fn parse(source: &str) -> SourceResult<Self> {
        let file = syn::parse_file(source)?;
        Ok(Self {
            file,
            source: Some(source.to_string()),
        })
    }

    /// Load and parse a Rust file.
    pub fn from_file(path: &Path) -> SourceResult<Self> {
        let source = std::fs::read_to_string(path)?;
        Self::parse(&source)
    }

    /// Get a reference to the underlying syn::File.
    pub fn file(&self) -> &syn::File {
        &self.file
    }

    /// Get a mutable reference to the underlying syn::File.
    pub fn file_mut(&mut self) -> &mut syn::File {
        &mut self.file
    }

    /// Get the original source if available.
    pub fn source(&self) -> Option<&str> {
        self.source.as_deref()
    }

    /// Convert the AST to a pretty-printed source code.
    pub fn to_string_pretty(&self) -> String {
        prettyplease::unparse(&self.file)
    }

    /// Get the token stream.
    pub fn to_token_stream(&self) -> TokenStream {
        self.file.to_token_stream()
    }

    // ==================== Analysis ====================

    /// Collect all use statements.
    pub fn collect_imports(&self) -> Vec<&syn::ItemUse> {
        let mut collector = ImportCollector::new();
        collector.visit_file(&self.file);
        collector.imports
    }

    /// Collect all identifiers used in the code (excluding imports).
    pub fn collect_used_identifiers(&self) -> std::collections::HashSet<String> {
        let mut collector = IdentifierCollector::new();
        collector.visit_file(&self.file);
        collector.identifiers
    }

    /// Find unused imports.
    pub fn find_unused_imports(&self) -> Vec<UnusedImport> {
        ops::RemoveUnusedImports::detect(self)
    }

    // ==================== Transformations ====================

    /// Remove all unused imports. Returns the removed imports.
    pub fn remove_unused_imports(&mut self) -> Vec<UnusedImport> {
        ops::RemoveUnusedImports::apply(self)
    }

    /// Apply a custom visitor mutation.
    pub fn visit_mut<V: VisitMut>(&mut self, visitor: &mut V) {
        visitor.visit_file_mut(&mut self.file);
    }

    /// Apply a custom visitor (read-only).
    pub fn visit<'a, V: Visit<'a>>(&'a self, visitor: &mut V) {
        visitor.visit_file(&self.file);
    }

    // ==================== Item Access ====================

    /// Get all items in the file.
    pub fn items(&self) -> &[syn::Item] {
        &self.file.items
    }

    /// Get mutable access to all items.
    pub fn items_mut(&mut self) -> &mut Vec<syn::Item> {
        &mut self.file.items
    }

    /// Filter items by type.
    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()
    }

    /// Get all functions.
    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()
    }

    /// Get all structs.
    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()
    }

    /// Get all impl blocks.
    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())
    }
}

/// Information about an unused import.
#[derive(Debug, Clone)]
pub struct UnusedImport {
    /// The imported path (e.g., "std::io").
    pub path: String,
    /// The name that would be used in code (e.g., "io").
    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"));
    }
}