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

1 files changed, 6 insertions, 7 deletions

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 {