axplat_macros/
lib.rs

1//! Macros to define and access a per-CPU data structure.
2
3use proc_macro::TokenStream;
4use proc_macro2::Span;
5use quote::quote;
6use syn::{Error, FnArg, ItemFn, ItemTrait, ReturnType, TraitItem};
7
8fn compiler_error(err: Error) -> TokenStream {
9    err.to_compile_error().into()
10}
11
12fn common_main(item: TokenStream, arg_num: usize, export_name: &str, err_msg: &str) -> TokenStream {
13    let main = syn::parse_macro_input!(item as ItemFn);
14    let mut err = if let ReturnType::Type(_, ty) = &main.sig.output {
15        quote! { #ty }.to_string() != "!"
16    } else {
17        true
18    };
19
20    let args = &main.sig.inputs;
21    for arg in args.iter() {
22        if let FnArg::Typed(pat) = arg {
23            let ty = &pat.ty;
24            if quote! { #ty }.to_string() != "usize" {
25                err = true;
26                break;
27            }
28        }
29    }
30    if args.len() != arg_num {
31        err = true;
32    }
33
34    if err {
35        compiler_error(Error::new(Span::call_site(), err_msg))
36    } else {
37        quote! {
38            #[unsafe(export_name = #export_name)]
39            #main
40        }
41        .into()
42    }
43}
44
45/// Marks a function to be called on the primary core after the platform
46/// initialization.
47///
48/// The function signature must be `fn(cpu_id: usize, arg: usize) -> !`, where
49/// `cpu_id` is the logical CPU ID (0, 1, ..., N-1, N is the number of CPU
50/// cores on the platform), and `arg` is passed from the bootloader (typically
51/// the device tree blob address).
52///
53/// # Example
54///
55/// ```rust
56/// # use axplat_macros as axplat;
57/// #[axplat::main]
58/// fn primary_main(cpu_id: usize, arg: usize) -> ! {
59///     todo!() // Your code here
60/// }
61#[proc_macro_attribute]
62pub fn main(attr: TokenStream, item: TokenStream) -> TokenStream {
63    if !attr.is_empty() {
64        return compiler_error(Error::new(
65            Span::call_site(),
66            "expect an empty attribute or `#[axplat::main]`",
67        ));
68    };
69    common_main(
70        item,
71        2,
72        "__axplat_main",
73        "expect a function with type `fn(cpu_id: usize, arg: usize) -> !`",
74    )
75}
76
77/// Marks a function to be called on the secondary cores after the platform
78/// initialization.
79///
80/// The function signature must be `fn(cpu_id: usize) -> !`, where `cpu_id` is
81/// the logical CPU ID (0, 1, ..., N-1, N is the number of CPU cores on the
82/// platform).
83///
84/// # Example
85///
86/// ```rust
87/// # use axplat_macros as axplat;
88/// #[axplat::secondary_main]
89/// fn secondary_main(cpu_id: usize) -> ! {
90///     todo!() // Your code here
91/// }
92#[proc_macro_attribute]
93pub fn secondary_main(attr: TokenStream, item: TokenStream) -> TokenStream {
94    if !attr.is_empty() {
95        return compiler_error(Error::new(
96            Span::call_site(),
97            "expect an empty attribute or `#[axplat::secondary_main]`",
98        ));
99    };
100    common_main(
101        item,
102        1,
103        "__axplat_secondary_main",
104        "expect a function with type `fn(cpu_id: usize) -> !`",
105    )
106}
107
108#[doc(hidden)]
109#[proc_macro_attribute]
110pub fn def_plat_interface(attr: TokenStream, item: TokenStream) -> TokenStream {
111    if !attr.is_empty() {
112        return compiler_error(Error::new(
113            Span::call_site(),
114            "expect an empty attribute: `#[def_plat_interface]`",
115        ));
116    }
117
118    let trait_ast = syn::parse_macro_input!(item as ItemTrait);
119    let trait_name = &trait_ast.ident;
120
121    let mut fn_list = vec![];
122    for item in &trait_ast.items {
123        if let TraitItem::Fn(method) = item {
124            let attrs = &method.attrs;
125            let sig = &method.sig;
126            let fn_name = &sig.ident;
127
128            let mut args = vec![];
129            for arg in &sig.inputs {
130                match arg {
131                    FnArg::Receiver(_) => {
132                        return compiler_error(Error::new_spanned(
133                            arg,
134                            "`self` is not allowed in the interface definition",
135                        ));
136                    }
137                    FnArg::Typed(ty) => args.push(ty.pat.clone()),
138                }
139            }
140
141            fn_list.push(quote! {
142                #(#attrs)*
143                #[inline]
144                pub #sig {
145                    crate::__priv::call_interface!(#trait_name::#fn_name, #(#args),* )
146                }
147            });
148        }
149    }
150
151    quote! {
152        #[crate::__priv::def_interface]
153        #trait_ast
154
155        #(#fn_list)*
156    }
157    .into()
158}