grid_rs_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::{quote, quote_spanned};
3use syn::spanned::Spanned;
4use syn::{parse_macro_input, Ident, ItemFn};
5
6/// Entry point for Wasm functions.
7#[proc_macro_attribute]
8pub fn main(_attr: TokenStream, item: TokenStream) -> TokenStream {
9    let ItemFn {
10        attrs,
11        vis,
12        mut sig,
13        block,
14    } = parse_macro_input!(item as ItemFn);
15
16    if sig.asyncness.is_some() {
17        return quote!(compile_error!("async functions are not supported")).into();
18    }
19
20    let orig_ident = sig.ident.clone();
21    let new_ident = Ident::new(&format!("__{orig_ident}"), sig.ident.span());
22    sig.ident = new_ident.clone();
23
24    let main_impl = quote_spanned! {sig.output.span()=>
25        #[cfg(target_arch = "wasm32")]
26        #[unsafe(no_mangle)]
27        pub extern "C" fn call(ptr: i32, len: i32) -> i32 {
28            let result = std::panic::catch_unwind(|| {
29                let input_data = unsafe {
30                    std::slice::from_raw_parts(ptr as *const u8, len as usize)
31                };
32                #new_ident(&input_data)
33            });
34
35            match result {
36                Ok(Ok(data)) => grid_rs::Output::write_all(&data) as i32,
37                Ok(Err(e)) => {
38                    let error = serde_json::json!({ "error": e });
39                    grid_rs::Output::write_all(&serde_json::to_vec(&error).unwrap()) as i32
40                },
41                Err(e) => {
42                    let error = serde_json::json!({
43                        "error": format!("Runtime panic: {:?}", e)
44                    });
45                    grid_rs::Output::write_all(&serde_json::to_vec(&error).unwrap()) as i32
46                }
47            }
48        }
49
50        fn main() {
51            // Native main that does nothing but satisfy the compiler
52            let input = std::env::args().nth(1).unwrap_or_default();
53            let result = #new_ident(input.as_bytes());
54            match result {
55                Ok(data) => println!("{}", String::from_utf8_lossy(&data)),
56                Err(e) => eprintln!("Error: {}", e),
57            }
58        }
59    };
60
61    quote! {
62        #(#attrs)*
63        #vis #sig #block
64
65        #main_impl
66    }
67    .into()
68}