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

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{Error, FnArg, ItemFn, PatType};

#[proc_macro_attribute]
pub fn space(_: TokenStream, input: TokenStream) -> TokenStream {
    // Parse ast from token stream
    let ast = syn::parse::<ItemFn>(input).expect("Place this attribute above a function");

    // Signature and output type
    let name = ast.sig.ident;
    let output = ast.sig.output;
    let stub_name = format_ident!("{name}_stub");

    // Stub
    let body = ast.block.stmts;
    let stub = match ast.sig.inputs.first() {
        Some(FnArg::Typed(PatType { pat, ty, .. })) => quote! {
            let #pat = ::space_lib::rmp_serde::from_slice::<#ty>(__input_bytes).unwrap();

            // Actual code
            fn #stub_name(#pat: #ty) #output {
                #(#body);*
            }
            let output = #stub_name(#pat);

            // Serialize output
            let __output_bytes = ::space_lib::rmp_serde::to_vec_named(&output).unwrap().leak();
            ::std::boxed::Box::new(::space_lib::SpaceSlice {
                len: __output_bytes.len(),
                ptr: __output_bytes.as_mut_ptr(),
            })
        },
        _ => Error::new(name.span(), "expected one argument").to_compile_error(),
    };

    quote! {
        #[no_mangle]
        fn #name(ptr: usize) -> ::std::boxed::Box<::space_lib::SpaceSlice> {
            // Deserialize input
            let __input_bytes = unsafe {
                let len = *(ptr as *const usize);
                let data = (ptr + 4) as *mut u8;
                ::std::slice::from_raw_parts(data, len)
            };

            #stub
        }
    }
    .into()
}