1 files changed, 114 insertions, 0 deletions

diff --git a/embassy-macros/src/lib.rs b/embassy-macros/src/lib.rs
new file mode 100644
index 000000000..1745311ba
--- /dev/null
+++ b/embassy-macros/src/lib.rs
@@ -0,0 +1,114 @@
+#![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()
+}