#![feature(proc_macro_span)] use std::cell::RefCell; use std::collections::HashMap; use lazy_static::lazy_static; use proc_macro::TokenStream; use quote::quote; use std::sync::Mutex; use syn::parse::Parse; use syn::parse::ParseStream; use syn::parse_macro_input; use syn::ItemStruct; lazy_static! { /// The registry, each key corresponds to an identifier that needs to be /// valid in the context of the [`genegenerate_registry`] macro. static ref REGISTRY: Mutex>>> = Mutex::new(RefCell::new(HashMap::new())); } /// Arguments for the [`auto_registry`] proc macro struct AutoRegistryArgs { /// The registry name registry: syn::LitStr, /// The absolute path to the struct, if not specified the macro will try /// to automatically infer the full path. path: Option, } /// Parser for [`AutoRegistryArgs`] impl Parse for AutoRegistryArgs { fn parse(input: ParseStream) -> syn::Result { let mut registry = None; let mut path = None; loop { let key: syn::Ident = input.parse()?; input.parse::()?; let value: syn::LitStr = input.parse()?; match key.to_string().as_str() { "registry" => registry = Some(value), "path" => path = Some(value), _ => { return Err(syn::Error::new( key.span(), format!( "Unknown attribute `{}`, excepted `registry` or `path`", key.to_string() ), )) } } if input.is_empty() { break; } input.parse::()?; } if registry.is_none() { return Err(syn::Error::new( input.span(), "Missing required attribute `registry`".to_string(), )); } Ok(AutoRegistryArgs { registry: registry.unwrap(), path, }) } } /// The proc macro used on a struct to add it to the registry /// /// # Attributes /// - registry: (String) Name of the registry to collect the struct into /// - path: (Optional String) The crate path in which the struct is located /// If left empty, the path will be try to be automatically-deduced #[proc_macro_attribute] pub fn auto_registry(attr: TokenStream, input: TokenStream) -> TokenStream { let args = parse_macro_input!(attr as AutoRegistryArgs); let input = parse_macro_input!(input as ItemStruct); let ident = &input.ident; let path = if let Some(path) = args.path { format!("{}::{}", path.value(), ident.to_string().as_str()) } else { // Attempt to get the path in a hacky way in case the path wasn't // specified as an attribute to the macro let path = match input .ident .span() .unwrap() .source_file() .path() .canonicalize() { Ok(path) => path, Err(e) => { return syn::Error::new( input.ident.span(), format!("Failed to canonicalize path: {}", e), ) .to_compile_error() .into(); } }; let crate_path = std::env::var("CARGO_MANIFEST_DIR").unwrap(); let relative_path = path.strip_prefix(&crate_path).unwrap(); let relative_path_str = relative_path.to_string_lossy(); // Remove the first path component e.g "src/" let pos = if let Some(pos) = relative_path_str.find("/") { pos + 1 } else { 0 }; let module_path = relative_path_str .split_at(pos) .1 .strip_suffix(".rs") .unwrap() .replace("/", "::"); if module_path.is_empty() { format!("crate::{}", ident.to_string()) } else { format!("crate::{module_path}::{}", ident.to_string()) } }; let reg_mtx = REGISTRY.lock().unwrap(); let mut reg_borrow = reg_mtx.borrow_mut(); if let Some(ref mut vec) = reg_borrow.get_mut(args.registry.value().as_str()) { vec.push(path); } else { reg_borrow.insert(args.registry.value(), vec![path]); } quote! { #input } .into() } /// Arguments for the [`generate_registry`] proc macro struct GenerateRegistryArgs { /// The registry name registry: syn::LitStr, /// The target, i.e the generated function name target: syn::Ident, /// The maker macro, takes all constructed items and processes them maker: syn::Expr, /// The return type for the function return_type: syn::Type, } /// Parser for [`GenerateRegistryArgs`] impl Parse for GenerateRegistryArgs { fn parse(input: ParseStream) -> syn::Result { let mut registry = None; let mut target = None; let mut maker = None; let mut return_type = None; loop { let key: syn::Ident = input.parse()?; input.parse::()?; match key.to_string().as_str() { "registry" => registry = Some(input.parse()?), "target" => target = Some(input.parse()?), "maker" => maker = Some(input.parse()?), "return_type" => return_type = Some(input.parse()?), _ => { return Err(syn::Error::new( key.span(), format!( "Unknown attribute `{}`, excepted `registry` or `target`", key.to_string() ), )) } } if input.is_empty() { break; } input.parse::()?; } if registry.is_none() { return Err(syn::Error::new( input.span(), "Missing required attribute `registry`".to_string(), )); } else if target.is_none() { return Err(syn::Error::new( input.span(), "Missing required attribute `target`".to_string(), )); } else if maker.is_none() { return Err(syn::Error::new( input.span(), "Missing required attribute `maker`".to_string(), )); } else if return_type.is_none() { return Err(syn::Error::new( input.span(), "Missing required attribute `return_type`".to_string(), )); } Ok(GenerateRegistryArgs { registry: registry.unwrap(), target: target.unwrap(), maker: maker.unwrap(), return_type: return_type.unwrap(), }) } } /// The proc macro that generates the function to build the registry /// /// # Attributes /// - registry: (String) Name of the registry to generate /// - target: (Identifier) Name of the resulting function /// - maker: (Macro) A macro that will take all the newly constructed objects /// comma-separated and create the resulting expression /// - return_type: (Type) The return type of the generated function. /// Must match the type of the macro invocation #[proc_macro_attribute] pub fn generate_registry(attr: TokenStream, input: TokenStream) -> TokenStream { let args = parse_macro_input!(attr as GenerateRegistryArgs); let reg_mtx = REGISTRY.lock().unwrap(); let mut stream = proc_macro2::TokenStream::new(); if let Some(names) = reg_mtx.borrow().get(args.registry.value().as_str()) { for name in names { let struct_name: proc_macro2::TokenStream = name.parse().unwrap(); stream.extend(quote::quote_spanned!(proc_macro2::Span::call_site() => #struct_name::new(), )); } } else { panic!( "Unable to find registry item with key=`{}`", args.registry.value() ); } let function = args.target; //proc_macro2::Ident::new(args.target.value().as_str(), proc_macro2::Span::call_site()); let return_type = args.return_type; let maker = args.maker; let rest: proc_macro2::TokenStream = input.into(); quote! { fn #function() -> #return_type { #maker!( #stream ) } #rest } .into() }