embassy/embassy-macros/src/lib.rs
2020-10-19 21:25:54 +02:00

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()
}