3 files changed, 95 insertions, 46 deletions

diff --git a/embassy-nrf/src/uarte.rs b/embassy-nrf/src/uarte.rs
index 8aee11c47..9daed2f4a 100644
--- a/embassy-nrf/src/uarte.rs
+++ b/embassy-nrf/src/uarte.rs
@@ -17,8 +17,8 @@ use crate::hal::gpio::Port as GpioPort;
 use crate::hal::pac;
 use crate::hal::prelude::*;
 use crate::hal::target_constants::EASY_DMA_SIZE;
-use crate::interrupt;
 use crate::interrupt::OwnedInterrupt;
+use crate::{interrupt, util};
 
 pub use crate::hal::uarte::Pins;
 // Re-export SVD variants to allow user to directly set values.
@@ -131,6 +131,8 @@ where
     }
 
     pub fn free(self) -> (T, T::Interrupt, Pins) {
+        // Wait for the peripheral to be disabled from the ISR.
+        while self.instance.enable.read().enable().is_enabled() {}
         (self.instance, self.irq, self.pins)
     }
 
@@ -156,6 +158,13 @@ where
             uarte.events_endtx.reset();
             trace!("endtx");
             compiler_fence(Ordering::SeqCst);
+
+            if uarte.events_txstarted.read().bits() != 0 {
+                // The ENDTX was signal triggered because DMA has finished.
+                uarte.events_txstarted.reset();
+                try_disable = true;
+            }
+
             T::state().tx_done.signal(());
         }
 
@@ -170,6 +179,13 @@ where
             trace!("endrx");
             let len = uarte.rxd.amount.read().bits();
             compiler_fence(Ordering::SeqCst);
+
+            if uarte.events_rxstarted.read().bits() != 0 {
+                // The ENDRX was signal triggered because DMA buffer is full.
+                uarte.events_rxstarted.reset();
+                try_disable = true;
+            }
+
             T::state().rx_done.signal(len);
         }
 
@@ -204,7 +220,7 @@ impl<T: Instance> embassy::uart::Uart for Uarte<T> {
         // `mem::forget()` on a previous future after polling it once.
         assert!(!self.tx_started());
 
-        self.enable();
+        T::state().tx_done.reset();
 
         SendFuture {
             uarte: self,
@@ -227,7 +243,7 @@ impl<T: Instance> embassy::uart::Uart for Uarte<T> {
         // `mem::forget()` on a previous future after polling it once.
         assert!(!self.rx_started());
 
-        self.enable();
+        T::state().rx_done.reset();
 
         ReceiveFuture {
             uarte: self,
@@ -241,7 +257,7 @@ pub struct SendFuture<'a, T>
 where
     T: Instance,
 {
-    uarte: &'a Uarte<T>,
+    uarte: &'a mut Uarte<T>,
     buf: &'a [u8],
 }
 
@@ -259,7 +275,9 @@ where
                 .instance
                 .tasks_stoptx
                 .write(|w| unsafe { w.bits(1) });
-            T::state().tx_done.blocking_wait();
+
+            // TX is stopped almost instantly, spinning is fine.
+            while !T::state().tx_done.signaled() {}
         }
     }
 }
@@ -273,28 +291,34 @@ where
     fn poll(self: core::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let Self { uarte, buf } = unsafe { self.get_unchecked_mut() };
 
-        if !uarte.tx_started() {
-            let uarte = &uarte.instance;
-
-            T::state().tx_done.reset();
-
+        if T::state().tx_done.poll_wait(cx).is_pending() {
             let ptr = buf.as_ptr();
             let len = buf.len();
             assert!(len <= EASY_DMA_SIZE);
             // TODO: panic if buffer is not in SRAM
 
+            uarte.enable();
+
             compiler_fence(Ordering::SeqCst);
-            uarte.txd.ptr.write(|w| unsafe { w.ptr().bits(ptr as u32) });
             uarte
+                .instance
+                .txd
+                .ptr
+                .write(|w| unsafe { w.ptr().bits(ptr as u32) });
+            uarte
+                .instance
                 .txd
                 .maxcnt
                 .write(|w| unsafe { w.maxcnt().bits(len as _) });
 
             trace!("starttx");
-            uarte.tasks_starttx.write(|w| unsafe { w.bits(1) });
-        }
+            uarte.instance.tasks_starttx.write(|w| unsafe { w.bits(1) });
+            while !uarte.tx_started() {} // Make sure transmission has started
 
-        T::state().tx_done.poll_wait(cx).map(|()| Ok(()))
+            Poll::Pending
+        } else {
+            Poll::Ready(Ok(()))
+        }
     }
 }
 
@@ -303,7 +327,7 @@ pub struct ReceiveFuture<'a, T>
 where
     T: Instance,
 {
-    uarte: &'a Uarte<T>,
+    uarte: &'a mut Uarte<T>,
     buf: &'a mut [u8],
 }
 
@@ -313,14 +337,15 @@ where
 {
     fn drop(self: &mut Self) {
         if self.uarte.rx_started() {
-            trace!("stoprx");
+            trace!("stoprx (drop)");
 
             self.uarte.instance.events_rxstarted.reset();
             self.uarte
                 .instance
                 .tasks_stoprx
                 .write(|w| unsafe { w.bits(1) });
-            T::state().rx_done.blocking_wait();
+
+            util::low_power_wait_until(|| T::state().rx_done.signaled())
         }
     }
 }
@@ -334,27 +359,35 @@ where
     fn poll(self: core::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let Self { uarte, buf } = unsafe { self.get_unchecked_mut() };
 
-        if !uarte.rx_started() {
-            let uarte = &uarte.instance;
-
-            T::state().rx_done.reset();
-
-            let ptr = buf.as_ptr();
-            let len = buf.len();
-            assert!(len <= EASY_DMA_SIZE);
-
-            compiler_fence(Ordering::SeqCst);
-            uarte.rxd.ptr.write(|w| unsafe { w.ptr().bits(ptr as u32) });
-            uarte
-                .rxd
-                .maxcnt
-                .write(|w| unsafe { w.maxcnt().bits(len as _) });
-
-            trace!("startrx");
-            uarte.tasks_startrx.write(|w| unsafe { w.bits(1) });
+        match T::state().rx_done.poll_wait(cx) {
+            Poll::Pending if !uarte.rx_started() => {
+                let ptr = buf.as_ptr();
+                let len = buf.len();
+                assert!(len <= EASY_DMA_SIZE);
+
+                uarte.enable();
+
+                compiler_fence(Ordering::SeqCst);
+                uarte
+                    .instance
+                    .rxd
+                    .ptr
+                    .write(|w| unsafe { w.ptr().bits(ptr as u32) });
+                uarte
+                    .instance
+                    .rxd
+                    .maxcnt
+                    .write(|w| unsafe { w.maxcnt().bits(len as _) });
+
+                trace!("startrx");
+                uarte.instance.tasks_startrx.write(|w| unsafe { w.bits(1) });
+                while !uarte.rx_started() {} // Make sure reception has started
+
+                Poll::Pending
+            }
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(_) => Poll::Ready(Ok(())),
         }
-
-        T::state().rx_done.poll_wait(cx).map(|_| Ok(()))
     }
 }
 
@@ -365,8 +398,19 @@ where
 {
     /// Stops the ongoing reception and returns the number of bytes received.
     pub async fn stop(self) -> usize {
-        drop(self);
-        let len = T::state().rx_done.wait().await;
+        let len = if self.uarte.rx_started() {
+            trace!("stoprx (stop)");
+
+            self.uarte.instance.events_rxstarted.reset();
+            self.uarte
+                .instance
+                .tasks_stoprx
+                .write(|w| unsafe { w.bits(1) });
+            T::state().rx_done.wait().await
+        } else {
+            // Transfer was stopped before it even started. No bytes were sent.
+            0
+        };
         len as _
     }
 }
diff --git a/embassy-nrf/src/util/mod.rs b/embassy-nrf/src/util/mod.rs
index 2fd5453d3..cf3306545 100644
--- a/embassy-nrf/src/util/mod.rs
+++ b/embassy-nrf/src/util/mod.rs
@@ -1,2 +1,12 @@
 pub mod peripheral;
 pub mod ring_buffer;
+
+/// Low power blocking wait loop using WFE/SEV.
+pub fn low_power_wait_until(mut condition: impl FnMut() -> bool) {
+    while !condition() {
+        // WFE might "eat" an event that would have otherwise woken the executor.
+        cortex_m::asm::wfe();
+    }
+    // Retrigger an event to be transparent to the executor.
+    cortex_m::asm::sev();
+}
diff --git a/embassy/src/util/signal.rs b/embassy/src/util/signal.rs
index cb6410611..8e778d1e3 100644
--- a/embassy/src/util/signal.rs
+++ b/embassy/src/util/signal.rs
@@ -63,12 +63,7 @@ impl<T: Send> Signal<T> {
         futures::future::poll_fn(move |cx| self.poll_wait(cx))
     }
 
-    /// Blocks until the signal has been received.
-    ///
-    /// Returns immediately when [`poll_wait()`] has not been called before.
-    pub fn blocking_wait(&self) {
-        while cortex_m::interrupt::free(|_| {
-            matches!(unsafe { &*self.state.get() }, State::Waiting(_))
-        }) {}
+    pub fn signaled(&self) -> bool {
+        cortex_m::interrupt::free(|_| matches!(unsafe { &*self.state.get() }, State::Signaled(_)))
     }
 }