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

1 files changed, 92 insertions, 3 deletions

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 _) }
+            })
+        }
+    }
+}