Export Macro

export_macro.rs

//!
//! An export macro for Rust.
//! This is my first procedural macro!
//! Feel free to use this however you like.
//! I am developing a better version on the side.
//!

// Macro usages

#[macro_use]
extern crate proc_macro_error;

// Usages

use proc_macro::TokenStream as PTokenStream;
use proc_macro2::{Delimiter, Ident, Span, TokenStream, TokenTree};
use quote::TokenStreamExt;

// Macros

/// # export
/// A quick `#[export]` macro.
/// Exports all or a set list of things from the given `mod` statement.
///
/// ## Examples:
///
/// ```rust
/// #[export] // Exports all items declared in `thing_1`.
/// pub mod thing_1;
///
/// #[export(utility_1, utility_2)] // Exports only `utility_1` and `utility_2` from `thing_3`.
/// pub mod thing_3;
/// ````
#[proc_macro_error]
#[proc_macro_attribute]
pub fn export(attr: PTokenStream, item: PTokenStream) -> PTokenStream {
    let attr = TokenStream::from(attr);
    let item = TokenStream::from(item);
    let item_iter = item.clone().into_iter().collect::<Vec<TokenTree>>();

    let mut pub_ = None;
    let mut pub_suffix = None;
    let mut is_mod = false;
    let mut mod_idx = 0;
    let mut name = String::new();

    if let Some(TokenTree::Ident(ident)) = item_iter.get(0) {
        if Ident::new("pub", ident.span()).eq(ident) {
            pub_ = Some(TokenTree::Ident(ident.clone()));
            mod_idx = 1;
        }
    }

    if let Some(TokenTree::Group(group)) = item_iter.get(1) {
        if group.delimiter().eq(&Delimiter::Parenthesis) {
            mod_idx = 2;
            pub_suffix = Some(TokenTree::Group(group.clone()));
        }
    }

    if let Some(TokenTree::Ident(ident)) = item_iter.get(mod_idx) {
        if Ident::new("mod", ident.span()).eq(ident) {
            is_mod = true;
        }
    }

    if !is_mod {
        abort_call_site!(
            "The provided input is not a `[pub|pub(crate)|pub(super)] mod` statement!"
        );
    }

    if let Some(TokenTree::Ident(ident)) = item_iter.get(mod_idx + 1) {
        name = ident.to_string();
    }

    if name == String::new() {
        abort_call_site!("Could not get a mod name! Expected: not nothing, got: nothing");
    }

    let name_i = Ident::new(&name, Span::call_site());

    let new_item = quote! {
        #pub_ #pub_suffix mod #name_i;
    };

    let attr = attr.into_iter().collect::<Vec<TokenTree>>();

    let attr_content = attr
        .split(|v| {
            if let TokenTree::Punct(p) = v {
                p.as_char() == ','
            } else {
                false
            }
        })
        .flatten()
        .map(|v| v.clone())
        .collect::<Vec<TokenTree>>();

    let mut stream = TokenStream::new();

    if attr_content.len() == 0 {
        stream.append_all(quote! {
            pub use #name_i :: *;
        });
    } else {
        for item in attr_content {
            stream.append_all(quote! {
                pub use #name_i :: #item;
            });
        }
    }

    PTokenStream::from(quote! {
        #new_item
        #stream
    })
}