126 lines
		
	
	
		
			4.2 KiB
		
	
	
	
		
			Rust
		
	
	
	
	
	
			
		
		
	
	
			126 lines
		
	
	
		
			4.2 KiB
		
	
	
	
		
			Rust
		
	
	
	
	
	
| use darling::export::NestedMeta;
 | |
| use darling::FromMeta;
 | |
| use proc_macro2::{Span, TokenStream};
 | |
| use quote::{format_ident, quote};
 | |
| use syn::{parse_quote, Expr, ExprLit, ItemFn, Lit, LitInt, ReturnType, Type};
 | |
| 
 | |
| use crate::util::ctxt::Ctxt;
 | |
| 
 | |
| #[derive(Debug, FromMeta)]
 | |
| struct Args {
 | |
|     #[darling(default)]
 | |
|     pool_size: Option<syn::Expr>,
 | |
| }
 | |
| 
 | |
| pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStream> {
 | |
|     let args = Args::from_list(args).map_err(|e| e.write_errors())?;
 | |
| 
 | |
|     let pool_size = args.pool_size.unwrap_or(Expr::Lit(ExprLit {
 | |
|         attrs: vec![],
 | |
|         lit: Lit::Int(LitInt::new("1", Span::call_site())),
 | |
|     }));
 | |
| 
 | |
|     let ctxt = Ctxt::new();
 | |
| 
 | |
|     if f.sig.asyncness.is_none() {
 | |
|         ctxt.error_spanned_by(&f.sig, "task functions must be async");
 | |
|     }
 | |
|     if !f.sig.generics.params.is_empty() {
 | |
|         ctxt.error_spanned_by(&f.sig, "task functions must not be generic");
 | |
|     }
 | |
|     if !f.sig.generics.where_clause.is_none() {
 | |
|         ctxt.error_spanned_by(&f.sig, "task functions must not have `where` clauses");
 | |
|     }
 | |
|     if !f.sig.abi.is_none() {
 | |
|         ctxt.error_spanned_by(&f.sig, "task functions must not have an ABI qualifier");
 | |
|     }
 | |
|     if !f.sig.variadic.is_none() {
 | |
|         ctxt.error_spanned_by(&f.sig, "task functions must not be variadic");
 | |
|     }
 | |
|     match &f.sig.output {
 | |
|         ReturnType::Default => {}
 | |
|         ReturnType::Type(_, ty) => match &**ty {
 | |
|             Type::Tuple(tuple) if tuple.elems.is_empty() => {}
 | |
|             Type::Never(_) => {}
 | |
|             _ => ctxt.error_spanned_by(
 | |
|                 &f.sig,
 | |
|                 "task functions must either not return a value, return `()` or return `!`",
 | |
|             ),
 | |
|         },
 | |
|     }
 | |
| 
 | |
|     let mut args = Vec::new();
 | |
|     let mut fargs = f.sig.inputs.clone();
 | |
| 
 | |
|     for arg in fargs.iter_mut() {
 | |
|         match arg {
 | |
|             syn::FnArg::Receiver(_) => {
 | |
|                 ctxt.error_spanned_by(arg, "task functions must not have receiver arguments");
 | |
|             }
 | |
|             syn::FnArg::Typed(t) => match t.pat.as_mut() {
 | |
|                 syn::Pat::Ident(id) => {
 | |
|                     id.mutability = None;
 | |
|                     args.push((id.clone(), t.attrs.clone()));
 | |
|                 }
 | |
|                 _ => {
 | |
|                     ctxt.error_spanned_by(arg, "pattern matching in task arguments is not yet supported");
 | |
|                 }
 | |
|             },
 | |
|         }
 | |
|     }
 | |
| 
 | |
|     ctxt.check()?;
 | |
| 
 | |
|     let task_ident = f.sig.ident.clone();
 | |
|     let task_inner_ident = format_ident!("__{}_task", task_ident);
 | |
| 
 | |
|     let mut task_inner = f;
 | |
|     let visibility = task_inner.vis.clone();
 | |
|     task_inner.vis = syn::Visibility::Inherited;
 | |
|     task_inner.sig.ident = task_inner_ident.clone();
 | |
| 
 | |
|     // assemble the original input arguments,
 | |
|     // including any attributes that may have
 | |
|     // been applied previously
 | |
|     let mut full_args = Vec::new();
 | |
|     for (arg, cfgs) in args {
 | |
|         full_args.push(quote!(
 | |
|             #(#cfgs)*
 | |
|             #arg
 | |
|         ));
 | |
|     }
 | |
| 
 | |
|     #[cfg(feature = "nightly")]
 | |
|     let mut task_outer: ItemFn = parse_quote! {
 | |
|         #visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
 | |
|             type Fut = impl ::core::future::Future + 'static;
 | |
|             const POOL_SIZE: usize = #pool_size;
 | |
|             static POOL: ::embassy_executor::raw::TaskPool<Fut, POOL_SIZE> = ::embassy_executor::raw::TaskPool::new();
 | |
|             unsafe { POOL._spawn_async_fn(move || #task_inner_ident(#(#full_args,)*)) }
 | |
|         }
 | |
|     };
 | |
|     #[cfg(not(feature = "nightly"))]
 | |
|     let mut task_outer: ItemFn = parse_quote! {
 | |
|         #visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
 | |
|             const POOL_SIZE: usize = #pool_size;
 | |
|             static POOL: ::embassy_executor::_export::TaskPoolRef = ::embassy_executor::_export::TaskPoolRef::new();
 | |
|             unsafe { POOL.get::<_, POOL_SIZE>()._spawn_async_fn(move || #task_inner_ident(#(#full_args,)*)) }
 | |
|         }
 | |
|     };
 | |
| 
 | |
|     task_outer.attrs.append(&mut task_inner.attrs.clone());
 | |
| 
 | |
|     let result = quote! {
 | |
|         // This is the user's task function, renamed.
 | |
|         // We put it outside the #task_ident fn below, because otherwise
 | |
|         // the items defined there (such as POOL) would be in scope
 | |
|         // in the user's code.
 | |
|         #[doc(hidden)]
 | |
|         #task_inner
 | |
| 
 | |
|         #task_outer
 | |
|     };
 | |
| 
 | |
|     Ok(result)
 | |
| }
 |