embassy-rp: defensive change to ensure wakers are registered

This change ensures that wakers are registered PRIOR to checking status
in i2c `wait_on` helpers.

2 files changed, 8 insertions, 4 deletions

diff --git a/embassy-rp/src/i2c.rs b/embassy-rp/src/i2c.rs
index adc38b73d..a983b7bc3 100644
--- a/embassy-rp/src/i2c.rs
+++ b/embassy-rp/src/i2c.rs
@@ -114,17 +114,19 @@ impl<'d, T: Instance> I2c<'d, T, Async> {
     }
 
     /// Calls `f` to check if we are ready or not.
-    /// If not, `g` is called once the waker is set (to eg enable the required interrupts).
+    /// If not, `g` is called once(to eg enable the required interrupts).
+    /// The waker will always be registered prior to calling `f`.
     async fn wait_on<F, U, G>(&mut self, mut f: F, mut g: G) -> U
     where
         F: FnMut(&mut Self) -> Poll<U>,
         G: FnMut(&mut Self),
     {
         future::poll_fn(|cx| {
+            // Register prior to checking the condition
+            T::waker().register(cx.waker());
             let r = f(self);
 
             if r.is_pending() {
-                T::waker().register(cx.waker());
                 g(self);
             }
             r
diff --git a/embassy-rp/src/i2c_slave.rs b/embassy-rp/src/i2c_slave.rs
index d420030d8..7bc14511d 100644
--- a/embassy-rp/src/i2c_slave.rs
+++ b/embassy-rp/src/i2c_slave.rs
@@ -159,7 +159,8 @@ impl<'d, T: Instance> I2cSlave<'d, T> {
     }
 
     /// Calls `f` to check if we are ready or not.
-    /// If not, `g` is called once the waker is set (to eg enable the required interrupts).
+    /// If not, `g` is called once(to eg enable the required interrupts).
+    /// The waker will always be registered prior to calling `f`.
     #[inline(always)]
     async fn wait_on<F, U, G>(&mut self, mut f: F, mut g: G) -> U
     where
@@ -167,10 +168,11 @@ impl<'d, T: Instance> I2cSlave<'d, T> {
         G: FnMut(&mut Self),
     {
         future::poll_fn(|cx| {
+            // Register prior to checking the condition
+            T::waker().register(cx.waker());
             let r = f(self);
 
             if r.is_pending() {
-                T::waker().register(cx.waker());
                 g(self);
             }