mu_macro 0.1.4

extern crate alloc;
use alloc::string::ToString;
use proc_macro::TokenStream;
use proc_macro2::Ident;
use quote::quote;
use syn::{parse_macro_input, ItemFn};

#[proc_macro_attribute]
pub fn mu_kernel(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(item as ItemFn);
    TokenStream::from(generate_kernel_code(input_fn, false))
}

#[proc_macro_attribute]
pub fn mu_non_blocking_kernel(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(item as ItemFn);
    TokenStream::from(generate_kernel_code(input_fn, true))
}

const KERNEL_TYPE_NON_BLOCKING: u64 = 1;

fn generate_kernel_code(input_fn: ItemFn, is_non_blocking: bool) -> proc_macro2::TokenStream {
    let fn_name = &input_fn.sig.ident;
    let mut fn_name_str = fn_name.to_string();
    fn_name_str.push('\0');
    let mu_kernel_func_name =
        Ident::new(&alloc::format!("__mu_kernel_{}", fn_name), fn_name.span());
    let kernel_info_name = Ident::new(&alloc::format!("__mu_func_{}", fn_name), fn_name.span());

    let inputs = &input_fn.sig.inputs;
    let output = &input_fn.sig.output;
    let body = &input_fn.block;

    let mut expanded = quote! {
        #[no_mangle]
        extern "C" fn #fn_name(#inputs) #output #body

        #[no_mangle]
        extern "C" fn #mu_kernel_func_name(mode: u32, ndarray_value_bit: u32, argv: *mut core::ffi::c_char) {
            use core::ffi::c_char;
            unsafe {
                let func_ptr: extern "C" fn(#inputs) #output = #fn_name;
                mu_lib::call_helper::call_main(func_ptr, mode, ndarray_value_bit, argv as *mut u8);
            }
        }

        #[used]
        #[no_mangle]
        static #kernel_info_name: mu_lib::mu::KernelInfo = mu_lib::mu::KernelInfo {
            name: #fn_name_str.as_ptr() as *const core::ffi::c_char,
            func: #mu_kernel_func_name,
        };
    };

    if is_non_blocking {
        expanded.extend(quote! {
            #[used]
            #[no_mangle]
            static __mu_kernel_type__: u64 = #KERNEL_TYPE_NON_BLOCKING;
        });
    }

    expanded
}