115 lines
3.3 KiB
Rust
115 lines
3.3 KiB
Rust
|
#![feature(proc_macro_diagnostic)]
|
||
|
|
|
||
|
|
extern crate proc_macro;
|
||
|
|
|
||
|
|
use darling::FromMeta;
|
||
|
|
use proc_macro::{Diagnostic, Level, Span, TokenStream};
|
||
|
|
use quote::{format_ident, quote};
|
||
|
|
use syn::spanned::Spanned;
|
||
|
|
|
||
|
|
#[derive(Debug, FromMeta)]
|
||
|
|
struct MacroArgs {
|
||
|
|
#[darling(default)]
|
||
|
|
pool_size: Option<usize>,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[proc_macro_attribute]
|
||
|
|
pub fn task(args: TokenStream, item: TokenStream) -> TokenStream {
|
||
|
|
let args = syn::parse_macro_input!(args as syn::AttributeArgs);
|
||
|
|
let mut task_fn = syn::parse_macro_input!(item as syn::ItemFn);
|
||
|
|
|
||
|
|
let args = match MacroArgs::from_list(&args) {
|
||
|
|
Ok(v) => v,
|
||
|
|
Err(e) => {
|
||
|
|
return TokenStream::from(e.write_errors());
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
let pool_size: usize = args.pool_size.unwrap_or(1);
|
||
|
|
|
||
|
|
let mut fail = false;
|
||
|
|
if task_fn.sig.asyncness.is_none() {
|
||
|
|
task_fn
|
||
|
|
.sig
|
||
|
|
.span()
|
||
|
|
.unwrap()
|
||
|
|
.error("task functions must be async")
|
||
|
|
.emit();
|
||
|
|
fail = true;
|
||
|
|
}
|
||
|
|
if task_fn.sig.generics.params.len() != 0 {
|
||
|
|
task_fn
|
||
|
|
.sig
|
||
|
|
.span()
|
||
|
|
.unwrap()
|
||
|
|
.error("task functions must not be generic")
|
||
|
|
.emit();
|
||
|
|
fail = true;
|
||
|
|
}
|
||
|
|
if pool_size < 1 {
|
||
|
|
Span::call_site()
|
||
|
|
.error("pool_size must be 1 or greater")
|
||
|
|
.emit();
|
||
|
|
fail = true
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut arg_names: syn::punctuated::Punctuated<syn::Ident, syn::Token![,]> =
|
||
|
|
syn::punctuated::Punctuated::new();
|
||
|
|
let args = &task_fn.sig.inputs;
|
||
|
|
|
||
|
|
for arg in args.iter() {
|
||
|
|
match arg {
|
||
|
|
syn::FnArg::Receiver(_) => {
|
||
|
|
arg.span()
|
||
|
|
.unwrap()
|
||
|
|
.error("task functions must not have receiver arguments")
|
||
|
|
.emit();
|
||
|
|
fail = true;
|
||
|
|
}
|
||
|
|
syn::FnArg::Typed(t) => match t.pat.as_ref() {
|
||
|
|
syn::Pat::Ident(i) => arg_names.push(i.ident.clone()),
|
||
|
|
_ => {
|
||
|
|
arg.span()
|
||
|
|
.unwrap()
|
||
|
|
.error("pattern matching in task arguments is not yet supporteds")
|
||
|
|
.emit();
|
||
|
|
fail = true;
|
||
|
|
}
|
||
|
|
},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if fail {
|
||
|
|
return TokenStream::new();
|
||
|
|
}
|
||
|
|
|
||
|
|
let name = task_fn.sig.ident.clone();
|
||
|
|
|
||
|
|
let type_name = format_ident!("__embassy_executor_type_{}", name);
|
||
|
|
let pool_name = format_ident!("__embassy_executor_pool_{}", name);
|
||
|
|
let task_fn_name = format_ident!("__embassy_executor_task_{}", name);
|
||
|
|
let create_fn_name = format_ident!("__embassy_executor_create_{}", name);
|
||
|
|
|
||
|
|
let visibility = &task_fn.vis;
|
||
|
|
|
||
|
|
task_fn.sig.ident = task_fn_name.clone();
|
||
|
|
|
||
|
|
let result = quote! {
|
||
|
|
#task_fn
|
||
|
|
#[allow(non_camel_case_types)]
|
||
|
|
type #type_name = impl ::core::future::Future + 'static;
|
||
|
|
|
||
|
|
fn #create_fn_name(#args) -> #type_name {
|
||
|
|
#task_fn_name(#arg_names)
|
||
|
|
}
|
||
|
|
|
||
|
|
#[allow(non_upper_case_globals)]
|
||
|
|
static #pool_name: [::embassy::executor::Task<#type_name>; #pool_size] = [::embassy::executor::Task::new(); #pool_size];
|
||
|
|
|
||
|
|
#visibility fn #name(#args) -> ::embassy::executor::SpawnToken {
|
||
|
|
unsafe { ::embassy::executor::Task::spawn(&#pool_name, || #create_fn_name(#arg_names)) }
|
||
|
|
}
|
||
|
|
};
|
||
|
|
result.into()
|
||
|
|
}
|