ableos/kernel/ktest_macro/src/lib.rs

extern crate proc_macro;
extern crate quote;
extern crate syn;

use {
    proc_macro::TokenStream,
    quote::quote,
    syn::{parse::Parse, parse_macro_input, Expr, ItemFn, Token}
};

struct KtestInput {
    lhs: Expr,
    _comma: Token![,],
    rhs: Expr,
}

impl Parse for KtestInput {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        Ok(Self {
            lhs: input.parse()?,
            _comma: input.parse()?,
            rhs: input.parse()?,
        })
    }
}

#[proc_macro]
pub fn ktest_eq(item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as KtestInput);

    let lhs = input.lhs;
    let rhs = input.rhs;

    let out = quote! {
        if #lhs != #rhs {
            return Err(name);
        }
    };
    TokenStream::from(out)
}

#[proc_macro]
pub fn ktest_neq(item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as KtestInput);

    let lhs = input.lhs;
    let rhs = input.rhs;

    let out = quote! {
        if #lhs == #rhs {
            return Err(name);
        }
    };
    TokenStream::from(out)
}

#[proc_macro_attribute]
pub fn ktest(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    let test_name = &input.sig.ident;
    let test_string = test_name.to_string();
    let static_var_name = syn::Ident::new(
        &format!("__ktest_{}", test_name).to_uppercase(),
        test_name.span(),
    );

    let block = &input.block;
    let out = quote! {
        #[cfg(feature = "ktest")]
        fn #test_name() -> Result<String, String> {
            use crate::alloc::string::ToString;
            let name = #test_string.to_string();

            #block

            return Ok(name);
        }

        #[cfg(feature = "ktest")]
        #[unsafe(link_section = ".note.ktest")]
        #[used]
        pub static #static_var_name: fn() -> Result<String, String> = #test_name;
    };

    TokenStream::from(out)
}