Repository: nexus-xyz/nexus-zkvm
File: runtime/macros/src/entry.rs (Rust)

Code before:
            }
        }
    })
}

pub fn profile(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input: ItemFn = syn::parse2(item).unwrap();
    let visibility = &input.vis;
    let name = &input.sig.ident;
    let inputs = &input.sig.inputs;
    let output = &input.sig.output;
    let block = &input.block;
    let generics = &input.sig.generics;
    let where_clause = &input.sig.generics.where_clause;

    let result = quote! {
        #visibility fn #name #generics (#inputs) #output #where_clause {
            profile_cycles!("^:{}", stringify!(#name));
            let result = (|| #block)();
            profile_cycles!("$:{}", stringify!(#name));
            result

        }
    };

    result
}


Review comment (suggestion):
1. this module is called `entry`, which is for the entry function. To keep it clean, let's move it to a separate one.
2. Just to be more strict, let's return an error if `attr` is not empty, like in `main`. Alternatively, users should be able to specify custom name for the function to avoid name clash. Or this macro should specify its full path.
3. use absolute path to external methods, like `nexus_rt::profile_cycles`. `nexus_rt` (very unlikely) but can happen to be renamed. https://crates.io/crates/proc-macro-crate
4. You should only unwrap in procedural macros when you can prove that it's practically unreachable. Making this macro to panic is very easy, a compile error should be returned instead.

Finally, this can be simplified.
```rust
let nexus_rt = {
    // get the import name
};
let input: ItemFn = syn::parse2(item)?;
let vis = &input.vis;
let block = &input.block;
let fn_name = &input.sig.ident.to_string();

quote! {
    #vis #sig {
        #nexus_rt :: profile_cycles!("^:{}", fn_name);
        let result = #block;
        #nexus_rt :: profile_cycles!("$:{}", fn_name);
        result
    }
}.into()
```

Code after:
use proc_macro2::TokenStream;
use quote::quote;
use syn::{Error, ItemFn};

use super::parse_args::parse_memory_limit;

pub fn main(args: TokenStream, input: TokenStream) -> Result<TokenStream, Error> {
    let func: ItemFn = syn::parse2(input)?;
    let memlimit = parse_memory_limit(args)?;

    let fn_sig = &func.sig;
    if &fn_sig.ident != "main" {
        return Err(Error::new_spanned(
            &fn_sig.ident,
            "function name must be `main`",
        ));
    }
    if fn_sig.asyncness.is_some() {
        return Err(Error::new_spanned(
            fn_sig,
            "`main` function is not allowed to be `async`",
        ));
    }
    if !fn_sig.generics.params.empty_or_trailing() {
        return Err(Error::new_spanned(
            &fn_sig.generics,
            "`main` function is not allowed to have generic parameters",
        ));
    }
    if func.sig.generics.where_clause.is_some() {
        return Err(Error::new_spanned(
            &fn_sig.generics.where_clause,
            "`main` function is not allowed to have a `where` clause",
        ));
    }

    Ok(quote! {
        const _: fn() = main;

        #[cfg_attr(target_arch = "riscv32", no_mangle)]
        #[allow(unused)]
        #func

        #[cfg(target_arch = "riscv32")]
        #[doc(hidden)]
        mod __private {
            #[no_mangle]
            pub fn get_stack_size() -> i32 {
                #memlimit
            }
        }
    })
}