From 4783222f675c1316d0e1a0fda92591f7aca63456 Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Fri, 1 Jan 2021 22:30:11 +0100
Subject: [PATCH] Rename WakerStore -> WakerRegistration.

---
 embassy-nrf/src/buffered_uarte.rs | 14 ++++++------
 embassy/src/util/mod.rs           |  4 ++--
 embassy/src/util/waker.rs         | 38 +++++++++++++++++++++++++++++++
 embassy/src/util/waker_store.rs   | 23 -------------------
 4 files changed, 47 insertions(+), 32 deletions(-)
 create mode 100644 embassy/src/util/waker.rs
 delete mode 100644 embassy/src/util/waker_store.rs

diff --git a/embassy-nrf/src/buffered_uarte.rs b/embassy-nrf/src/buffered_uarte.rs
index db6a83fb..6227e332 100644
--- a/embassy-nrf/src/buffered_uarte.rs
+++ b/embassy-nrf/src/buffered_uarte.rs
@@ -26,7 +26,7 @@ use crate::pac::{uarte0, UARTE0};
 pub use uarte0::{baudrate::BAUDRATE_A as Baudrate, config::PARITY_A as Parity};
 
 use embassy::io::{AsyncBufRead, AsyncWrite, Result};
-use embassy::util::WakerStore;
+use embassy::util::WakerRegistration;
 
 use crate::fmt::{assert, panic, todo, *};
 
@@ -147,11 +147,11 @@ pub struct UarteState<T: Instance> {
 
     rx: RingBuf,
     rx_state: RxState,
-    rx_waker: WakerStore,
+    rx_waker: WakerRegistration,
 
     tx: RingBuf,
     tx_state: TxState,
-    tx_waker: WakerStore,
+    tx_waker: WakerRegistration,
 
     _pin: PhantomPinned,
 }
@@ -233,11 +233,11 @@ impl<T: Instance> BufferedUarte<T> {
 
                 rx: RingBuf::new(),
                 rx_state: RxState::Idle,
-                rx_waker: WakerStore::new(),
+                rx_waker: WakerRegistration::new(),
 
                 tx: RingBuf::new(),
                 tx_state: TxState::Idle,
-                tx_waker: WakerStore::new(),
+                tx_waker: WakerRegistration::new(),
 
                 _pin: PhantomPinned,
             }),
@@ -327,7 +327,7 @@ impl<T: Instance> UarteState<T> {
             this.inner.tasks_stoprx.write(|w| unsafe { w.bits(1) });
         }
 
-        this.rx_waker.store(cx.waker());
+        this.rx_waker.register(cx.waker());
         Poll::Pending
     }
 
@@ -346,7 +346,7 @@ impl<T: Instance> UarteState<T> {
         let tx_buf = this.tx.push_buf();
         if tx_buf.len() == 0 {
             trace!("poll_write: pending");
-            this.tx_waker.store(cx.waker());
+            this.tx_waker.register(cx.waker());
             return Poll::Pending;
         }
 
diff --git a/embassy/src/util/mod.rs b/embassy/src/util/mod.rs
index 94745b83..5694d6bf 100644
--- a/embassy/src/util/mod.rs
+++ b/embassy/src/util/mod.rs
@@ -2,10 +2,10 @@ mod drop_bomb;
 mod forever;
 mod portal;
 mod signal;
-mod waker_store;
+mod waker;
 
 pub use drop_bomb::*;
 pub use forever::*;
 pub use portal::*;
 pub use signal::*;
-pub use waker_store::*;
+pub use waker::*;
diff --git a/embassy/src/util/waker.rs b/embassy/src/util/waker.rs
new file mode 100644
index 00000000..9735438b
--- /dev/null
+++ b/embassy/src/util/waker.rs
@@ -0,0 +1,38 @@
+use core::task::Context;
+use core::task::Waker;
+
+/// Utility struct to register and wake a waker.
+#[derive(Debug)]
+pub struct WakerRegistration {
+    waker: Option<Waker>,
+}
+
+impl WakerRegistration {
+    pub const fn new() -> Self {
+        Self { waker: None }
+    }
+
+    /// Register a waker. Overwrites the previous waker, if any.
+    pub fn register(&mut self, w: &Waker) {
+        match self.waker {
+            // Optimization: If both the old and new Wakers wake the same task, we can simply
+            // keep the old waker, skipping the clone. (In most executor implementations,
+            // cloning a waker is somewhat expensive, comparable to cloning an Arc).
+            Some(ref w2) if (w2.will_wake(w)) => {}
+            // In all other cases
+            // - we have no waker registered
+            // - we have a waker registered but it's for a different task.
+            // then clone the new waker and store it
+            _ => self.waker = Some(w.clone()),
+        }
+    }
+
+    /// Wake the registered waker, if any.
+    pub fn wake(&mut self) {
+        self.waker.take().map(|w| w.wake());
+    }
+
+    pub fn context(&self) -> Option<Context<'_>> {
+        self.waker.as_ref().map(|w| Context::from_waker(w))
+    }
+}
\ No newline at end of file
diff --git a/embassy/src/util/waker_store.rs b/embassy/src/util/waker_store.rs
deleted file mode 100644
index 0b2f09f4..00000000
--- a/embassy/src/util/waker_store.rs
+++ /dev/null
@@ -1,23 +0,0 @@
-use core::task::Waker;
-
-pub struct WakerStore {
-    waker: Option<Waker>,
-}
-
-impl WakerStore {
-    pub const fn new() -> Self {
-        Self { waker: None }
-    }
-
-    pub fn store(&mut self, w: &Waker) {
-        match self.waker {
-            Some(ref w2) if (w2.will_wake(w)) => {}
-            Some(_) => panic!("Waker overflow"),
-            None => self.waker = Some(w.clone()),
-        }
-    }
-
-    pub fn wake(&mut self) {
-        self.waker.take().map(|w| w.wake());
-    }
-}