From 7671605452658cedd6e460bb7c86ef450fe6a04c Mon Sep 17 00:00:00 2001
From: Ulf Lilleengen <ulf.lilleengen@gmail.com>
Date: Mon, 12 Apr 2021 10:41:21 +0200
Subject: [PATCH] Add std version of embassy::main

---
 embassy-macros/Cargo.toml |  1 +
 embassy-macros/src/lib.rs | 66 +++++++++++++++++++++++++++++++++++++++
 embassy-std/Cargo.toml    |  3 +-
 3 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/embassy-macros/Cargo.toml b/embassy-macros/Cargo.toml
index 20dc7c38..6da2a586 100644
--- a/embassy-macros/Cargo.toml
+++ b/embassy-macros/Cargo.toml
@@ -17,3 +17,4 @@ proc-macro = true
 stm32 = []
 nrf = []
 rp = []
+std = []
diff --git a/embassy-macros/src/lib.rs b/embassy-macros/src/lib.rs
index 885eac11..8f35cc9e 100644
--- a/embassy-macros/src/lib.rs
+++ b/embassy-macros/src/lib.rs
@@ -280,3 +280,69 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
     };
     result.into()
 }
+
+#[cfg(feature = "std")]
+#[proc_macro_attribute]
+pub fn main(_: TokenStream, item: TokenStream) -> TokenStream {
+    let task_fn = syn::parse_macro_input!(item as syn::ItemFn);
+
+    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.is_empty() {
+        task_fn
+            .sig
+            .span()
+            .unwrap()
+            .error("main function must not be generic")
+            .emit();
+        fail = true;
+    }
+
+    let args = task_fn.sig.inputs.clone();
+
+    if args.len() != 1 {
+        task_fn
+            .sig
+            .span()
+            .unwrap()
+            .error("main function must have one argument")
+            .emit();
+        fail = true;
+    }
+
+    if fail {
+        return TokenStream::new();
+    }
+
+    let task_fn_body = task_fn.block.clone();
+
+    let result = quote! {
+        #[embassy::task]
+        async fn __embassy_main(#args) {
+            #task_fn_body
+        }
+
+        fn main() -> ! {
+            unsafe fn make_static<T>(t: &mut T) -> &'static mut T {
+                ::core::mem::transmute(t)
+            }
+
+            let mut executor = ::embassy_std::Executor::new();
+            let executor = unsafe { make_static(&mut executor) };
+
+            executor.run(|spawner| {
+                spawner.spawn(__embassy_main(spawner)).unwrap();
+            })
+
+        }
+    };
+    result.into()
+}
diff --git a/embassy-std/Cargo.toml b/embassy-std/Cargo.toml
index 0a59999c..2a95838e 100644
--- a/embassy-std/Cargo.toml
+++ b/embassy-std/Cargo.toml
@@ -6,4 +6,5 @@ edition = "2018"
 
 [dependencies]
 embassy     = { version = "0.1.0", path = "../embassy", features = ["std"] }
-lazy_static = "1.4.0"
\ No newline at end of file
+embassy-macros = { version = "0.1.0", path = "../embassy-macros", features = ["std"]}
+lazy_static = "1.4.0"