From 72bfe0733ba10e8198b0fbf01dc504aec4675ae9 Mon Sep 17 00:00:00 2001
From: notgull <jtnunley01@gmail.com>
Date: Sun, 2 Apr 2023 11:24:24 -0700
Subject: [PATCH] Review comments

---
 src/lib.rs    |  2 +-
 src/no_std.rs | 85 ++++++++++++++++++++++-----------------------------
 src/std.rs    | 12 +++++---
 3 files changed, 44 insertions(+), 55 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 6351d7f..0868105 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -862,6 +862,7 @@ enum TaskRef<'a> {
 
 impl TaskRef<'_> {
     /// Tells if this task will wake up the other task.
+    #[allow(unreachable_patterns)]
     fn will_wake(self, other: Self) -> bool {
         match (self, other) {
             (Self::Waker(a), Self::Waker(b)) => a.will_wake(b),
@@ -870,7 +871,6 @@ impl TaskRef<'_> {
                 // TODO: Use unreleased will_unpark API.
                 false
             }
-            #[cfg(feature = "std")]
             _ => false,
         }
     }
diff --git a/src/no_std.rs b/src/no_std.rs
index 051aeeb..36727d2 100644
--- a/src/no_std.rs
+++ b/src/no_std.rs
@@ -287,38 +287,46 @@ enum Entry {
     Sentinel,
 }
 
-impl fmt::Debug for Entry {
+struct TakenState<'a> {
+    slot: &'a Cell<State>,
+    state: State,
+}
+
+impl Drop for TakenState<'_> {
+    fn drop(&mut self) {
+        self.slot
+            .set(mem::replace(&mut self.state, State::NotifiedTaken));
+    }
+}
+
+impl fmt::Debug for TakenState<'_> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        struct TakenState<'a> {
-            state: Option<State>,
-            cell: &'a Cell<State>,
-        }
+        fmt::Debug::fmt(&self.state, f)
+    }
+}
 
-        impl Drop for TakenState<'_> {
-            fn drop(&mut self) {
-                self.cell.set(self.state.take().unwrap());
-            }
-        }
+impl PartialEq for TakenState<'_> {
+    fn eq(&self, other: &Self) -> bool {
+        self.state == other.state
+    }
+}
 
-        impl fmt::Debug for TakenState<'_> {
-            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-                fmt::Debug::fmt(self.state.as_ref().unwrap(), f)
-            }
-        }
+impl<'a> TakenState<'a> {
+    fn new(slot: &'a Cell<State>) -> Self {
+        let state = slot.replace(State::NotifiedTaken);
+        Self { slot, state }
+    }
+}
 
+impl fmt::Debug for Entry {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            Entry::Listener { state, next, prev } => {
-                let taken = TakenState {
-                    state: Some(state.replace(State::Created)),
-                    cell: state,
-                };
-
-                f.debug_struct("Listener")
-                    .field("state", &taken)
-                    .field("prev", prev)
-                    .field("next", next)
-                    .finish()
-            }
+            Entry::Listener { state, next, prev } => f
+                .debug_struct("Listener")
+                .field("state", &TakenState::new(state))
+                .field("prev", prev)
+                .field("next", next)
+                .finish(),
             Entry::Empty(next) => f.debug_tuple("Empty").field(next).finish(),
             Entry::Sentinel => f.debug_tuple("Sentinel").finish(),
         }
@@ -327,17 +335,6 @@ impl fmt::Debug for Entry {
 
 impl PartialEq for Entry {
     fn eq(&self, other: &Entry) -> bool {
-        struct RestoreState<'a> {
-            state: Option<State>,
-            cell: &'a Cell<State>,
-        }
-
-        impl Drop for RestoreState<'_> {
-            fn drop(&mut self) {
-                self.cell.set(self.state.take().unwrap());
-            }
-        }
-
         match (self, other) {
             (
                 Self::Listener {
@@ -351,17 +348,7 @@ impl PartialEq for Entry {
                     next: next2,
                 },
             ) => {
-                let taken1 = RestoreState {
-                    state: Some(state1.replace(State::Created)),
-                    cell: state1,
-                };
-
-                let taken2 = RestoreState {
-                    state: Some(state2.replace(State::Created)),
-                    cell: state2,
-                };
-
-                if taken1.state != taken2.state {
+                if TakenState::new(state1) != TakenState::new(state2) {
                     return false;
                 }
 
diff --git a/src/std.rs b/src/std.rs
index bc87a3a..f317c18 100644
--- a/src/std.rs
+++ b/src/std.rs
@@ -141,11 +141,13 @@ impl crate::Inner {
 
             State::Task(other_task) => {
                 // Only replace the task if it's different.
-                if !task.will_wake(other_task.as_task_ref()) {
-                    entry.state.set(State::Task(task.into_task()));
-                } else {
-                    entry.state.set(State::Task(other_task));
-                }
+                entry.state.set(State::Task({
+                    if !task.will_wake(other_task.as_task_ref()) {
+                        task.into_task()
+                    } else {
+                        other_task
+                    }
+                }));
 
                 Some(false)
             }