Make some things more consistent

2 files changed, 31 insertions, 53 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 9b8a4ea8a..2e5941ef7 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -40,7 +40,6 @@ use core::marker::PhantomData;
 use core::mem;
 use core::pin::Pin;
 use core::ptr::NonNull;
-#[cfg(feature = "drs-scheduler")]
 use core::ptr::addr_of_mut;
 #[cfg(not(feature = "arch-avr"))]
 use core::sync::atomic::AtomicPtr;
@@ -51,9 +50,7 @@ use embassy_executor_timer_queue::TimerQueueItem;
 #[cfg(feature = "arch-avr")]
 use portable_atomic::AtomicPtr;
 
-use self::run_queue::RunQueue;
-#[cfg(not(feature = "drs-scheduler"))]
-use self::run_queue::RunQueueItem;
+use self::run_queue::{RunQueue, RunQueueItem};
 use self::state::State;
 use self::util::{SyncUnsafeCell, UninitCell};
 pub use self::waker::task_from_waker;
@@ -66,9 +63,6 @@ extern "Rust" fn __embassy_time_queue_item_from_waker(waker: &Waker) -> &'static
 }
 
 #[cfg(feature = "drs-scheduler")]
-use cordyceps::{stack, Linked};
-
-#[cfg(feature = "drs-scheduler")]
 pub use run_queue::Deadline;
 
 /// Raw task header for use in task pointers.
@@ -110,31 +104,14 @@ pub use run_queue::Deadline;
 /// - 4: A run-queued task exits - `TaskStorage::poll -> Poll::Ready`
 /// - 5: Task is dequeued. The task's future is not polled, because exiting the task replaces its `poll_fn`.
 /// - 6: A task is waken when it is not spawned - `wake_task -> State::run_enqueue`
-#[cfg_attr(feature = "drs-scheduler", repr(C))]
 pub(crate) struct TaskHeader {
-    // TODO(AJM): Make a decision whether we want to support the spicier "pointer recast"/"type punning"
-    // method of implementing the `cordyceps::Linked` trait or not.
-    //
-    // Currently, I do the safer version with `addr_of_mut!`, which doesn't REQUIRE that the first
-    // element is the `links` field, at the potential cost of a little extra pointer math.
-    //
-    // The optimizer *might* (total guess) notice that we are always doing an offset of zero in the
-    // call to `addr_of_mut` in the `impl Linked for TaskHeader` below, and get the best of both worlds,
-    // but right now this is maybe a little over cautious.
-    //
-    // See https://docs.rs/cordyceps/latest/cordyceps/trait.Linked.html#implementing-linkedlinks for
-    // more context on the choices here.
-    #[cfg(feature = "drs-scheduler")]
-    pub(crate) links: stack::Links<TaskHeader>,
+    pub(crate) run_queue_item: RunQueueItem,
 
     #[cfg(feature = "drs-scheduler")]
+    /// Deadline Rank Scheduler Deadline. This field should not be accessed outside the context of
+    /// the task itself as it being polled by the executor.
     pub(crate) deadline: SyncUnsafeCell<u64>,
 
-    // TODO(AJM): We could potentially replace RunQueueItem for other runqueue impls, though
-    // right now cordyceps doesn't work on non-atomic systems
-    #[cfg(not(feature = "drs-scheduler"))]
-    pub(crate) run_queue_item: RunQueueItem,
-
     pub(crate) state: State,
     pub(crate) executor: AtomicPtr<SyncExecutor>,
     poll_fn: SyncUnsafeCell<Option<unsafe fn(TaskRef)>>,
@@ -148,25 +125,6 @@ pub(crate) struct TaskHeader {
     all_tasks_next: AtomicPtr<TaskHeader>,
 }
 
-#[cfg(feature = "drs-scheduler")]
-unsafe impl Linked<stack::Links<TaskHeader>> for TaskHeader {
-    type Handle = TaskRef;
-
-    fn into_ptr(r: Self::Handle) -> NonNull<Self> {
-        r.ptr.cast()
-    }
-
-    unsafe fn from_ptr(ptr: NonNull<Self>) -> Self::Handle {
-        let ptr: NonNull<TaskHeader> = ptr;
-        TaskRef { ptr }
-    }
-
-    unsafe fn links(ptr: NonNull<Self>) -> NonNull<stack::Links<TaskHeader>> {
-        let ptr: *mut TaskHeader = ptr.as_ptr();
-        NonNull::new_unchecked(addr_of_mut!((*ptr).links))
-    }
-}
-
 /// This is essentially a `&'static TaskStorage<F>` where the type of the future has been erased.
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub struct TaskRef {
@@ -257,11 +215,8 @@ impl<F: Future + 'static> TaskStorage<F> {
     pub const fn new() -> Self {
         Self {
             raw: TaskHeader {
-                #[cfg(not(feature = "drs-scheduler"))]
                 run_queue_item: RunQueueItem::new(),
                 #[cfg(feature = "drs-scheduler")]
-                links: stack::Links::new(),
-                #[cfg(feature = "drs-scheduler")]
                 deadline: SyncUnsafeCell::new(0u64),
                 state: State::new(),
                 executor: AtomicPtr::new(core::ptr::null_mut()),
diff --git a/embassy-executor/src/raw/run_queue_drs_atomics.rs b/embassy-executor/src/raw/run_queue_drs_atomics.rs
index 69b7b3bf0..047265954 100644
--- a/embassy-executor/src/raw/run_queue_drs_atomics.rs
+++ b/embassy-executor/src/raw/run_queue_drs_atomics.rs
@@ -2,6 +2,29 @@ use super::{TaskHeader, TaskRef};
 use cordyceps::{SortedList, TransferStack};
 use core::future::{Future, poll_fn};
 use core::task::Poll;
+use core::ptr::{addr_of_mut, NonNull};
+use cordyceps::sorted_list::Links;
+use cordyceps::Linked;
+
+pub(crate) type RunQueueItem = Links<TaskHeader>;
+
+unsafe impl Linked<Links<TaskHeader>> for super::TaskHeader {
+    type Handle = TaskRef;
+
+    fn into_ptr(r: Self::Handle) -> NonNull<Self> {
+        r.ptr.cast()
+    }
+
+    unsafe fn from_ptr(ptr: NonNull<Self>) -> Self::Handle {
+        let ptr: NonNull<TaskHeader> = ptr;
+        TaskRef { ptr }
+    }
+
+    unsafe fn links(ptr: NonNull<Self>) -> NonNull<Links<TaskHeader>> {
+        let ptr: *mut TaskHeader = ptr.as_ptr();
+        NonNull::new_unchecked(addr_of_mut!((*ptr).run_queue_item))
+    }
+}
 
 /// Atomic task queue using a very, very simple lock-free linked-list queue:
 ///
@@ -39,10 +62,10 @@ impl RunQueue {
     /// 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(TaskRef)) {
-        let mut sorted = SortedList::<TaskHeader>::new(|lhs, rhs| unsafe {
-            // TODO: Do we need any kind of access control here? Not if we say that
-            // tasks can only set their own priority, which they can't do if we're in
-            // the scheduler
+        // SAFETY: `deadline` can only be set through the `Deadline` interface, which
+        // only allows access to this value while the given task is being polled.
+        // This acts as mutual exclusion for access.
+        let mut sorted = SortedList::<TaskHeader>::new_custom(|lhs, rhs| unsafe {
             lhs.deadline.get().cmp(&rhs.deadline.get())
         });