ryo-source 0.1.0

High-speed Rust AST manipulation engine
Documentation
//! AST visitors for analysis.

use std::collections::HashSet;
use syn::visit::Visit;

/// Collects all `use` statements from the AST.
pub struct ImportCollector<'ast> {
    pub imports: Vec<&'ast syn::ItemUse>,
}

impl<'ast> ImportCollector<'ast> {
    pub fn new() -> Self {
        Self { imports: vec![] }
    }
}

impl<'ast> Visit<'ast> for ImportCollector<'ast> {
    fn visit_item_use(&mut self, node: &'ast syn::ItemUse) {
        self.imports.push(node);
        syn::visit::visit_item_use(self, node);
    }
}

/// Collects all identifiers used in the code (excluding use statements).
pub struct IdentifierCollector {
    pub identifiers: HashSet<String>,
    /// Track if we're inside a use statement (to skip those identifiers).
    in_use_statement: bool,
}

impl IdentifierCollector {
    pub fn new() -> Self {
        Self {
            identifiers: HashSet::new(),
            in_use_statement: false,
        }
    }
}

impl<'ast> Visit<'ast> for IdentifierCollector {
    fn visit_item_use(&mut self, node: &'ast syn::ItemUse) {
        // Skip identifiers inside use statements
        self.in_use_statement = true;
        syn::visit::visit_item_use(self, node);
        self.in_use_statement = false;
    }

    fn visit_ident(&mut self, ident: &'ast proc_macro2::Ident) {
        if !self.in_use_statement {
            self.identifiers.insert(ident.to_string());
        }
    }

    fn visit_path(&mut self, path: &'ast syn::Path) {
        if !self.in_use_statement {
            // For paths like `io::stdin()`, we want to capture "io"
            if let Some(first_segment) = path.segments.first() {
                self.identifiers.insert(first_segment.ident.to_string());
            }
        }
        syn::visit::visit_path(self, path);
    }
}

/// Extracts the "used name" from a UseTree.
/// For `use std::io;` this returns "io".
/// For `use std::io as stdio;` this returns "stdio".
pub fn extract_use_names(tree: &syn::UseTree) -> Vec<String> {
    let mut names = Vec::new();
    extract_use_names_recursive(tree, &mut names);
    names
}

fn extract_use_names_recursive(tree: &syn::UseTree, names: &mut Vec<String>) {
    match tree {
        syn::UseTree::Path(path) => {
            // `use std::io;` - recurse into the path
            extract_use_names_recursive(&path.tree, names);
        }
        syn::UseTree::Name(name) => {
            // `use foo;` or end of path - this is the imported name
            names.push(name.ident.to_string());
        }
        syn::UseTree::Rename(rename) => {
            // `use foo as bar;` - the alias is what's used
            names.push(rename.rename.to_string());
        }
        syn::UseTree::Glob(_) => {
            // `use foo::*;` - can't track individual names
            // We could mark as "always used" to be safe
        }
        syn::UseTree::Group(group) => {
            // `use foo::{bar, baz};` - recurse into each
            for tree in &group.items {
                extract_use_names_recursive(tree, names);
            }
        }
    }
}

/// Gets the full path of a UseTree as a string.
pub fn use_tree_to_path(tree: &syn::UseTree) -> String {
    match tree {
        syn::UseTree::Path(path) => {
            let rest = use_tree_to_path(&path.tree);
            if rest.is_empty() {
                path.ident.to_string()
            } else {
                format!("{}::{}", path.ident, rest)
            }
        }
        syn::UseTree::Name(name) => name.ident.to_string(),
        syn::UseTree::Rename(rename) => {
            format!("{} as {}", rename.ident, rename.rename)
        }
        syn::UseTree::Glob(_) => "*".to_string(),
        syn::UseTree::Group(group) => {
            let items: Vec<_> = group.items.iter().map(use_tree_to_path).collect();
            format!("{{{}}}", items.join(", "))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_simple_use() {
        let use_tree: syn::UseTree = syn::parse_quote!(std::io);
        let names = extract_use_names(&use_tree);
        assert_eq!(names, vec!["io"]);
    }

    #[test]
    fn test_extract_renamed_use() {
        let use_tree: syn::UseTree = syn::parse_quote!(std::io as stdio);
        let names = extract_use_names(&use_tree);
        assert_eq!(names, vec!["stdio"]);
    }

    #[test]
    fn test_extract_group_use() {
        let use_tree: syn::UseTree = syn::parse_quote!(std::{io, fs});
        let names = extract_use_names(&use_tree);
        assert_eq!(names, vec!["io", "fs"]);
    }

    #[test]
    fn test_use_tree_to_path() {
        let use_tree: syn::UseTree = syn::parse_quote!(std::io);
        assert_eq!(use_tree_to_path(&use_tree), "std::io");
    }
}