rustextract

Cargo.toml

[package]
name = "rustextract"
version = "0.1.0"
edition = "2021"

[dependencies]
syn = { version = "2.0", features = ["full", "visit-mut", "parsing", "printing"] }
quote = "1.0"
proc-macro2 = "1.0"
walkdir = "2.3"
prettyplease = "0.2"

main.rs

use std::path::Path;
use std::fs;
use std::io::{self, Read};
use syn::{parse_file, Item, ImplItem, TraitItem, Attribute};
use syn::visit_mut::{self, VisitMut};
use walkdir::WalkDir;
use quote::ToTokens;
use proc_macro2::TokenStream;

struct CodeProcessor;

impl VisitMut for CodeProcessor {
    fn visit_file_mut(&mut self, file: &mut syn::File) {
        // Process each item in the file
        for item in &mut file.items {
            self.visit_item_mut(item);
        }

        // Remove empty items
        file.items.retain(|item| !is_empty_item(item));
    }

    fn visit_item_mut(&mut self, item: &mut Item) {
        // Check if this item should be removed
        if should_remove_item(item) {
            *item = create_empty_item();
            return;
        }

        // Process item based on its type
        match item {
            Item::Fn(func) => {
                // Keep signature, remove body
                func.block.stmts.clear();
            },
            Item::Mod(module) => {
                // Process module contents
                if let Some((_, items)) = &mut module.content {
                    for sub_item in items.iter_mut() {
                        self.visit_item_mut(sub_item);
                    }
                    items.retain(|sub_item| !is_empty_item(sub_item));
                }
            },
            Item::Impl(impl_block) => {
                // Process impl items
                for impl_item in &mut impl_block.items {
                    self.visit_impl_item_mut(impl_item);
                }

                // Remove empty impl items
                impl_block.items.retain(|impl_item| !is_empty_impl_item(impl_item));
            },
            Item::Trait(trait_def) => {
                // Process trait items
                for trait_item in &mut trait_def.items {
                    self.visit_trait_item_mut(trait_item);
                }

                // Remove empty trait items
                trait_def.items.retain(|trait_item| !is_empty_trait_item(trait_item));
            },
            _ => {
                // Process other item types recursively
                visit_mut::visit_item_mut(self, item);
            }
        }
    }

    fn visit_impl_item_mut(&mut self, item: &mut ImplItem) {
        match item {
            ImplItem::Fn(method) => {
                // Check if method should be removed
                if is_test(&method.attrs) || is_non_windows(&method.attrs) || !is_public(&method.vis) {
                    *item = create_empty_impl_item();
                } else {
                    // Keep signature, remove body
                    method.block.stmts.clear();
                }
            },
            ImplItem::Const(constant) => {
                // Keep only public constants
                if !is_public(&constant.vis) {
                    *item = create_empty_impl_item();
                }
            },
            ImplItem::Type(type_item) => {
                // Keep only public type items
                if !is_public(&type_item.vis) {
                    *item = create_empty_impl_item();
                }
            },
            _ => visit_mut::visit_impl_item_mut(self, item),
        }
    }

    fn visit_trait_item_mut(&mut self, item: &mut TraitItem) {
        match item {
            TraitItem::Fn(method) => {
                // Check if method should be removed
                if is_test(&method.attrs) || is_non_windows(&method.attrs) {
                    *item = create_empty_trait_item();
                } else if let Some(block) = &mut method.default {
                    // Keep signature, remove default implementation body
                    block.stmts.clear();
                }
            },
            _ => visit_mut::visit_trait_item_mut(self, item),
        }
    }
}

// Helper functions

// Create empty items for removal
fn create_empty_item() -> Item {
    Item::Verbatim(TokenStream::new())
}

fn create_empty_impl_item() -> ImplItem {
    ImplItem::Verbatim(TokenStream::new())
}

fn create_empty_trait_item() -> TraitItem {
    TraitItem::Verbatim(TokenStream::new())
}

// Check if an item is empty
fn is_empty_item(item: &Item) -> bool {
    matches!(item, Item::Verbatim(ts) if ts.is_empty())
}

fn is_empty_impl_item(item: &ImplItem) -> bool {
    matches!(item, ImplItem::Verbatim(ts) if ts.is_empty())
}

fn is_empty_trait_item(item: &TraitItem) -> bool {
    matches!(item, TraitItem::Verbatim(ts) if ts.is_empty())
}

// Check if an item should be removed
fn should_remove_item(item: &Item) -> bool {
    match item {
        // Keep only public structs
        Item::Struct(struct_item) => {
            is_test(&struct_item.attrs) || is_non_windows(&struct_item.attrs) || !is_public(&struct_item.vis)
        },
        // Keep only public enums
        Item::Enum(enum_item) => {
            is_test(&enum_item.attrs) || is_non_windows(&enum_item.attrs) || !is_public(&enum_item.vis)
        },
        // Keep only public traits
        Item::Trait(trait_item) => {
            is_test(&trait_item.attrs) || is_non_windows(&trait_item.attrs) || !is_public(&trait_item.vis)
        },
        // Keep modules since they might contain public structs, enums or traits
        Item::Mod(mod_item) => {
            is_test(&mod_item.attrs) || is_non_windows(&mod_item.attrs) || mod_item.ident == "tests"
        },
        // Keep impls since they might be for public structs, enums or traits
        Item::Impl(impl_item) => {
            is_test(&impl_item.attrs) || is_non_windows(&impl_item.attrs)
        },
        // Remove everything else
        _ => true,
    }
}

// Check if an item is public
fn is_public(vis: &syn::Visibility) -> bool {
    matches!(vis, syn::Visibility::Public(_))
}

// Check if an item is a test
fn is_test(attrs: &[Attribute]) -> bool {
    attrs.iter().any(|attr| {
        let path_str = attr.path().to_token_stream().to_string();

        // Direct test attribute
        if path_str == "test" {
            return true;
        }

        // cfg(test) attribute
        if path_str == "cfg" {
            let tokens = attr.to_token_stream().to_string();
            return tokens.contains("test");
        }

        // cfg_attr with test
        if path_str == "cfg_attr" {
            let tokens = attr.to_token_stream().to_string();
            return tokens.contains("test");
        }

        false
    }) || false
}

// Check if an item is for a non-Windows platform
fn is_non_windows(attrs: &[Attribute]) -> bool {
    attrs.iter().any(|attr| {
        let path_str = attr.path().to_token_stream().to_string();

        if path_str == "cfg" {
            let attr_str = attr.to_token_stream().to_string();

            // Check for non-Windows platforms
            let non_windows_platforms = ["unix", "macos", "linux", "android", "ios", "wasm"];
            for platform in non_windows_platforms.iter() {
                if attr_str.contains(platform) {
                    return true;
                }
            }

            // Check for negated Windows
            if attr_str.contains("not(windows)") || attr_str.contains("not (windows)") {
                return true;
            }

            // Not a non-Windows attribute
            return false;
        }

        false
    }) || false
}

// Process a single Rust file
fn process_file(file_path: &Path) -> io::Result<String> {
    // Read the file
    let mut content = String::new();
    fs::File::open(file_path)?.read_to_string(&mut content)?;

    // Parse the Rust code
    let mut syntax_tree = match parse_file(&content) {
        Ok(syntax_tree) => syntax_tree,
        Err(e) => {
            eprintln!("Error parsing {}: {}", file_path.display(), e);
            return Ok(String::new());
        }
    };

    // Process the syntax tree
    let mut processor = CodeProcessor;
    processor.visit_file_mut(&mut syntax_tree);

    // Use prettyplease to format the code correctly
    let processed_code = prettyplease::unparse(&syntax_tree);

    Ok(processed_code)
}

// Main function to process all Rust files in a directory
fn main() -> io::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    let dir_path = if args.len() > 1 {
        &args[1]
    } else {
        "."
    };

    // Find all Rust files in the directory
    for entry in WalkDir::new(dir_path).into_iter().filter_map(|e| e.ok()) {
        let path = entry.path();
        if path.is_file() && path.extension().map_or(false, |ext| ext == "rs") {
            match process_file(path) {
                Ok(processed_code) => {
                    if !processed_code.is_empty() {
                        println!("// Processed from: {}", path.display());
                        println!("{}", processed_code);
                        println!();
                    }
                },
                Err(e) => eprintln!("Error processing {}: {}", path.display(), e),
            }
        }
    }

    Ok(())
}