Allow arbitrary expressions as pool_size

This commit is contained in:
Dániel Buga 2023-06-25 22:32:24 +02:00
parent aa0ab06645
commit 2809e926cf

View File

@ -1,21 +1,24 @@
use darling::export::NestedMeta;
use darling::FromMeta;
use proc_macro2::TokenStream;
use proc_macro2::{Span, TokenStream};
use quote::{format_ident, quote};
use syn::{parse_quote, ItemFn, ReturnType, Type};
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<usize>,
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: usize = args.pool_size.unwrap_or(1);
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();
@ -46,10 +49,6 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
},
}
if pool_size < 1 {
ctxt.error_spanned_by(&f.sig, "pool_size must be 1 or greater");
}
let mut arg_names = Vec::new();
let mut fargs = f.sig.inputs.clone();
@ -83,7 +82,8 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
let mut task_outer: ItemFn = parse_quote! {
#visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
type Fut = impl ::core::future::Future + 'static;
static POOL: ::embassy_executor::raw::TaskPool<Fut, #pool_size> = ::embassy_executor::raw::TaskPool::new();
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(#(#arg_names,)*)) }
}
};