fix: wait for UART TX FIFO drain in async TX

vorago-shared-hal/src/uart/tx_async.rs

@@ -23,6 +23,7 @@ static TX_CONTEXTS: [Mutex<RefCell<TxContext>>; 2] =
 // Completion flag. Kept outside of the context structure as an atomic to avoid
 // critical section.
 static TX_DONE: [AtomicBool; 2] = [const { AtomicBool::new(false) }; 2];
+const EMPTY_SLICE: &[u8] = &[];
 
 /// This is a generic interrupt handler to handle asynchronous UART TX operations for a given
 /// UART bank.
@@ -48,7 +49,8 @@ pub fn on_interrupt_tx(bank: Bank) {
     // Safety: We documented that the user provided slice must outlive the future, so we convert
     // the raw pointer back to the slice here.
     let slice = unsafe { context.slice.get().unwrap() };
-    if context.progress >= slice.len() && !tx_status.tx_busy() {
+    let tx_fifo_empty = uart.read_state().tx_fifo().value() == 0;
+    if context.progress >= slice.len() && tx_fifo_empty && !tx_status.tx_busy() {
         uart.modify_irq_enabled(|mut value| {
             value.set_tx(false);
             value.set_tx_empty(false);
@@ -170,6 +172,49 @@ impl Drop for TxFuture {
     }
 }
 
+pub struct TxFlushFuture {
+    id: Bank,
+}
+
+impl TxFlushFuture {
+    pub fn new(tx: &mut Tx) -> Self {
+        TX_DONE[tx.id as usize].store(false, core::sync::atomic::Ordering::Relaxed);
+        tx.disable_interrupts();
+
+        critical_section::with(|cs| {
+            let context_ref = TX_CONTEXTS[tx.id as usize].borrow(cs);
+            let mut context = context_ref.borrow_mut();
+            unsafe { context.slice.set(EMPTY_SLICE) };
+            context.progress = 0;
+
+            // Ensure those are enabled inside a critical section at the same time. Can lead to
+            // weird glitches otherwise.
+            tx.enable_interrupts(
+                #[cfg(feature = "vor4x")]
+                true,
+            );
+            tx.enable();
+        });
+
+        Self { id: tx.id }
+    }
+}
+
+impl Future for TxFlushFuture {
+    type Output = Result<(), TxOverrunError>;
+
+    fn poll(
+        self: core::pin::Pin<&mut Self>,
+        cx: &mut core::task::Context<'_>,
+    ) -> core::task::Poll<Self::Output> {
+        UART_TX_WAKERS[self.id as usize].register(cx.waker());
+        if TX_DONE[self.id as usize].swap(false, core::sync::atomic::Ordering::Relaxed) {
+            return core::task::Poll::Ready(Ok(()));
+        }
+        core::task::Poll::Pending
+    }
+}
+
 pub struct TxAsync(Tx);
 
 impl TxAsync {
@@ -198,6 +243,11 @@ impl TxAsync {
         fut.await
     }
 
+    pub async fn flush(&mut self) -> Result<(), TxOverrunError> {
+        let fut = TxFlushFuture::new(&mut self.0);
+        fut.await
+    }
+
     pub fn release(self) -> Tx {
         self.0
     }
@@ -228,6 +278,6 @@ impl Write for TxAsync {
     }
 
     async fn flush(&mut self) -> Result<(), Self::Error> {
-        Ok(())
+        self.flush().await
     }
 }