kona_common_proc/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
extern crate proc_macro;

use proc_macro::TokenStream;
use quote::quote;
use syn::{
    parse::{Parse, ParseStream},
    parse_macro_input, ItemFn, LitInt, Result,
};

/// The arguments for the `#[client_entry]` attribute proc macro
struct MacroArgs {
    /// The heap size to allocate
    heap_size: LitInt,
}

impl Parse for MacroArgs {
    fn parse(input: ParseStream) -> Result<Self> {
        let heap_size = input.parse()?;
        Ok(MacroArgs { heap_size })
    }
}

#[proc_macro_attribute]
pub fn client_entry(attr: TokenStream, input: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as MacroArgs);
    let input_fn = parse_macro_input!(input as ItemFn);

    let heap_size = args.heap_size;
    let fn_body = &input_fn.block;
    let fn_name = &input_fn.sig.ident;

    let expanded = quote! {
        fn #fn_name() -> Result<(), String> {
            match #fn_body {
                Ok(_) => kona_common::io::exit(0),
                Err(e) => {
                    kona_common::io::print_err(alloc::format!("Program encountered fatal error: {:?}\n", e).as_ref());
                    kona_common::io::exit(1);
                }
            }
        }

        cfg_if::cfg_if! {
            if #[cfg(any(target_arch = "mips", target_arch = "riscv64"))] {
                const HEAP_SIZE: usize = #heap_size;

                #[doc = "Program entry point"]
                #[no_mangle]
                pub extern "C" fn _start() {
                    kona_common::alloc_heap!(HEAP_SIZE);
                    let _ = #fn_name();
                }

                #[panic_handler]
                fn panic(info: &core::panic::PanicInfo) -> ! {
                    let msg = alloc::format!("Panic: {}", info);
                    kona_common::io::print_err(msg.as_ref());
                    kona_common::io::exit(2)
                }
            }
        }
    };

    TokenStream::from(expanded)
}