kona_common_proc/
lib.rs

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

kona_common_proc/
lib.rs
