Rebase on master

9 files changed, 192 insertions, 153 deletions

diff --git a/embassy-executor/src/arch/riscv32.rs b/embassy-executor/src/arch/riscv32.rs
index 2a4b006da..e97a56cda 100644
--- a/embassy-executor/src/arch/riscv32.rs
+++ b/embassy-executor/src/arch/riscv32.rs
@@ -1,7 +1,6 @@
 use core::marker::PhantomData;
 use core::ptr;
-
-use atomic_polyfill::{AtomicBool, Ordering};
+use core::sync::atomic::{AtomicBool, Ordering};
 
 use super::{raw, Spawner};
 
diff --git a/embassy-executor/src/arch/xtensa.rs b/embassy-executor/src/arch/xtensa.rs
index f908aaa70..4ee0d9f78 100644
--- a/embassy-executor/src/arch/xtensa.rs
+++ b/embassy-executor/src/arch/xtensa.rs
@@ -1,7 +1,6 @@
 use core::marker::PhantomData;
 use core::ptr;
-
-use atomic_polyfill::{AtomicBool, Ordering};
+use core::sync::atomic::{AtomicBool, Ordering};
 
 use super::{raw, Spawner};
 
diff --git a/embassy-executor/src/lib.rs b/embassy-executor/src/lib.rs
index e4cbd04b9..4c7e2f4cd 100644
--- a/embassy-executor/src/lib.rs
+++ b/embassy-executor/src/lib.rs
@@ -8,18 +8,22 @@
 pub(crate) mod fmt;
 
 #[cfg(feature = "nightly")]
-pub use embassy_macros::{main, task};
+pub use embassy_macros::task;
 
 cfg_if::cfg_if! {
     if #[cfg(cortex_m)] {
         #[path="arch/cortex_m.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_cortex_m as main;
     }
     else if #[cfg(target_arch="riscv32")] {
         #[path="arch/riscv32.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_riscv as main;
     }
     else if #[cfg(all(target_arch="xtensa", feature = "nightly"))] {
         #[path="arch/xtensa.rs"]
@@ -30,11 +34,15 @@ cfg_if::cfg_if! {
         #[path="arch/wasm.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_wasm as main;
     }
     else if #[cfg(feature="std")] {
         #[path="arch/std.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_std as main;
     }
 }
 
diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index e1258ebb5..42bd82262 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -15,10 +15,10 @@ mod waker;
 
 use core::cell::Cell;
 use core::future::Future;
+use core::mem;
 use core::pin::Pin;
 use core::ptr::NonNull;
 use core::task::{Context, Poll};
-use core::{mem, ptr};
 
 use atomic_polyfill::{AtomicU32, Ordering};
 use critical_section::CriticalSection;
@@ -43,14 +43,11 @@ pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 1;
 pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 2;
 
 /// Raw task header for use in task pointers.
-///
-/// This is an opaque struct, used for raw pointers to tasks, for use
-/// with funtions like [`wake_task`] and [`task_from_waker`].
-pub struct TaskHeader {
+pub(crate) struct TaskHeader {
     pub(crate) state: AtomicU32,
     pub(crate) run_queue_item: RunQueueItem,
-    pub(crate) executor: Cell<*const Executor>, // Valid if state != 0
-    pub(crate) poll_fn: UninitCell<unsafe fn(NonNull<TaskHeader>)>, // Valid if STATE_SPAWNED
+    pub(crate) executor: Cell<Option<&'static Executor>>,
+    poll_fn: Cell<Option<unsafe fn(TaskRef)>>,
 
     #[cfg(feature = "integrated-timers")]
     pub(crate) expires_at: Cell<Instant>,
@@ -58,20 +55,34 @@ pub struct TaskHeader {
     pub(crate) timer_queue_item: timer_queue::TimerQueueItem,
 }
 
-impl TaskHeader {
-    pub(crate) const fn new() -> Self {
+/// This is essentially a `&'static TaskStorage<F>` where the type of the future has been erased.
+#[derive(Clone, Copy)]
+pub struct TaskRef {
+    ptr: NonNull<TaskHeader>,
+}
+
+impl TaskRef {
+    fn new<F: Future + 'static>(task: &'static TaskStorage<F>) -> Self {
         Self {
-            state: AtomicU32::new(0),
-            run_queue_item: RunQueueItem::new(),
-            executor: Cell::new(ptr::null()),
-            poll_fn: UninitCell::uninit(),
+            ptr: NonNull::from(task).cast(),
+        }
+    }
 
-            #[cfg(feature = "integrated-timers")]
-            expires_at: Cell::new(Instant::from_ticks(0)),
-            #[cfg(feature = "integrated-timers")]
-            timer_queue_item: timer_queue::TimerQueueItem::new(),
+    /// Safety: The pointer must have been obtained with `Task::as_ptr`
+    pub(crate) unsafe fn from_ptr(ptr: *const TaskHeader) -> Self {
+        Self {
+            ptr: NonNull::new_unchecked(ptr as *mut TaskHeader),
         }
     }
+
+    pub(crate) fn header(self) -> &'static TaskHeader {
+        unsafe { self.ptr.as_ref() }
+    }
+
+    /// The returned pointer is valid for the entire TaskStorage.
+    pub(crate) fn as_ptr(self) -> *const TaskHeader {
+        self.ptr.as_ptr()
+    }
 }
 
 /// Raw storage in which a task can be spawned.
@@ -101,7 +112,18 @@ impl<F: Future + 'static> TaskStorage<F> {
     /// Create a new TaskStorage, in not-spawned state.
     pub const fn new() -> Self {
         Self {
-            raw: TaskHeader::new(),
+            raw: TaskHeader {
+                state: AtomicU32::new(0),
+                run_queue_item: RunQueueItem::new(),
+                executor: Cell::new(None),
+                // Note: this is lazily initialized so that a static `TaskStorage` will go in `.bss`
+                poll_fn: Cell::new(None),
+
+                #[cfg(feature = "integrated-timers")]
+                expires_at: Cell::new(Instant::from_ticks(0)),
+                #[cfg(feature = "integrated-timers")]
+                timer_queue_item: timer_queue::TimerQueueItem::new(),
+            },
             future: UninitCell::uninit(),
         }
     }
@@ -120,29 +142,17 @@ impl<F: Future + 'static> TaskStorage<F> {
     /// Once the task has finished running, you may spawn it again. It is allowed to spawn it
     /// on a different executor.
     pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
-        if self.spawn_mark_used() {
-            return unsafe { SpawnToken::<F>::new(self.spawn_initialize(future)) };
+        let task = AvailableTask::claim(self);
+        match task {
+            Some(task) => {
+                let task = task.initialize(future);
+                unsafe { SpawnToken::<F>::new(task) }
+            }
+            None => SpawnToken::new_failed(),
         }
-
-        SpawnToken::<F>::new_failed()
-    }
-
-    fn spawn_mark_used(&'static self) -> bool {
-        let state = STATE_SPAWNED | STATE_RUN_QUEUED;
-        self.raw
-            .state
-            .compare_exchange(0, state, Ordering::AcqRel, Ordering::Acquire)
-            .is_ok()
-    }
-
-    unsafe fn spawn_initialize(&'static self, future: impl FnOnce() -> F) -> NonNull<TaskHeader> {
-        // Initialize the task
-        self.raw.poll_fn.write(Self::poll);
-        self.future.write(future());
-        NonNull::new_unchecked(self as *const TaskStorage<F> as *const TaskHeader as *mut TaskHeader)
     }
 
-    unsafe fn poll(p: NonNull<TaskHeader>) {
+    unsafe fn poll(p: TaskRef) {
         let this = &*(p.as_ptr() as *const TaskStorage<F>);
 
         let future = Pin::new_unchecked(this.future.as_mut());
@@ -164,6 +174,28 @@ impl<F: Future + 'static> TaskStorage<F> {
 
 unsafe impl<F: Future + 'static> Sync for TaskStorage<F> {}
 
+struct AvailableTask<F: Future + 'static> {
+    task: &'static TaskStorage<F>,
+}
+
+impl<F: Future + 'static> AvailableTask<F> {
+    fn claim(task: &'static TaskStorage<F>) -> Option<Self> {
+        task.raw
+            .state
+            .compare_exchange(0, STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::AcqRel, Ordering::Acquire)
+            .ok()
+            .map(|_| Self { task })
+    }
+
+    fn initialize(self, future: impl FnOnce() -> F) -> TaskRef {
+        unsafe {
+            self.task.raw.poll_fn.set(Some(TaskStorage::<F>::poll));
+            self.task.future.write(future());
+        }
+        TaskRef::new(self.task)
+    }
+}
+
 /// Raw storage that can hold up to N tasks of the same type.
 ///
 /// This is essentially a `[TaskStorage<F>; N]`.
@@ -187,13 +219,14 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
     /// is currently free. If none is free, a "poisoned" SpawnToken is returned,
     /// which will cause [`Spawner::spawn()`](super::Spawner::spawn) to return the error.
     pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
-        for task in &self.pool {
-            if task.spawn_mark_used() {
-                return unsafe { SpawnToken::<F>::new(task.spawn_initialize(future)) };
+        let task = self.pool.iter().find_map(AvailableTask::claim);
+        match task {
+            Some(task) => {
+                let task = task.initialize(future);
+                unsafe { SpawnToken::<F>::new(task) }
             }
+            None => SpawnToken::new_failed(),
         }
-
-        SpawnToken::<F>::new_failed()
     }
 
     /// Like spawn(), but allows the task to be send-spawned if the args are Send even if
@@ -235,13 +268,14 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
         // This ONLY holds for `async fn` futures. The other `spawn` methods can be called directly
         // by the user, with arbitrary hand-implemented futures. This is why these return `SpawnToken<F>`.
 
-        for task in &self.pool {
-            if task.spawn_mark_used() {
-                return SpawnToken::<FutFn>::new(task.spawn_initialize(future));
+        let task = self.pool.iter().find_map(AvailableTask::claim);
+        match task {
+            Some(task) => {
+                let task = task.initialize(future);
+                unsafe { SpawnToken::<FutFn>::new(task) }
             }
+            None => SpawnToken::new_failed(),
         }
-
-        SpawnToken::<FutFn>::new_failed()
     }
 }
 
@@ -307,7 +341,7 @@ impl Executor {
     /// - `task` must be set up to run in this executor.
     /// - `task` must NOT be already enqueued (in this executor or another one).
     #[inline(always)]
-    unsafe fn enqueue(&self, cs: CriticalSection, task: NonNull<TaskHeader>) {
+    unsafe fn enqueue(&self, cs: CriticalSection, task: TaskRef) {
         #[cfg(feature = "rtos-trace")]
         trace::task_ready_begin(task.as_ptr() as u32);
 
@@ -325,8 +359,8 @@ impl Executor {
     /// It is OK to use `unsafe` to call this from a thread that's not the executor thread.
     /// In this case, the task's Future must be Send. This is because this is effectively
     /// sending the task to the executor thread.
-    pub(super) unsafe fn spawn(&'static self, task: NonNull<TaskHeader>) {
-        task.as_ref().executor.set(self);
+    pub(super) unsafe fn spawn(&'static self, task: TaskRef) {
+        task.header().executor.set(Some(self));
 
         #[cfg(feature = "rtos-trace")]
         trace::task_new(task.as_ptr() as u32);
@@ -354,46 +388,54 @@ impl Executor {
     /// somehow schedule for `poll()` to be called later, at a time you know for sure there's
     /// no `poll()` already running.
     pub unsafe fn poll(&'static self) {
-        #[cfg(feature = "integrated-timers")]
-        self.timer_queue.dequeue_expired(Instant::now(), |task| wake_task(task));
+        loop {
+            #[cfg(feature = "integrated-timers")]
+            self.timer_queue.dequeue_expired(Instant::now(), |task| wake_task(task));
 
-        self.run_queue.dequeue_all(|p| {
-            let task = p.as_ref();
+            self.run_queue.dequeue_all(|p| {
+                let task = p.header();
 
-            #[cfg(feature = "integrated-timers")]
-            task.expires_at.set(Instant::MAX);
-
-            let state = task.state.fetch_and(!STATE_RUN_QUEUED, Ordering::AcqRel);
-            if state & STATE_SPAWNED == 0 {
-                // If task is not running, ignore it. This can happen in the following scenario:
-                //   - Task gets dequeued, poll starts
-                //   - While task is being polled, it gets woken. It gets placed in the queue.
-                //   - Task poll finishes, returning done=true
-                //   - RUNNING bit is cleared, but the task is already in the queue.
-                return;
-            }
+                #[cfg(feature = "integrated-timers")]
+                task.expires_at.set(Instant::MAX);
+
+                let state = task.state.fetch_and(!STATE_RUN_QUEUED, Ordering::AcqRel);
+                if state & STATE_SPAWNED == 0 {
+                    // If task is not running, ignore it. This can happen in the following scenario:
+                    //   - Task gets dequeued, poll starts
+                    //   - While task is being polled, it gets woken. It gets placed in the queue.
+                    //   - Task poll finishes, returning done=true
+                    //   - RUNNING bit is cleared, but the task is already in the queue.
+                    return;
+                }
+
+                #[cfg(feature = "rtos-trace")]
+                trace::task_exec_begin(p.as_ptr() as u32);
 
-            #[cfg(feature = "rtos-trace")]
-            trace::task_exec_begin(p.as_ptr() as u32);
+                // Run the task
+                task.poll_fn.get().unwrap_unchecked()(p);
 
-            // Run the task
-            task.poll_fn.read()(p as _);
+                #[cfg(feature = "rtos-trace")]
+                trace::task_exec_end();
 
-            #[cfg(feature = "rtos-trace")]
-            trace::task_exec_end();
+                // Enqueue or update into timer_queue
+                #[cfg(feature = "integrated-timers")]
+                self.timer_queue.update(p);
+            });
 
-            // Enqueue or update into timer_queue
             #[cfg(feature = "integrated-timers")]
-            self.timer_queue.update(p);
-        });
+            {
+                // If this is already in the past, set_alarm might return false
+                // In that case do another poll loop iteration.
+                let next_expiration = self.timer_queue.next_expiration();
+                if driver::set_alarm(self.alarm, next_expiration.as_ticks()) {
+                    break;
+                }
+            }
 
-        #[cfg(feature = "integrated-timers")]
-        {
-            // If this is already in the past, set_alarm will immediately trigger the alarm.
-            // This will cause `signal_fn` to be called, which will cause `poll()` to be called again,
-            // so we immediately do another poll loop iteration.
-            let next_expiration = self.timer_queue.next_expiration();
-            driver::set_alarm(self.alarm, next_expiration.as_ticks());
+            #[cfg(not(feature = "integrated-timers"))]
+            {
+                break;
+            }
         }
 
         #[cfg(feature = "rtos-trace")]
@@ -409,16 +451,12 @@ impl Executor {
     }
 }
 
-/// Wake a task by raw pointer.
+/// Wake a task by `TaskRef`.
 ///
-/// You can obtain task pointers from `Waker`s using [`task_from_waker`].
-///
-/// # Safety
-///
-/// `task` must be a valid task pointer obtained from [`task_from_waker`].
-pub unsafe fn wake_task(task: NonNull<TaskHeader>) {
+/// You can obtain a `TaskRef` from a `Waker` using [`task_from_waker`].
+pub fn wake_task(task: TaskRef) {
     critical_section::with(|cs| {
-        let header = task.as_ref();
+        let header = task.header();
         let state = header.state.load(Ordering::Relaxed);
 
         // If already scheduled, or if not started,
@@ -430,20 +468,29 @@ pub unsafe fn wake_task(task: NonNull<TaskHeader>) {
         header.state.store(state | STATE_RUN_QUEUED, Ordering::Relaxed);
 
         // We have just marked the task as scheduled, so enqueue it.
-        let executor = &*header.executor.get();
-        executor.enqueue(cs, task);
+        unsafe {
+            let executor = header.executor.get().unwrap_unchecked();
+            executor.enqueue(cs, task);
+        }
     })
 }
 
 #[cfg(feature = "integrated-timers")]
-#[no_mangle]
-unsafe fn _embassy_time_schedule_wake(at: Instant, waker: &core::task::Waker) {
-    let task = waker::task_from_waker(waker);
-    let task = task.as_ref();
-    let expires_at = task.expires_at.get();
-    task.expires_at.set(expires_at.min(at));
+struct TimerQueue;
+
+#[cfg(feature = "integrated-timers")]
+impl embassy_time::queue::TimerQueue for TimerQueue {
+    fn schedule_wake(&'static self, at: Instant, waker: &core::task::Waker) {
+        let task = waker::task_from_waker(waker);
+        let task = task.header();
+        let expires_at = task.expires_at.get();
+        task.expires_at.set(expires_at.min(at));
+    }
 }
 
+#[cfg(feature = "integrated-timers")]
+embassy_time::timer_queue_impl!(static TIMER_QUEUE: TimerQueue = TimerQueue);
+
 #[cfg(feature = "rtos-trace")]
 impl rtos_trace::RtosTraceOSCallbacks for Executor {
     fn task_list() {
diff --git a/embassy-executor/src/raw/run_queue.rs b/embassy-executor/src/raw/run_queue.rs
index ed8c82a5c..362157535 100644
--- a/embassy-executor/src/raw/run_queue.rs
+++ b/embassy-executor/src/raw/run_queue.rs
@@ -4,7 +4,7 @@ use core::ptr::NonNull;
 use atomic_polyfill::{AtomicPtr, Ordering};
 use critical_section::CriticalSection;
 
-use super::TaskHeader;
+use super::{TaskHeader, TaskRef};
 
 pub(crate) struct RunQueueItem {
     next: AtomicPtr<TaskHeader>,
@@ -46,25 +46,26 @@ impl RunQueue {
     ///
     /// `item` must NOT be already enqueued in any queue.
     #[inline(always)]
-    pub(crate) unsafe fn enqueue(&self, _cs: CriticalSection, task: NonNull<TaskHeader>) -> bool {
+    pub(crate) unsafe fn enqueue(&self, _cs: CriticalSection, task: TaskRef) -> bool {
         let prev = self.head.load(Ordering::Relaxed);
-        task.as_ref().run_queue_item.next.store(prev, Ordering::Relaxed);
-        self.head.store(task.as_ptr(), Ordering::Relaxed);
+        task.header().run_queue_item.next.store(prev, Ordering::Relaxed);
+        self.head.store(task.as_ptr() as _, Ordering::Relaxed);
         prev.is_null()
     }
 
     /// Empty the queue, then call `on_task` for each task that was in the queue.
     /// NOTE: It is OK for `on_task` to enqueue more tasks. In this case they're left in the queue
     /// and will be processed by the *next* call to `dequeue_all`, *not* the current one.
-    pub(crate) fn dequeue_all(&self, on_task: impl Fn(NonNull<TaskHeader>)) {
+    pub(crate) fn dequeue_all(&self, on_task: impl Fn(TaskRef)) {
         // Atomically empty the queue.
         let mut ptr = self.head.swap(ptr::null_mut(), Ordering::AcqRel);
 
         // Iterate the linked list of tasks that were previously in the queue.
         while let Some(task) = NonNull::new(ptr) {
+            let task = unsafe { TaskRef::from_ptr(task.as_ptr()) };
             // If the task re-enqueues itself, the `next` pointer will get overwritten.
             // Therefore, first read the next pointer, and only then process the task.
-            let next = unsafe { task.as_ref() }.run_queue_item.next.load(Ordering::Relaxed);
+            let next = task.header().run_queue_item.next.load(Ordering::Relaxed);
 
             on_task(task);
 
diff --git a/embassy-executor/src/raw/timer_queue.rs b/embassy-executor/src/raw/timer_queue.rs
index 24c31892a..57d6d3cda 100644
--- a/embassy-executor/src/raw/timer_queue.rs
+++ b/embassy-executor/src/raw/timer_queue.rs
@@ -1,45 +1,39 @@
 use core::cell::Cell;
 use core::cmp::min;
-use core::ptr;
-use core::ptr::NonNull;
 
 use atomic_polyfill::Ordering;
 use embassy_time::Instant;
 
-use super::{TaskHeader, STATE_TIMER_QUEUED};
+use super::{TaskRef, STATE_TIMER_QUEUED};
 
 pub(crate) struct TimerQueueItem {
-    next: Cell<*mut TaskHeader>,
+    next: Cell<Option<TaskRef>>,
 }
 
 impl TimerQueueItem {
     pub const fn new() -> Self {
-        Self {
-            next: Cell::new(ptr::null_mut()),
-        }
+        Self { next: Cell::new(None) }
     }
 }
 
 pub(crate) struct TimerQueue {
-    head: Cell<*mut TaskHeader>,
+    head: Cell<Option<TaskRef>>,
 }
 
 impl TimerQueue {
     pub const fn new() -> Self {
-        Self {
-            head: Cell::new(ptr::null_mut()),
-        }
+        Self { head: Cell::new(None) }
     }
 
-    pub(crate) unsafe fn update(&self, p: NonNull<TaskHeader>) {
-        let task = p.as_ref();
+    pub(crate) unsafe fn update(&self, p: TaskRef) {
+        let task = p.header();
         if task.expires_at.get() != Instant::MAX {
             let old_state = task.state.fetch_or(STATE_TIMER_QUEUED, Ordering::AcqRel);
             let is_new = old_state & STATE_TIMER_QUEUED == 0;
 
             if is_new {
                 task.timer_queue_item.next.set(self.head.get());
-                self.head.set(p.as_ptr());
+                self.head.set(Some(p));
             }
         }
     }
@@ -47,7 +41,7 @@ impl TimerQueue {
     pub(crate) unsafe fn next_expiration(&self) -> Instant {
         let mut res = Instant::MAX;
         self.retain(|p| {
-            let task = p.as_ref();
+            let task = p.header();
             let expires = task.expires_at.get();
             res = min(res, expires);
             expires != Instant::MAX
@@ -55,9 +49,9 @@ impl TimerQueue {
         res
     }
 
-    pub(crate) unsafe fn dequeue_expired(&self, now: Instant, on_task: impl Fn(NonNull<TaskHeader>)) {
+    pub(crate) unsafe fn dequeue_expired(&self, now: Instant, on_task: impl Fn(TaskRef)) {
         self.retain(|p| {
-            let task = p.as_ref();
+            let task = p.header();
             if task.expires_at.get() <= now {
                 on_task(p);
                 false
@@ -67,11 +61,10 @@ impl TimerQueue {
         });
     }
 
-    pub(crate) unsafe fn retain(&self, mut f: impl FnMut(NonNull<TaskHeader>) -> bool) {
+    pub(crate) unsafe fn retain(&self, mut f: impl FnMut(TaskRef) -> bool) {
         let mut prev = &self.head;
-        while !prev.get().is_null() {
-            let p = NonNull::new_unchecked(prev.get());
-            let task = &*p.as_ptr();
+        while let Some(p) = prev.get() {
+            let task = p.header();
             if f(p) {
                 // Skip to next
                 prev = &task.timer_queue_item.next;
diff --git a/embassy-executor/src/raw/util.rs b/embassy-executor/src/raw/util.rs
index ed5822188..2b1f6b6f3 100644
--- a/embassy-executor/src/raw/util.rs
+++ b/embassy-executor/src/raw/util.rs
@@ -25,9 +25,3 @@ impl<T> UninitCell<T> {
         ptr::drop_in_place(self.as_mut_ptr())
     }
 }
-
-impl<T: Copy> UninitCell<T> {
-    pub unsafe fn read(&self) -> T {
-        ptr::read(self.as_mut_ptr())
-    }
-}
diff --git a/embassy-executor/src/raw/waker.rs b/embassy-executor/src/raw/waker.rs
index 5765259f2..400b37fa9 100644
--- a/embassy-executor/src/raw/waker.rs
+++ b/embassy-executor/src/raw/waker.rs
@@ -1,8 +1,7 @@
 use core::mem;
-use core::ptr::NonNull;
 use core::task::{RawWaker, RawWakerVTable, Waker};
 
-use super::{wake_task, TaskHeader};
+use super::{wake_task, TaskHeader, TaskRef};
 
 const VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake, drop);
 
@@ -11,14 +10,14 @@ unsafe fn clone(p: *const ()) -> RawWaker {
 }
 
 unsafe fn wake(p: *const ()) {
-    wake_task(NonNull::new_unchecked(p as *mut TaskHeader))
+    wake_task(TaskRef::from_ptr(p as *const TaskHeader))
 }
 
 unsafe fn drop(_: *const ()) {
     // nop
 }
 
-pub(crate) unsafe fn from_task(p: NonNull<TaskHeader>) -> Waker {
+pub(crate) unsafe fn from_task(p: TaskRef) -> Waker {
     Waker::from_raw(RawWaker::new(p.as_ptr() as _, &VTABLE))
 }
 
@@ -33,7 +32,7 @@ pub(crate) unsafe fn from_task(p: NonNull<TaskHeader>) -> Waker {
 /// # Panics
 ///
 /// Panics if the waker is not created by the Embassy executor.
-pub fn task_from_waker(waker: &Waker) -> NonNull<TaskHeader> {
+pub fn task_from_waker(waker: &Waker) -> TaskRef {
     // safety: OK because WakerHack has the same layout as Waker.
     // This is not really guaranteed because the structs are `repr(Rust)`, it is
     // indeed the case in the current implementation.
@@ -43,8 +42,8 @@ pub fn task_from_waker(waker: &Waker) -> NonNull<TaskHeader> {
         panic!("Found waker not created by the Embassy executor. `embassy_time::Timer` only works with the Embassy executor.")
     }
 
-    // safety: we never create a waker with a null data pointer.
-    unsafe { NonNull::new_unchecked(hack.data as *mut TaskHeader) }
+    // safety: our wakers are always created with `TaskRef::as_ptr`
+    unsafe { TaskRef::from_ptr(hack.data as *const TaskHeader) }
 }
 
 struct WakerHack {
diff --git a/embassy-executor/src/spawner.rs b/embassy-executor/src/spawner.rs
index 400d973ff..7c0a0183c 100644
--- a/embassy-executor/src/spawner.rs
+++ b/embassy-executor/src/spawner.rs
@@ -1,7 +1,6 @@
 use core::future::poll_fn;
 use core::marker::PhantomData;
 use core::mem;
-use core::ptr::NonNull;
 use core::task::Poll;
 
 use super::raw;
@@ -22,12 +21,12 @@ use super::raw;
 /// Once you've invoked a task function and obtained a SpawnToken, you *must* spawn it.
 #[must_use = "Calling a task function does nothing on its own. You must spawn the returned SpawnToken, typically with Spawner::spawn()"]
 pub struct SpawnToken<S> {
-    raw_task: Option<NonNull<raw::TaskHeader>>,
+    raw_task: Option<raw::TaskRef>,
     phantom: PhantomData<*mut S>,
 }
 
 impl<S> SpawnToken<S> {
-    pub(crate) unsafe fn new(raw_task: NonNull<raw::TaskHeader>) -> Self {
+    pub(crate) unsafe fn new(raw_task: raw::TaskRef) -> Self {
         Self {
             raw_task: Some(raw_task),
             phantom: PhantomData,
@@ -90,10 +89,10 @@ impl Spawner {
     ///
     /// Panics if the current executor is not an Embassy executor.
     pub async fn for_current_executor() -> Self {
-        poll_fn(|cx| unsafe {
+        poll_fn(|cx| {
             let task = raw::task_from_waker(cx.waker());
-            let executor = (*task.as_ptr()).executor.get();
-            Poll::Ready(Self::new(&*executor))
+            let executor = unsafe { task.header().executor.get().unwrap_unchecked() };
+            Poll::Ready(Self::new(executor))
         })
         .await
     }
@@ -166,10 +165,10 @@ impl SendSpawner {
     ///
     /// Panics if the current executor is not an Embassy executor.
     pub async fn for_current_executor() -> Self {
-        poll_fn(|cx| unsafe {
+        poll_fn(|cx| {
             let task = raw::task_from_waker(cx.waker());
-            let executor = (*task.as_ptr()).executor.get();
-            Poll::Ready(Self::new(&*executor))
+            let executor = unsafe { task.header().executor.get().unwrap_unchecked() };
+            Poll::Ready(Self::new(executor))
         })
         .await
     }