executor,sync: add support for turbo-wakers.

This is a `core` patch to make wakers 1 word (the task pointer) instead of 2 (task pointer + vtable). It allows having the "waker optimization" we had a while back on `WakerRegistration/AtomicWaker`, but EVERYWHERE, without patching all crates.

Advantages:
- Less memory usage.
- Faster.
- `AtomicWaker` can actually use atomics to load/store the waker, No critical section needed.
- No `dyn` call, which means `cargo-call-stack` can now see through wakes.

Disadvantages:
- You have to patch `core`...
- Breaks all executors and other things that create wakers, unless they opt in to using the new `from_ptr` API.

How to use:

- Run this shell script to patch `core`. https://gist.github.com/Dirbaio/c67da7cf318515181539122c9d32b395
- Enable `build-std`
- Enable `build-std-features = core/turbowakers`
- Enable feature `turbowakers` in `embassy-executor`, `embassy-sync`.
- Make sure you have no other crate creating wakers other than `embassy-executor`. These will panic at runtime.

Note that the patched `core` is equivalent to the unpached one when the `turbowakers` feature is not enabled, so it should be fine to leave it there.

8 files changed, 115 insertions, 42 deletions

diff --git a/embassy-executor/Cargo.toml b/embassy-executor/Cargo.toml
index c2868eb98..8ad3fd698 100644
--- a/embassy-executor/Cargo.toml
+++ b/embassy-executor/Cargo.toml
@@ -38,6 +38,8 @@ wasm = ["dep:wasm-bindgen", "dep:js-sys"]
 # Enable nightly-only features
 nightly = []
 
+turbowakers = []
+
 integrated-timers = ["dep:embassy-time"]
 
 # Trace interrupt invocations with rtos-trace.
diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 15ff18fc8..72c367c33 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -11,6 +11,7 @@ mod run_queue;
 #[cfg(feature = "integrated-timers")]
 mod timer_queue;
 pub(crate) mod util;
+#[cfg_attr(feature = "turbowakers", path = "waker_turbo.rs")]
 mod waker;
 
 use core::future::Future;
diff --git a/embassy-executor/src/raw/waker_turbo.rs b/embassy-executor/src/raw/waker_turbo.rs
new file mode 100644
index 000000000..435a0ff7e
--- /dev/null
+++ b/embassy-executor/src/raw/waker_turbo.rs
@@ -0,0 +1,34 @@
+use core::ptr::NonNull;
+use core::task::Waker;
+
+use super::{wake_task, TaskHeader, TaskRef};
+
+pub(crate) unsafe fn from_task(p: TaskRef) -> Waker {
+    Waker::from_turbo_ptr(NonNull::new_unchecked(p.as_ptr() as _))
+}
+
+/// Get a task pointer from a waker.
+///
+/// This can be used as an optimization in wait queues to store task pointers
+/// (1 word) instead of full Wakers (2 words). This saves a bit of RAM and helps
+/// avoid dynamic dispatch.
+///
+/// You can use the returned task pointer to wake the task with [`wake_task`](super::wake_task).
+///
+/// # Panics
+///
+/// Panics if the waker is not created by the Embassy executor.
+pub fn task_from_waker(waker: &Waker) -> TaskRef {
+    let ptr = waker.as_turbo_ptr().as_ptr();
+
+    // safety: our wakers are always created with `TaskRef::as_ptr`
+    unsafe { TaskRef::from_ptr(ptr as *const TaskHeader) }
+}
+
+#[inline(never)]
+#[no_mangle]
+fn _turbo_wake(ptr: NonNull<()>) {
+    // safety: our wakers are always created with `TaskRef::as_ptr`
+    let task = unsafe { TaskRef::from_ptr(ptr.as_ptr() as *const TaskHeader) };
+    wake_task(task)
+}
diff --git a/embassy-sync/Cargo.toml b/embassy-sync/Cargo.toml
index 7b5d3ce48..e4871e718 100644
--- a/embassy-sync/Cargo.toml
+++ b/embassy-sync/Cargo.toml
@@ -25,6 +25,7 @@ features = ["nightly"]
 [features]
 nightly = ["embedded-io/async"]
 std = []
+turbowakers = []
 
 [dependencies]
 defmt = { version = "0.3", optional = true }
diff --git a/embassy-sync/src/waitqueue/atomic_waker.rs b/embassy-sync/src/waitqueue/atomic_waker.rs
new file mode 100644
index 000000000..63fe04a6e
--- /dev/null
+++ b/embassy-sync/src/waitqueue/atomic_waker.rs
@@ -0,0 +1,41 @@
+use core::cell::Cell;
+use core::task::Waker;
+
+use crate::blocking_mutex::raw::CriticalSectionRawMutex;
+use crate::blocking_mutex::Mutex;
+
+/// Utility struct to register and wake a waker.
+pub struct AtomicWaker {
+    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
+}
+
+impl AtomicWaker {
+    /// Create a new `AtomicWaker`.
+    pub const fn new() -> Self {
+        Self {
+            waker: Mutex::const_new(CriticalSectionRawMutex::new(), Cell::new(None)),
+        }
+    }
+
+    /// Register a waker. Overwrites the previous waker, if any.
+    pub fn register(&self, w: &Waker) {
+        critical_section::with(|cs| {
+            let cell = self.waker.borrow(cs);
+            cell.set(match cell.replace(None) {
+                Some(w2) if (w2.will_wake(w)) => Some(w2),
+                _ => Some(w.clone()),
+            })
+        })
+    }
+
+    /// Wake the registered waker, if any.
+    pub fn wake(&self) {
+        critical_section::with(|cs| {
+            let cell = self.waker.borrow(cs);
+            if let Some(w) = cell.replace(None) {
+                w.wake_by_ref();
+                cell.set(Some(w));
+            }
+        })
+    }
+}
diff --git a/embassy-sync/src/waitqueue/atomic_waker_turbo.rs b/embassy-sync/src/waitqueue/atomic_waker_turbo.rs
new file mode 100644
index 000000000..5c6a96ec8
--- /dev/null
+++ b/embassy-sync/src/waitqueue/atomic_waker_turbo.rs
@@ -0,0 +1,30 @@
+use core::ptr;
+use core::ptr::NonNull;
+use core::sync::atomic::{AtomicPtr, Ordering};
+use core::task::Waker;
+
+/// Utility struct to register and wake a waker.
+pub struct AtomicWaker {
+    waker: AtomicPtr<()>,
+}
+
+impl AtomicWaker {
+    /// Create a new `AtomicWaker`.
+    pub const fn new() -> Self {
+        Self {
+            waker: AtomicPtr::new(ptr::null_mut()),
+        }
+    }
+
+    /// Register a waker. Overwrites the previous waker, if any.
+    pub fn register(&self, w: &Waker) {
+        self.waker.store(w.as_turbo_ptr().as_ptr() as _, Ordering::Release);
+    }
+
+    /// Wake the registered waker, if any.
+    pub fn wake(&self) {
+        if let Some(ptr) = NonNull::new(self.waker.load(Ordering::Acquire)) {
+            unsafe { Waker::from_turbo_ptr(ptr) }.wake();
+        }
+    }
+}
diff --git a/embassy-sync/src/waitqueue/mod.rs b/embassy-sync/src/waitqueue/mod.rs
index 6661a6b61..6b0b0c64e 100644
--- a/embassy-sync/src/waitqueue/mod.rs
+++ b/embassy-sync/src/waitqueue/mod.rs
@@ -1,7 +1,11 @@
 //! Async low-level wait queues
 
-mod waker;
-pub use waker::*;
+#[cfg_attr(feature = "turbowakers", path = "atomic_waker_turbo.rs")]
+mod atomic_waker;
+pub use atomic_waker::*;
+
+mod waker_registration;
+pub use waker_registration::*;
 
 mod multi_waker;
 pub use multi_waker::*;
diff --git a/embassy-sync/src/waitqueue/waker.rs b/embassy-sync/src/waitqueue/waker_registration.rs
index 9ce94a089..9b666e7c4 100644
--- a/embassy-sync/src/waitqueue/waker.rs
+++ b/embassy-sync/src/waitqueue/waker_registration.rs
@@ -1,10 +1,6 @@
-use core::cell::Cell;
 use core::mem;
 use core::task::Waker;
 
-use crate::blocking_mutex::raw::CriticalSectionRawMutex;
-use crate::blocking_mutex::Mutex;
-
 /// Utility struct to register and wake a waker.
 #[derive(Debug, Default)]
 pub struct WakerRegistration {
@@ -54,39 +50,3 @@ impl WakerRegistration {
         self.waker.is_some()
     }
 }
-
-/// Utility struct to register and wake a waker.
-pub struct AtomicWaker {
-    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
-}
-
-impl AtomicWaker {
-    /// Create a new `AtomicWaker`.
-    pub const fn new() -> Self {
-        Self {
-            waker: Mutex::const_new(CriticalSectionRawMutex::new(), Cell::new(None)),
-        }
-    }
-
-    /// Register a waker. Overwrites the previous waker, if any.
-    pub fn register(&self, w: &Waker) {
-        critical_section::with(|cs| {
-            let cell = self.waker.borrow(cs);
-            cell.set(match cell.replace(None) {
-                Some(w2) if (w2.will_wake(w)) => Some(w2),
-                _ => Some(w.clone()),
-            })
-        })
-    }
-
-    /// Wake the registered waker, if any.
-    pub fn wake(&self) {
-        critical_section::with(|cs| {
-            let cell = self.waker.borrow(cs);
-            if let Some(w) = cell.replace(None) {
-                w.wake_by_ref();
-                cell.set(Some(w));
-            }
-        })
-    }
-}