executor: Replace `NonNull<TaskHeader>` with `TaskRef`

1 files changed, 14 insertions, 21 deletions

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;