Allow arbitrary expressions as pool_size

1 files changed, 9 insertions, 9 deletions

diff --git a/embassy-macros/src/macros/task.rs b/embassy-macros/src/macros/task.rs
index 8c4bf7265..1d30434e9 100644
--- a/embassy-macros/src/macros/task.rs
+++ b/embassy-macros/src/macros/task.rs
@@ -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,)*)) }
         }
     };