stm32: fix refcounts for usart, spi, and i2c

7 files changed, 60 insertions, 17 deletions

diff --git a/embassy-stm32/build.rs b/embassy-stm32/build.rs
index 2b66e7dad..d3f6452b4 100644
--- a/embassy-stm32/build.rs
+++ b/embassy-stm32/build.rs
@@ -347,9 +347,7 @@ fn main() {
                 TokenStream::new()
             };
 
-            let ptype = p
-                .name
-                .replace(&['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'][..], "");
+            let ptype = p.name.replace(|c| char::is_numeric(c), "");
             let pname = format_ident!("{}", p.name);
             let clk = format_ident!("{}", rcc.clock.to_ascii_lowercase());
             let en_reg = format_ident!("{}", en.register.to_ascii_lowercase());
@@ -363,16 +361,14 @@ fn main() {
 
                 (
                     quote! {
-                        use atomic_polyfill::Ordering;
-
-                        if refcount_statics::#refcount_static.fetch_add(1, Ordering::SeqCst) > 0 {
+                        unsafe { refcount_statics::#refcount_static += 1 };
+                        if unsafe { refcount_statics::#refcount_static } > 1 {
                             return;
                         }
                     },
                     quote! {
-                        use atomic_polyfill::Ordering;
-
-                        if refcount_statics::#refcount_static.fetch_sub(1, Ordering::SeqCst) > 1 {
+                        unsafe { refcount_statics::#refcount_static -= 1 };
+                        if unsafe { refcount_statics::#refcount_static } > 0  {
                             return;
                         }
                     },
@@ -416,14 +412,12 @@ fn main() {
     let mut refcount_mod = TokenStream::new();
     for refcount_static in refcount_statics {
         refcount_mod.extend(quote! {
-            pub(crate) static #refcount_static: AtomicU8 = AtomicU8::new(0);
+            pub(crate) static mut #refcount_static: u8 = 0;
         });
     }
 
     g.extend(quote! {
         mod refcount_statics {
-            use atomic_polyfill::AtomicU8;
-
             #refcount_mod
         }
     });
diff --git a/embassy-stm32/src/i2c/v1.rs b/embassy-stm32/src/i2c/v1.rs
index 618d85af2..f32dd0f0c 100644
--- a/embassy-stm32/src/i2c/v1.rs
+++ b/embassy-stm32/src/i2c/v1.rs
@@ -339,6 +339,12 @@ impl<'d, T: Instance, TXDMA, RXDMA> I2c<'d, T, TXDMA, RXDMA> {
     }
 }
 
+impl<'d, T: Instance, TXDMA, RXDMA> Drop for I2c<'d, T, TXDMA, RXDMA> {
+    fn drop(&mut self) {
+        T::disable();
+    }
+}
+
 impl<'d, T: Instance> embedded_hal_02::blocking::i2c::Read for I2c<'d, T> {
     type Error = Error;
 
diff --git a/embassy-stm32/src/i2c/v2.rs b/embassy-stm32/src/i2c/v2.rs
index 4327899bb..36f70e32e 100644
--- a/embassy-stm32/src/i2c/v2.rs
+++ b/embassy-stm32/src/i2c/v2.rs
@@ -838,6 +838,12 @@ impl<'d, T: Instance, TXDMA, RXDMA> I2c<'d, T, TXDMA, RXDMA> {
     }
 }
 
+impl<'d, T: Instance, TXDMA, RXDMA> Drop for I2c<'d, T, TXDMA, RXDMA> {
+    fn drop(&mut self) {
+        T::disable();
+    }
+}
+
 mod eh02 {
     use super::*;
 
diff --git a/embassy-stm32/src/spi/mod.rs b/embassy-stm32/src/spi/mod.rs
index e2bc8d7f2..853de98f9 100644
--- a/embassy-stm32/src/spi/mod.rs
+++ b/embassy-stm32/src/spi/mod.rs
@@ -646,6 +646,8 @@ impl<'d, T: Instance, Tx, Rx> Drop for Spi<'d, T, Tx, Rx> {
         self.sck.as_ref().map(|x| x.set_as_disconnected());
         self.mosi.as_ref().map(|x| x.set_as_disconnected());
         self.miso.as_ref().map(|x| x.set_as_disconnected());
+
+        T::disable();
     }
 }
 
diff --git a/embassy-stm32/src/usart/buffered.rs b/embassy-stm32/src/usart/buffered.rs
index 596d40bf9..989c88205 100644
--- a/embassy-stm32/src/usart/buffered.rs
+++ b/embassy-stm32/src/usart/buffered.rs
@@ -124,6 +124,8 @@ impl<'d, T: BasicInstance> BufferedUart<'d, T> {
         rx_buffer: &'d mut [u8],
         config: Config,
     ) -> BufferedUart<'d, T> {
+        // UartRx and UartTx have one refcount ea.
+        T::enable();
         T::enable();
         T::reset();
 
@@ -143,6 +145,8 @@ impl<'d, T: BasicInstance> BufferedUart<'d, T> {
     ) -> BufferedUart<'d, T> {
         into_ref!(cts, rts);
 
+        // UartRx and UartTx have one refcount ea.
+        T::enable();
         T::enable();
         T::reset();
 
@@ -169,6 +173,8 @@ impl<'d, T: BasicInstance> BufferedUart<'d, T> {
     ) -> BufferedUart<'d, T> {
         into_ref!(de);
 
+        // UartRx and UartTx have one refcount ea.
+        T::enable();
         T::enable();
         T::reset();
 
@@ -382,6 +388,8 @@ impl<'d, T: BasicInstance> Drop for BufferedUartRx<'d, T> {
                 T::Interrupt::disable();
             }
         }
+
+        T::disable();
     }
 }
 
@@ -397,6 +405,8 @@ impl<'d, T: BasicInstance> Drop for BufferedUartTx<'d, T> {
                 T::Interrupt::disable();
             }
         }
+
+        T::disable();
     }
 }
 
diff --git a/embassy-stm32/src/usart/mod.rs b/embassy-stm32/src/usart/mod.rs
index 255ddfd4b..bfb056718 100644
--- a/embassy-stm32/src/usart/mod.rs
+++ b/embassy-stm32/src/usart/mod.rs
@@ -618,6 +618,18 @@ impl<'d, T: BasicInstance, RxDma> UartRx<'d, T, RxDma> {
     }
 }
 
+impl<'d, T: BasicInstance, TxDma> Drop for UartTx<'d, T, TxDma> {
+    fn drop(&mut self) {
+        T::disable();
+    }
+}
+
+impl<'d, T: BasicInstance, TxDma> Drop for UartRx<'d, T, TxDma> {
+    fn drop(&mut self) {
+        T::disable();
+    }
+}
+
 impl<'d, T: BasicInstance, TxDma, RxDma> Uart<'d, T, TxDma, RxDma> {
     pub fn new(
         peri: impl Peripheral<P = T> + 'd,
@@ -628,6 +640,8 @@ impl<'d, T: BasicInstance, TxDma, RxDma> Uart<'d, T, TxDma, RxDma> {
         rx_dma: impl Peripheral<P = RxDma> + 'd,
         config: Config,
     ) -> Self {
+        // UartRx and UartTx have one refcount ea.
+        T::enable();
         T::enable();
         T::reset();
 
@@ -647,6 +661,8 @@ impl<'d, T: BasicInstance, TxDma, RxDma> Uart<'d, T, TxDma, RxDma> {
     ) -> Self {
         into_ref!(cts, rts);
 
+        // UartRx and UartTx have one refcount ea.
+        T::enable();
         T::enable();
         T::reset();
 
@@ -672,6 +688,8 @@ impl<'d, T: BasicInstance, TxDma, RxDma> Uart<'d, T, TxDma, RxDma> {
     ) -> Self {
         into_ref!(de);
 
+        // UartRx and UartTx have one refcount ea.
+        T::enable();
         T::enable();
         T::reset();
 
diff --git a/embassy-stm32/src/usart/ringbuffered.rs b/embassy-stm32/src/usart/ringbuffered.rs
index b3f570624..e990eaca8 100644
--- a/embassy-stm32/src/usart/ringbuffered.rs
+++ b/embassy-stm32/src/usart/ringbuffered.rs
@@ -1,4 +1,5 @@
 use core::future::poll_fn;
+use core::mem;
 use core::sync::atomic::{compiler_fence, Ordering};
 use core::task::Poll;
 
@@ -24,12 +25,16 @@ impl<'d, T: BasicInstance, RxDma: super::RxDma<T>> UartRx<'d, T, RxDma> {
         let request = self.rx_dma.request();
         let opts = Default::default();
 
-        let ring_buf = unsafe { ReadableRingBuffer::new_read(self.rx_dma, request, rdr(T::regs()), dma_buf, opts) };
+        // Safety: we forget the struct before this function returns.
+        let rx_dma = unsafe { self.rx_dma.clone_unchecked() };
+        let _peri = unsafe { self._peri.clone_unchecked() };
 
-        RingBufferedUartRx {
-            _peri: self._peri,
-            ring_buf,
-        }
+        let ring_buf = unsafe { ReadableRingBuffer::new_read(rx_dma, request, rdr(T::regs()), dma_buf, opts) };
+
+        // Don't disable the clock
+        mem::forget(self);
+
+        RingBufferedUartRx { _peri, ring_buf }
     }
 }
 
@@ -186,6 +191,8 @@ impl<'d, T: BasicInstance, RxDma: super::RxDma<T>> RingBufferedUartRx<'d, T, RxD
 impl<T: BasicInstance, RxDma: super::RxDma<T>> Drop for RingBufferedUartRx<'_, T, RxDma> {
     fn drop(&mut self) {
         self.teardown_uart();
+
+        T::disable();
     }
 }
 /// Return an error result if the Sr register has errors