executor: add support for main/task macros in stable (allocates tasks in an arena)

10 files changed, 105 insertions, 10 deletions

diff --git a/.github/ci/test.sh b/.github/ci/test.sh
index e48d6e65b..dfc2b08ce 100755
--- a/.github/ci/test.sh
+++ b/.github/ci/test.sh
@@ -15,6 +15,7 @@ export CARGO_NET_GIT_FETCH_WITH_CLI=true
 hashtime restore /ci/cache/filetime.json || true
 hashtime save /ci/cache/filetime.json
 
+MIRIFLAGS=-Zmiri-ignore-leaks cargo miri test --manifest-path ./embassy-executor/Cargo.toml
 MIRIFLAGS=-Zmiri-ignore-leaks cargo miri test --manifest-path ./embassy-executor/Cargo.toml --features nightly
 
 cargo test --manifest-path ./embassy-sync/Cargo.toml 
diff --git a/embassy-executor/Cargo.toml b/embassy-executor/Cargo.toml
index f0e85a2c8..3ff8440ca 100644
--- a/embassy-executor/Cargo.toml
+++ b/embassy-executor/Cargo.toml
@@ -43,7 +43,7 @@ executor-thread = []
 executor-interrupt = []
 
 # Enable nightly-only features
-nightly = []
+nightly = ["embassy-macros/nightly"]
 
 turbowakers = []
 
diff --git a/embassy-executor/src/arch/cortex_m.rs b/embassy-executor/src/arch/cortex_m.rs
index 55299c94f..4a6d58575 100644
--- a/embassy-executor/src/arch/cortex_m.rs
+++ b/embassy-executor/src/arch/cortex_m.rs
@@ -51,7 +51,6 @@ mod thread {
     use core::arch::asm;
     use core::marker::PhantomData;
 
-    #[cfg(feature = "nightly")]
     pub use embassy_macros::main_cortex_m as main;
 
     use crate::{raw, Spawner};
diff --git a/embassy-executor/src/arch/riscv32.rs b/embassy-executor/src/arch/riscv32.rs
index 6814e7844..ca12c3403 100644
--- a/embassy-executor/src/arch/riscv32.rs
+++ b/embassy-executor/src/arch/riscv32.rs
@@ -7,7 +7,6 @@ pub use thread::*;
 mod thread {
     use core::marker::PhantomData;
 
-    #[cfg(feature = "nightly")]
     pub use embassy_macros::main_riscv as main;
     use portable_atomic::{AtomicBool, Ordering};
 
diff --git a/embassy-executor/src/arch/std.rs b/embassy-executor/src/arch/std.rs
index 5b2f7e2e4..598bb0509 100644
--- a/embassy-executor/src/arch/std.rs
+++ b/embassy-executor/src/arch/std.rs
@@ -8,7 +8,6 @@ mod thread {
     use std::marker::PhantomData;
     use std::sync::{Condvar, Mutex};
 
-    #[cfg(feature = "nightly")]
     pub use embassy_macros::main_std as main;
 
     use crate::{raw, Spawner};
diff --git a/embassy-executor/src/arch/wasm.rs b/embassy-executor/src/arch/wasm.rs
index 15aed867a..3faa92575 100644
--- a/embassy-executor/src/arch/wasm.rs
+++ b/embassy-executor/src/arch/wasm.rs
@@ -8,7 +8,6 @@ mod thread {
 
     use core::marker::PhantomData;
 
-    #[cfg(feature = "nightly")]
     pub use embassy_macros::main_wasm as main;
     use js_sys::Promise;
     use wasm_bindgen::prelude::*;
diff --git a/embassy-executor/src/arch/xtensa.rs b/embassy-executor/src/arch/xtensa.rs
index d335594e2..6ed9f9e72 100644
--- a/embassy-executor/src/arch/xtensa.rs
+++ b/embassy-executor/src/arch/xtensa.rs
@@ -8,7 +8,6 @@ mod thread {
     use core::marker::PhantomData;
     use core::sync::atomic::{AtomicBool, Ordering};
 
-    #[cfg(feature = "nightly")]
     pub use embassy_macros::main_riscv as main;
 
     use crate::{raw, Spawner};
diff --git a/embassy-executor/src/lib.rs b/embassy-executor/src/lib.rs
index 3b61b4ba5..ac7dbb035 100644
--- a/embassy-executor/src/lib.rs
+++ b/embassy-executor/src/lib.rs
@@ -1,5 +1,5 @@
 #![cfg_attr(not(any(feature = "arch-std", feature = "arch-wasm")), no_std)]
-#![cfg_attr(all(feature = "nightly", feature = "arch-xtensa"), feature(asm_experimental_arch))]
+#![cfg_attr(feature = "arch-xtensa", feature(asm_experimental_arch))]
 #![allow(clippy::new_without_default)]
 #![doc = include_str!("../README.md")]
 #![warn(missing_docs)]
@@ -7,7 +7,6 @@
 // This mod MUST go first, so that the others see its macros.
 pub(crate) mod fmt;
 
-#[cfg(feature = "nightly")]
 pub use embassy_macros::task;
 
 macro_rules! check_at_most_one {
@@ -44,4 +43,94 @@ pub use spawner::*;
 /// Implementation details for embassy macros.
 /// Do not use. Used for macros and HALs only. Not covered by semver guarantees.
 #[doc(hidden)]
-pub mod _export {}
+#[cfg(not(feature = "nightly"))]
+pub mod _export {
+    use core::alloc::Layout;
+    use core::cell::{Cell, UnsafeCell};
+    use core::future::Future;
+    use core::mem::MaybeUninit;
+    use core::ptr::null_mut;
+
+    use critical_section::{CriticalSection, Mutex};
+
+    use crate::raw::TaskPool;
+
+    struct Arena<const N: usize> {
+        buf: UnsafeCell<MaybeUninit<[u8; N]>>,
+        ptr: Mutex<Cell<*mut u8>>,
+    }
+
+    unsafe impl<const N: usize> Sync for Arena<N> {}
+    unsafe impl<const N: usize> Send for Arena<N> {}
+
+    impl<const N: usize> Arena<N> {
+        const fn new() -> Self {
+            Self {
+                buf: UnsafeCell::new(MaybeUninit::uninit()),
+                ptr: Mutex::new(Cell::new(null_mut())),
+            }
+        }
+
+        fn alloc<T>(&'static self, cs: CriticalSection) -> &'static mut MaybeUninit<T> {
+            let layout = Layout::new::<T>();
+
+            let start = self.buf.get().cast::<u8>();
+            let end = unsafe { start.add(N) };
+
+            let mut ptr = self.ptr.borrow(cs).get();
+            if ptr.is_null() {
+                ptr = self.buf.get().cast::<u8>();
+            }
+
+            let bytes_left = (end as usize) - (ptr as usize);
+            let align_offset = (ptr as usize).next_multiple_of(layout.align()) - (ptr as usize);
+
+            if align_offset + layout.size() > bytes_left {
+                panic!("arena full");
+            }
+
+            let res = unsafe { ptr.add(align_offset) };
+            let ptr = unsafe { ptr.add(align_offset + layout.size()) };
+
+            self.ptr.borrow(cs).set(ptr);
+
+            unsafe { &mut *(res as *mut MaybeUninit<T>) }
+        }
+    }
+
+    const ARENA_SIZE: usize = 16 * 1024;
+    static ARENA: Arena<ARENA_SIZE> = Arena::new();
+
+    pub struct TaskPoolRef {
+        // type-erased `&'static mut TaskPool<F, N>`
+        // Needed because statics can't have generics.
+        ptr: Mutex<Cell<*mut ()>>,
+    }
+    unsafe impl Sync for TaskPoolRef {}
+    unsafe impl Send for TaskPoolRef {}
+
+    impl TaskPoolRef {
+        pub const fn new() -> Self {
+            Self {
+                ptr: Mutex::new(Cell::new(null_mut())),
+            }
+        }
+
+        /// Get the pool for this ref, allocating it from the arena the first time.
+        ///
+        /// safety: for a given TaskPoolRef instance, must always call with the exact
+        /// same generic params.
+        pub unsafe fn get<F: Future, const N: usize>(&'static self) -> &'static TaskPool<F, N> {
+            critical_section::with(|cs| {
+                let ptr = self.ptr.borrow(cs);
+                if ptr.get().is_null() {
+                    let pool = ARENA.alloc::<TaskPool<F, N>>(cs);
+                    pool.write(TaskPool::new());
+                    ptr.set(pool as *mut _ as _);
+                }
+
+                unsafe { &*(ptr.get() as *const _) }
+            })
+        }
+    }
+}
diff --git a/embassy-macros/Cargo.toml b/embassy-macros/Cargo.toml
index 4e94bf834..e13104cfd 100644
--- a/embassy-macros/Cargo.toml
+++ b/embassy-macros/Cargo.toml
@@ -20,4 +20,5 @@ proc-macro2 = "1.0.29"
 [lib]
 proc-macro = true
 
-[features]
\ No newline at end of file
+[features]
+nightly = []
\ No newline at end of file
diff --git a/embassy-macros/src/macros/task.rs b/embassy-macros/src/macros/task.rs
index 1d30434e9..5161e1020 100644
--- a/embassy-macros/src/macros/task.rs
+++ b/embassy-macros/src/macros/task.rs
@@ -79,6 +79,7 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
     task_inner.vis = syn::Visibility::Inherited;
     task_inner.sig.ident = task_inner_ident.clone();
 
+    #[cfg(feature = "nightly")]
     let mut task_outer: ItemFn = parse_quote! {
         #visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
             type Fut = impl ::core::future::Future + 'static;
@@ -87,6 +88,14 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
             unsafe { POOL._spawn_async_fn(move || #task_inner_ident(#(#arg_names,)*)) }
         }
     };
+    #[cfg(not(feature = "nightly"))]
+    let mut task_outer: ItemFn = parse_quote! {
+        #visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
+            const POOL_SIZE: usize = #pool_size;
+            static POOL: ::embassy_executor::_export::TaskPoolRef = ::embassy_executor::_export::TaskPoolRef::new();
+            unsafe { POOL.get::<_, POOL_SIZE>()._spawn_async_fn(move || #task_inner_ident(#(#arg_names,)*)) }
+        }
+    };
 
     task_outer.attrs.append(&mut task_inner.attrs.clone());