From 6862ac56cb1c1a29bad3e872b45386774c8edd58 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=BE=D0=BC=D0=B0=D0=BD=20=D0=9A=D1=80=D0=B8=D0=B2?=
 =?UTF-8?q?=D0=B5=D0=BD=D0=BA=D0=BE=D0=B2?= <qwerty19106@gmail.com>
Date: Mon, 16 Sep 2024 11:02:27 +0400
Subject: [PATCH] Stm32: implement async flush for UART

---
 embassy-stm32/src/usart/mod.rs | 71 ++++++++++++++++++++++++++++------
 1 file changed, 59 insertions(+), 12 deletions(-)

diff --git a/embassy-stm32/src/usart/mod.rs b/embassy-stm32/src/usart/mod.rs
index 6f838cce5..333e01e36 100644
--- a/embassy-stm32/src/usart/mod.rs
+++ b/embassy-stm32/src/usart/mod.rs
@@ -69,6 +69,12 @@ unsafe fn on_interrupt(r: Regs, s: &'static State) {
             // disable idle line detection
             w.set_idleie(false);
         });
+    } else if cr1.tcie() && sr.tc() {
+        // Transmission complete detected
+        r.cr1().modify(|w| {
+            // disable Transmission complete interrupt
+            w.set_tcie(false);
+        });
     } else if cr1.rxneie() {
         // We cannot check the RXNE flag as it is auto-cleared by the DMA controller
 
@@ -420,7 +426,7 @@ impl<'d> UartTx<'d, Async> {
 
     /// Wait until transmission complete
     pub async fn flush(&mut self) -> Result<(), Error> {
-        self.blocking_flush()
+        flush(&self.info, &self.state).await
     }
 }
 
@@ -531,16 +537,40 @@ impl<'d, M: Mode> UartTx<'d, M> {
     }
 }
 
+/// Wait until transmission complete
+async fn flush(info: &Info, state: &State) -> Result<(), Error> {
+    let r = info.regs;
+    if r.cr1().read().te() && !sr(r).read().tc() {
+        r.cr1().modify(|w| {
+            // enable Transmission Complete interrupt
+            w.set_tcie(true);
+        });
+
+        compiler_fence(Ordering::SeqCst);
+
+        // future which completes when Transmission complete is detected
+        let abort = poll_fn(move |cx| {
+            state.rx_waker.register(cx.waker());
+
+            let sr = sr(r).read();
+            if sr.tc() {
+                // Transmission complete detected
+                return Poll::Ready(());
+            }
+
+            Poll::Pending
+        });
+
+        abort.await;
+    }
+
+    Ok(())
+}
+
 fn blocking_flush(info: &Info) -> Result<(), Error> {
     let r = info.regs;
-    while !sr(r).read().tc() {}
-
-    // Disable Transmitter and enable receiver after transmission complete for Half-Duplex mode
-    if r.cr3().read().hdsel() {
-        r.cr1().modify(|reg| {
-            reg.set_te(false);
-            reg.set_re(true);
-        });
+    if r.cr1().read().te() {
+        while !sr(r).read().tc() {}
     }
 
     Ok(())
@@ -621,7 +651,13 @@ impl<'d> UartRx<'d, Async> {
         // Call flush for Half-Duplex mode if some bytes were written and flush was not called.
         // It prevents reading of bytes which have just been written.
         if r.cr3().read().hdsel() && r.cr1().read().te() {
-            blocking_flush(self.info)?;
+            flush(&self.info, &self.state).await?;
+
+            // Disable Transmitter and enable Receiver after flush
+            r.cr1().modify(|reg| {
+                reg.set_re(true);
+                reg.set_te(false);
+            });
         }
 
         // make sure USART state is restored to neutral state when this future is dropped
@@ -960,6 +996,12 @@ impl<'d, M: Mode> UartRx<'d, M> {
         // It prevents reading of bytes which have just been written.
         if r.cr3().read().hdsel() && r.cr1().read().te() {
             blocking_flush(self.info)?;
+
+            // Disable Transmitter and enable Receiver after flush
+            r.cr1().modify(|reg| {
+                reg.set_re(true);
+                reg.set_te(false);
+            });
         }
 
         for b in buffer {
@@ -1155,6 +1197,11 @@ impl<'d> Uart<'d, Async> {
         self.tx.write(buffer).await
     }
 
+    /// Wait until transmission complete
+    pub async fn flush(&mut self) -> Result<(), Error> {
+        self.tx.flush().await
+    }
+
     /// Perform an asynchronous read into `buffer`
     pub async fn read(&mut self, buffer: &mut [u8]) -> Result<(), Error> {
         self.rx.read(buffer).await
@@ -1733,7 +1780,7 @@ impl embedded_io_async::Write for Uart<'_, Async> {
     }
 
     async fn flush(&mut self) -> Result<(), Self::Error> {
-        self.blocking_flush()
+        self.flush().await
     }
 }
 
@@ -1744,7 +1791,7 @@ impl embedded_io_async::Write for UartTx<'_, Async> {
     }
 
     async fn flush(&mut self) -> Result<(), Self::Error> {
-        self.blocking_flush()
+        self.flush().await
     }
 }