1use std::str::FromStr;
2
3use darling::export::NestedMeta;
4use darling::FromMeta;
5use proc_macro2::{Span, TokenStream};
6use quote::{format_ident, quote};
7use syn::visit::{self, Visit};
8use syn::{Expr, ExprLit, Lit, LitInt, ReturnType, Type};
9
10use crate::util::*;
11
12#[derive(Debug, FromMeta, Default)]
13struct Args {
14 #[darling(default)]
15 pool_size: Option<syn::Expr>,
16 /// Use this to override the `embassy_executor` crate path. Defaults to `::embassy_executor`.
17 #[darling(default)]
18 embassy_executor: Option<syn::Expr>,
19}
20
21pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
22 let mut errors = TokenStream::new();
23
24 // If any of the steps for this macro fail, we still want to expand to an item that is as close
25 // to the expected output as possible. This helps out IDEs such that completions and other
26 // related features keep working.
27 let f: ItemFn = match syn::parse2(item.clone()) {
28 Ok(x) => x,
29 Err(e) => return token_stream_with_error(item, e),
30 };
31
32 let args = match NestedMeta::parse_meta_list(args) {
33 Ok(x) => x,
34 Err(e) => return token_stream_with_error(item, e),
35 };
36
37 let args = match Args::from_list(&args) {
38 Ok(x) => x,
39 Err(e) => {
40 errors.extend(e.write_errors());
41 Args::default()
42 }
43 };
44
45 let pool_size = args.pool_size.unwrap_or(Expr::Lit(ExprLit {
46 attrs: vec![],
47 lit: Lit::Int(LitInt::new("1", Span::call_site())),
48 }));
49
50 let embassy_executor = args
51 .embassy_executor
52 .unwrap_or(Expr::Verbatim(TokenStream::from_str("::embassy_executor").unwrap()));
53
54 if f.sig.asyncness.is_none() {
55 error(&mut errors, &f.sig, "task functions must be async");
56 }
57 if !f.sig.generics.params.is_empty() {
58 error(&mut errors, &f.sig, "task functions must not be generic");
59 }
60 if !f.sig.generics.where_clause.is_none() {
61 error(&mut errors, &f.sig, "task functions must not have `where` clauses");
62 }
63 if !f.sig.abi.is_none() {
64 error(&mut errors, &f.sig, "task functions must not have an ABI qualifier");
65 }
66 if !f.sig.variadic.is_none() {
67 error(&mut errors, &f.sig, "task functions must not be variadic");
68 }
69 match &f.sig.output {
70 ReturnType::Default => {}
71 ReturnType::Type(_, ty) => match &**ty {
72 Type::Tuple(tuple) if tuple.elems.is_empty() => {}
73 Type::Never(_) => {}
74 _ => error(
75 &mut errors,
76 &f.sig,
77 "task functions must either not return a value, return `()` or return `!`",
78 ),
79 },
80 }
81
82 let mut args = Vec::new();
83 let mut fargs = f.sig.inputs.clone();
84
85 for arg in fargs.iter_mut() {
86 match arg {
87 syn::FnArg::Receiver(_) => {
88 error(&mut errors, arg, "task functions must not have `self` arguments");
89 }
90 syn::FnArg::Typed(t) => {
91 check_arg_ty(&mut errors, &t.ty);
92 match t.pat.as_mut() {
93 syn::Pat::Ident(id) => {
94 id.mutability = None;
95 args.push((id.clone(), t.attrs.clone()));
96 }
97 _ => {
98 error(
99 &mut errors,
100 arg,
101 "pattern matching in task arguments is not yet supported",
102 );
103 }
104 }
105 }
106 }
107 }
108
109 let task_ident = f.sig.ident.clone();
110 let task_inner_ident = format_ident!("__{}_task", task_ident);
111
112 let mut task_inner = f.clone();
113 let visibility = task_inner.vis.clone();
114 task_inner.vis = syn::Visibility::Inherited;
115 task_inner.sig.ident = task_inner_ident.clone();
116
117 // assemble the original input arguments,
118 // including any attributes that may have
119 // been applied previously
120 let mut full_args = Vec::new();
121 for (arg, cfgs) in args {
122 full_args.push(quote!(
123 #(#cfgs)*
124 #arg
125 ));
126 }
127
128 #[cfg(feature = "nightly")]
129 let mut task_outer_body = quote! {
130 trait _EmbassyInternalTaskTrait {
131 type Fut: ::core::future::Future + 'static;
132 fn construct(#fargs) -> Self::Fut;
133 }
134
135 impl _EmbassyInternalTaskTrait for () {
136 type Fut = impl core::future::Future + 'static;
137 fn construct(#fargs) -> Self::Fut {
138 #task_inner_ident(#(#full_args,)*)
139 }
140 }
141
142 const POOL_SIZE: usize = #pool_size;
143 static POOL: #embassy_executor::raw::TaskPool<<() as _EmbassyInternalTaskTrait>::Fut, POOL_SIZE> = #embassy_executor::raw::TaskPool::new();
144 unsafe { POOL._spawn_async_fn(move || <() as _EmbassyInternalTaskTrait>::construct(#(#full_args,)*)) }
145 };
146 #[cfg(not(feature = "nightly"))]
147 let mut task_outer_body = quote! {
148 const POOL_SIZE: usize = #pool_size;
149 static POOL: #embassy_executor::_export::TaskPoolRef = #embassy_executor::_export::TaskPoolRef::new();
150 unsafe { POOL.get::<_, POOL_SIZE>()._spawn_async_fn(move || #task_inner_ident(#(#full_args,)*)) }
151 };
152
153 let task_outer_attrs = task_inner.attrs.clone();
154
155 if !errors.is_empty() {
156 task_outer_body = quote! {
157 #![allow(unused_variables, unreachable_code)]
158 let _x: #embassy_executor::SpawnToken<()> = ::core::todo!();
159 _x
160 };
161 }
162
163 // Copy the generics + where clause to avoid more spurious errors.
164 let generics = &f.sig.generics;
165 let where_clause = &f.sig.generics.where_clause;
166
167 let result = quote! {
168 // This is the user's task function, renamed.
169 // We put it outside the #task_ident fn below, because otherwise
170 // the items defined there (such as POOL) would be in scope
171 // in the user's code.
172 #[doc(hidden)]
173 #task_inner
174
175 #(#task_outer_attrs)*
176 #visibility fn #task_ident #generics (#fargs) -> #embassy_executor::SpawnToken<impl Sized> #where_clause{
177 #task_outer_body
178 }
179
180 #errors
181 };
182
183 result
184}
185
186fn check_arg_ty(errors: &mut TokenStream, ty: &Type) {
187 struct Visitor<'a> {
188 errors: &'a mut TokenStream,
189 }
190
191 impl<'a, 'ast> Visit<'ast> for Visitor<'a> {
192 fn visit_type_reference(&mut self, i: &'ast syn::TypeReference) {
193 // only check for elided lifetime here. If not elided, it's checked by `visit_lifetime`.
194 if i.lifetime.is_none() {
195 error(
196 self.errors,
197 i.and_token,
198 "Arguments for tasks must live forever. Try using the `'static` lifetime.",
199 )
200 }
201 visit::visit_type_reference(self, i);
202 }
203
204 fn visit_lifetime(&mut self, i: &'ast syn::Lifetime) {
205 if i.ident.to_string() != "static" {
206 error(
207 self.errors,
208 i,
209 "Arguments for tasks must live forever. Try using the `'static` lifetime.",
210 )
211 }
212 }
213
214 fn visit_type_impl_trait(&mut self, i: &'ast syn::TypeImplTrait) {
215 error(self.errors, i, "`impl Trait` is not allowed in task arguments. It is syntax sugar for generics, and tasks can't be generic.");
216 }
217 }
218
219 Visit::visit_type(&mut Visitor { errors }, ty);
220}
221