stm32/low-power: refactor refcount

6 files changed, 29 insertions, 35 deletions

diff --git a/embassy-stm32/build.rs b/embassy-stm32/build.rs
index 938e44b14..3400529c9 100644
--- a/embassy-stm32/build.rs
+++ b/embassy-stm32/build.rs
@@ -564,7 +564,7 @@ fn main() {
                     fn enable_and_reset_with_cs(_cs: critical_section::CriticalSection) {
                         #before_enable
                         #[cfg(feature = "low-power")]
-                        crate::rcc::clock_refcount_add(_cs);
+                        unsafe { crate::rcc::REFCOUNT_STOP2 += 1 };
                         crate::pac::RCC.#en_reg().modify(|w| w.#set_en_field(true));
                         #after_enable
                         #rst
@@ -573,7 +573,7 @@ fn main() {
                         #before_disable
                         crate::pac::RCC.#en_reg().modify(|w| w.#set_en_field(false));
                         #[cfg(feature = "low-power")]
-                        crate::rcc::clock_refcount_sub(_cs);
+                        unsafe { crate::rcc::REFCOUNT_STOP2 -= 1 };
                     }
                 }
 
diff --git a/embassy-stm32/src/lib.rs b/embassy-stm32/src/lib.rs
index 372246f87..5c7067d2d 100644
--- a/embassy-stm32/src/lib.rs
+++ b/embassy-stm32/src/lib.rs
@@ -226,9 +226,9 @@ pub fn init(config: Config) -> Peripherals {
             time_driver::init(cs);
 
             #[cfg(feature = "low-power")]
-            while !crate::rcc::low_power_ready() {
-                crate::rcc::clock_refcount_sub(cs);
-            }
+            {
+                crate::rcc::REFCOUNT_STOP2 = 0
+            };
         }
 
         p
diff --git a/embassy-stm32/src/low_power.rs b/embassy-stm32/src/low_power.rs
index 861a59d7b..d5846f530 100644
--- a/embassy-stm32/src/low_power.rs
+++ b/embassy-stm32/src/low_power.rs
@@ -6,7 +6,6 @@ use cortex_m::peripheral::SCB;
 use embassy_executor::*;
 
 use crate::interrupt;
-use crate::rcc::low_power_ready;
 use crate::time_driver::{get_driver, RtcDriver};
 
 const THREAD_PENDER: usize = usize::MAX;
@@ -33,6 +32,15 @@ pub fn stop_with_rtc(rtc: &'static Rtc) {
     unsafe { EXECUTOR.as_mut().unwrap() }.stop_with_rtc(rtc)
 }
 
+pub fn stop_ready(stop_mode: StopMode) -> bool {
+    unsafe { EXECUTOR.as_mut().unwrap() }.stop_ready(stop_mode)
+}
+
+#[non_exhaustive]
+pub enum StopMode {
+    Stop2,
+}
+
 /// Thread mode executor, using WFE/SEV.
 ///
 /// This is the simplest and most common kind of executor. It runs on
@@ -80,12 +88,18 @@ impl Executor {
         trace!("low power: stop with rtc configured");
     }
 
+    fn stop_ready(&self, stop_mode: StopMode) -> bool {
+        match stop_mode {
+            StopMode::Stop2 => unsafe { crate::rcc::REFCOUNT_STOP2 == 0 },
+        }
+    }
+
     fn configure_pwr(&mut self) {
         self.scb.clear_sleepdeep();
 
         compiler_fence(Ordering::SeqCst);
 
-        if !low_power_ready() {
+        if !self.stop_ready(StopMode::Stop2) {
             trace!("low power: not ready to stop");
         } else if self.time_driver.pause_time().is_err() {
             trace!("low power: failed to pause time");
diff --git a/embassy-stm32/src/rcc/mod.rs b/embassy-stm32/src/rcc/mod.rs
index 8cf2d6ab0..3b19e4b95 100644
--- a/embassy-stm32/src/rcc/mod.rs
+++ b/embassy-stm32/src/rcc/mod.rs
@@ -23,8 +23,6 @@ pub use mco::*;
 #[cfg_attr(rcc_u5, path = "u5.rs")]
 #[cfg_attr(rcc_wba, path = "wba.rs")]
 mod _version;
-#[cfg(feature = "low-power")]
-use core::sync::atomic::{AtomicU32, Ordering};
 
 pub use _version::*;
 
@@ -183,27 +181,7 @@ pub struct Clocks {
 }
 
 #[cfg(feature = "low-power")]
-static CLOCK_REFCOUNT: AtomicU32 = AtomicU32::new(0);
-
-#[cfg(feature = "low-power")]
-pub fn low_power_ready() -> bool {
-    // trace!("clock refcount: {}", CLOCK_REFCOUNT.load(Ordering::SeqCst));
-    CLOCK_REFCOUNT.load(Ordering::SeqCst) == 0
-}
-
-#[cfg(feature = "low-power")]
-pub(crate) fn clock_refcount_add(_cs: critical_section::CriticalSection) {
-    // We don't check for overflow because constructing more than u32 peripherals is unlikely
-    let n = CLOCK_REFCOUNT.load(Ordering::Relaxed);
-    CLOCK_REFCOUNT.store(n + 1, Ordering::Relaxed);
-}
-
-#[cfg(feature = "low-power")]
-pub(crate) fn clock_refcount_sub(_cs: critical_section::CriticalSection) {
-    let n = CLOCK_REFCOUNT.load(Ordering::Relaxed);
-    assert!(n != 0);
-    CLOCK_REFCOUNT.store(n - 1, Ordering::Relaxed);
-}
+pub(crate) static mut REFCOUNT_STOP2: u32 = 0;
 
 /// Frozen clock frequencies
 ///
diff --git a/embassy-stm32/src/rtc/mod.rs b/embassy-stm32/src/rtc/mod.rs
index d77443a9c..3334262f9 100644
--- a/embassy-stm32/src/rtc/mod.rs
+++ b/embassy-stm32/src/rtc/mod.rs
@@ -187,7 +187,9 @@ impl Rtc {
         critical_section::with(|cs| {
             <RTC as crate::rcc::sealed::RccPeripheral>::enable_and_reset_with_cs(cs);
             #[cfg(feature = "low-power")]
-            crate::rcc::clock_refcount_sub(cs);
+            unsafe {
+                crate::rcc::REFCOUNT_STOP2 -= 1
+            };
         });
 
         let mut this = Self {
diff --git a/tests/stm32/src/bin/stop.rs b/tests/stm32/src/bin/stop.rs
index 4f62ecae2..b9810673a 100644
--- a/tests/stm32/src/bin/stop.rs
+++ b/tests/stm32/src/bin/stop.rs
@@ -10,8 +10,8 @@ use chrono::NaiveDate;
 use common::*;
 use cortex_m_rt::entry;
 use embassy_executor::Spawner;
-use embassy_stm32::low_power::{stop_with_rtc, Executor};
-use embassy_stm32::rcc::{low_power_ready, LsConfig};
+use embassy_stm32::low_power::{stop_ready, stop_with_rtc, Executor, StopMode};
+use embassy_stm32::rcc::LsConfig;
 use embassy_stm32::rtc::{Rtc, RtcConfig};
 use embassy_stm32::Config;
 use embassy_time::Timer;
@@ -28,7 +28,7 @@ fn main() -> ! {
 async fn task_1() {
     for _ in 0..9 {
         info!("task 1: waiting for 500ms...");
-        defmt::assert!(low_power_ready());
+        defmt::assert!(stop_ready(StopMode::Stop2));
         Timer::after_millis(500).await;
     }
 }
@@ -37,7 +37,7 @@ async fn task_1() {
 async fn task_2() {
     for _ in 0..5 {
         info!("task 2: waiting for 1000ms...");
-        defmt::assert!(low_power_ready());
+        defmt::assert!(stop_ready(StopMode::Stop2));
         Timer::after_millis(1000).await;
     }