From 85fe0425b958707b4d4a3fe91a88ca8b6a592870 Mon Sep 17 00:00:00 2001 From: ef3d0c3e Date: Thu, 8 Aug 2024 11:21:58 +0200 Subject: [PATCH] Added auto-registry proc macro --- crates/auto-registry/Cargo.lock | 54 +++++++ crates/auto-registry/Cargo.toml | 13 ++ crates/auto-registry/src/lib.rs | 267 ++++++++++++++++++++++++++++++++ 3 files changed, 334 insertions(+) create mode 100644 crates/auto-registry/Cargo.lock create mode 100644 crates/auto-registry/Cargo.toml create mode 100644 crates/auto-registry/src/lib.rs diff --git a/crates/auto-registry/Cargo.lock b/crates/auto-registry/Cargo.lock new file mode 100644 index 0000000..a5cb56e --- /dev/null +++ b/crates/auto-registry/Cargo.lock @@ -0,0 +1,54 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "auto-registry" +version = "0.0.4" +dependencies = [ + "lazy_static", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "proc-macro2" +version = "1.0.86" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fa76aaf39101c457836aec0ce2316dbdc3ab723cdda1c6bd4e6ad4208acaca7" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "unicode-ident" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" diff --git a/crates/auto-registry/Cargo.toml b/crates/auto-registry/Cargo.toml new file mode 100644 index 0000000..fcf8751 --- /dev/null +++ b/crates/auto-registry/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "auto-registry" +version = "0.0.4" +edition = "2021" + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = { version = "1.0"} +quote = "1.0" +syn = { version = "1.0", features = [ "full" ] } +lazy_static = "1.5.0" diff --git a/crates/auto-registry/src/lib.rs b/crates/auto-registry/src/lib.rs new file mode 100644 index 0000000..5a9196a --- /dev/null +++ b/crates/auto-registry/src/lib.rs @@ -0,0 +1,267 @@ +#![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() +}