kona_std_fpvm_proc/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{
6    ItemFn, LitInt, Result,
7    parse::{Parse, ParseStream},
8    parse_macro_input,
9};
10
11/// The arguments for the `#[client_entry]` attribute proc macro
12struct MacroArgs {
13    /// The heap size to allocate
14    heap_size: LitInt,
15}
16
17impl Parse for MacroArgs {
18    fn parse(input: ParseStream) -> Result<Self> {
19        let heap_size = input.parse()?;
20        Ok(MacroArgs { heap_size })
21    }
22}
23
24#[proc_macro_attribute]
25pub fn client_entry(attr: TokenStream, input: TokenStream) -> TokenStream {
26    let args = parse_macro_input!(attr as MacroArgs);
27    let input_fn = parse_macro_input!(input as ItemFn);
28
29    let heap_size = args.heap_size;
30    let fn_body = &input_fn.block;
31    let fn_name = &input_fn.sig.ident;
32
33    let expanded = quote! {
34        fn #fn_name() -> Result<(), String> {
35            match #fn_body {
36                Ok(_) => kona_std_fpvm::io::exit(0),
37                Err(e) => {
38                    kona_std_fpvm::io::print_err(alloc::format!("Program encountered fatal error: {:?}\n", e).as_ref());
39                    kona_std_fpvm::io::exit(1);
40                }
41            }
42        }
43
44        cfg_if::cfg_if! {
45            if #[cfg(any(target_arch = "mips64", target_arch = "riscv64"))] {
46                const HEAP_SIZE: usize = #heap_size;
47
48                #[doc = "Program entry point"]
49                #[unsafe(no_mangle)]
50                pub extern "C" fn _start() {
51                    kona_std_fpvm::alloc_heap!(HEAP_SIZE);
52                    let _ = #fn_name();
53                }
54
55                #[panic_handler]
56                fn panic(info: &core::panic::PanicInfo) -> ! {
57                    let msg = alloc::format!("Panic: {}", info);
58                    kona_std_fpvm::io::print_err(msg.as_ref());
59                    kona_std_fpvm::io::exit(2)
60                }
61            }
62        }
63    };
64
65    TokenStream::from(expanded)
66}