1 files changed, 7 insertions, 6 deletions

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);