clusterflux-public/crates/disasmer-macros/src/lib.rs

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse::Parser, parse_macro_input, Expr, ItemFn, Lit, Meta, Token};

#[proc_macro_attribute]
pub fn main(attr: TokenStream, item: TokenStream) -> TokenStream {
    let function = parse_macro_input!(item as ItemFn);
    let function_name = function.sig.ident.to_string();
    let entrypoint_name = descriptor_name(
        attr,
        function_name
            .strip_suffix("_main")
            .unwrap_or(&function_name),
    );
    let descriptor = format_ident!(
        "__DISASMER_ENTRYPOINT_{}",
        function_name.to_ascii_uppercase()
    );

    quote! {
        #function

        #[doc(hidden)]
        pub const #descriptor: ::disasmer::EntrypointDescriptor = ::disasmer::EntrypointDescriptor {
            name: #entrypoint_name,
            function: #function_name,
        };
    }
    .into()
}

#[proc_macro_attribute]
pub fn task(attr: TokenStream, item: TokenStream) -> TokenStream {
    let function = parse_macro_input!(item as ItemFn);
    let function_name = function.sig.ident.to_string();
    let task_name = descriptor_name(attr, &function_name);
    let descriptor = format_ident!("__DISASMER_TASK_{}", function_name.to_ascii_uppercase());

    quote! {
        #function

        #[doc(hidden)]
        pub const #descriptor: ::disasmer::TaskDescriptor = ::disasmer::TaskDescriptor {
            name: #task_name,
            function: #function_name,
            remotely_startable: true,
        };
    }
    .into()
}

fn descriptor_name(attr: TokenStream, default: &str) -> String {
    if attr.is_empty() {
        return default.to_owned();
    }

    let parser = syn::punctuated::Punctuated::<Meta, Token![,]>::parse_terminated;
    let Ok(args) = parser.parse(attr) else {
        return default.to_owned();
    };

    for meta in args {
        let Meta::NameValue(name_value) = meta else {
            continue;
        };
        if !name_value.path.is_ident("name") {
            continue;
        }
        if let Expr::Lit(expr) = name_value.value {
            if let Lit::Str(name) = expr.lit {
                return name.value();
            }
        }
    }

    default.to_owned()
}