stm32: avoid creating many tiny critical sections in init.

Saves 292 bytes on stm32f0 bilnky with max optimizations (from 3132 to 2840).

14 files changed, 182 insertions, 148 deletions

diff --git a/embassy-hal-internal/src/interrupt.rs b/embassy-hal-internal/src/interrupt.rs
index b970aa2cd..19dabcf6f 100644
--- a/embassy-hal-internal/src/interrupt.rs
+++ b/embassy-hal-internal/src/interrupt.rs
@@ -4,6 +4,7 @@ use core::sync::atomic::{compiler_fence, Ordering};
 
 use cortex_m::interrupt::InterruptNumber;
 use cortex_m::peripheral::NVIC;
+use critical_section::CriticalSection;
 
 /// Generate a standard `mod interrupt` for a HAL.
 #[macro_export]
@@ -91,6 +92,12 @@ macro_rules! interrupt_mod {
                     fn set_priority(prio: crate::interrupt::Priority) {
                         Self::IRQ.set_priority(prio)
                     }
+
+                    /// Set the interrupt priority with an already-acquired critical section
+                    #[inline]
+                    fn set_priority_with_cs(cs: critical_section::CriticalSection, prio: crate::interrupt::Priority) {
+                        Self::IRQ.set_priority_with_cs(cs, prio)
+                    }
                 }
 
                 $(
@@ -195,10 +202,29 @@ pub unsafe trait InterruptExt: InterruptNumber + Copy {
     /// Set the interrupt priority.
     #[inline]
     fn set_priority(self, prio: Priority) {
-        critical_section::with(|_| unsafe {
+        unsafe {
+            let mut nvic: cortex_m::peripheral::NVIC = mem::transmute(());
+
+            // On thumbv6, set_priority must do a RMW to change 8bit in a 32bit reg.
+            #[cfg(armv6m)]
+            critical_section::with(|_| nvic.set_priority(self, prio.into()));
+            // On thumbv7+, set_priority does an atomic 8bit write, so no CS needed.
+            #[cfg(not(armv6m))]
+            nvic.set_priority(self, prio.into());
+        }
+    }
+
+    /// Set the interrupt priority with an already-acquired critical section
+    ///
+    /// Equivalent to `set_priority`, except you pass a `CriticalSection` to prove
+    /// you've already acquired a critical section. This prevents acquiring another
+    /// one, which saves code size.
+    #[inline]
+    fn set_priority_with_cs(self, _cs: CriticalSection, prio: Priority) {
+        unsafe {
             let mut nvic: cortex_m::peripheral::NVIC = mem::transmute(());
-            nvic.set_priority(self, prio.into())
-        })
+            nvic.set_priority(self, prio.into());
+        }
     }
 }
 
diff --git a/embassy-hal-internal/src/macros.rs b/embassy-hal-internal/src/macros.rs
index 0eea4b667..97df38954 100644
--- a/embassy-hal-internal/src/macros.rs
+++ b/embassy-hal-internal/src/macros.rs
@@ -48,17 +48,23 @@ macro_rules! peripherals_struct {
             ///Returns all the peripherals *once*
             #[inline]
             pub(crate) fn take() -> Self {
+                critical_section::with(Self::take_with_cs)
+            }
 
+            ///Returns all the peripherals *once*
+            #[inline]
+            pub(crate) fn take_with_cs(_cs: critical_section::CriticalSection) -> Self {
                 #[no_mangle]
                 static mut _EMBASSY_DEVICE_PERIPHERALS: bool = false;
 
-                critical_section::with(|_| unsafe {
+                // safety: OK because we're inside a CS.
+                unsafe {
                     if _EMBASSY_DEVICE_PERIPHERALS {
                         panic!("init called more than once!")
                     }
                     _EMBASSY_DEVICE_PERIPHERALS = true;
                     Self::steal()
-                })
+                }
             }
         }
 
diff --git a/embassy-stm32/build.rs b/embassy-stm32/build.rs
index bbdb1250c..8e680fb60 100644
--- a/embassy-stm32/build.rs
+++ b/embassy-stm32/build.rs
@@ -554,23 +554,19 @@ fn main() {
                     fn frequency() -> crate::time::Hertz {
                         #clock_frequency
                     }
-                    fn enable_and_reset() {
-                        critical_section::with(|_cs| {
-                            #before_enable
-                            #[cfg(feature = "low-power")]
-                            crate::rcc::clock_refcount_add(_cs);
-                            crate::pac::RCC.#en_reg().modify(|w| w.#set_en_field(true));
-                            #after_enable
-                            #rst
-                        })
+                    fn enable_and_reset_with_cs(_cs: critical_section::CriticalSection) {
+                        #before_enable
+                        #[cfg(feature = "low-power")]
+                        crate::rcc::clock_refcount_add(_cs);
+                        crate::pac::RCC.#en_reg().modify(|w| w.#set_en_field(true));
+                        #after_enable
+                        #rst
                     }
-                    fn disable() {
-                        critical_section::with(|_cs| {
-                            #before_disable
-                            crate::pac::RCC.#en_reg().modify(|w| w.#set_en_field(false));
-                            #[cfg(feature = "low-power")]
-                            crate::rcc::clock_refcount_sub(_cs);
-                        })
+                    fn disable_with_cs(_cs: critical_section::CriticalSection) {
+                        #before_disable
+                        crate::pac::RCC.#en_reg().modify(|w| w.#set_en_field(false));
+                        #[cfg(feature = "low-power")]
+                        crate::rcc::clock_refcount_sub(_cs);
                     }
                 }
 
diff --git a/embassy-stm32/src/dac/mod.rs b/embassy-stm32/src/dac/mod.rs
index 6458572f2..a3c7823cf 100644
--- a/embassy-stm32/src/dac/mod.rs
+++ b/embassy-stm32/src/dac/mod.rs
@@ -567,18 +567,14 @@ foreach_peripheral!(
                         critical_section::with(|_| unsafe { crate::rcc::get_freqs().apb1 })
                     }
 
-                    fn enable_and_reset() {
-                        critical_section::with(|_| {
-                            crate::pac::RCC.apb1lrstr().modify(|w| w.set_dac12rst(true));
-                            crate::pac::RCC.apb1lrstr().modify(|w| w.set_dac12rst(false));
-                            crate::pac::RCC.apb1lenr().modify(|w| w.set_dac12en(true));
-                        })
+                    fn enable_and_reset_with_cs(_cs: critical_section::CriticalSection) {
+                        crate::pac::RCC.apb1lrstr().modify(|w| w.set_dac12rst(true));
+                        crate::pac::RCC.apb1lrstr().modify(|w| w.set_dac12rst(false));
+                        crate::pac::RCC.apb1lenr().modify(|w| w.set_dac12en(true));
                     }
 
-                    fn disable() {
-                        critical_section::with(|_| {
-                            crate::pac::RCC.apb1lenr().modify(|w| w.set_dac12en(false))
-                        })
+                    fn disable_with_cs(_cs: critical_section::CriticalSection) {
+                        crate::pac::RCC.apb1lenr().modify(|w| w.set_dac12en(false))
                     }
                 }
 
diff --git a/embassy-stm32/src/dma/bdma.rs b/embassy-stm32/src/dma/bdma.rs
index 62eb65b1c..a7422f66b 100644
--- a/embassy-stm32/src/dma/bdma.rs
+++ b/embassy-stm32/src/dma/bdma.rs
@@ -77,10 +77,10 @@ impl State {
 static STATE: State = State::new();
 
 /// safety: must be called only once
-pub(crate) unsafe fn init(irq_priority: Priority) {
+pub(crate) unsafe fn init(cs: critical_section::CriticalSection, irq_priority: Priority) {
     foreach_interrupt! {
         ($peri:ident, bdma, $block:ident, $signal_name:ident, $irq:ident) => {
-            crate::interrupt::typelevel::$irq::set_priority(irq_priority);
+            crate::interrupt::typelevel::$irq::set_priority_with_cs(cs, irq_priority);
             crate::interrupt::typelevel::$irq::enable();
         };
     }
diff --git a/embassy-stm32/src/dma/dma.rs b/embassy-stm32/src/dma/dma.rs
index 5033ae477..cce0407c1 100644
--- a/embassy-stm32/src/dma/dma.rs
+++ b/embassy-stm32/src/dma/dma.rs
@@ -154,10 +154,10 @@ impl State {
 static STATE: State = State::new();
 
 /// safety: must be called only once
-pub(crate) unsafe fn init(irq_priority: Priority) {
+pub(crate) unsafe fn init(cs: critical_section::CriticalSection, irq_priority: Priority) {
     foreach_interrupt! {
         ($peri:ident, dma, $block:ident, $signal_name:ident, $irq:ident) => {
-            interrupt::typelevel::$irq::set_priority(irq_priority);
+            interrupt::typelevel::$irq::set_priority_with_cs(cs, irq_priority);
             interrupt::typelevel::$irq::enable();
         };
     }
diff --git a/embassy-stm32/src/dma/dmamux.rs b/embassy-stm32/src/dma/dmamux.rs
index 36fc03403..20601dc86 100644
--- a/embassy-stm32/src/dma/dmamux.rs
+++ b/embassy-stm32/src/dma/dmamux.rs
@@ -47,6 +47,6 @@ foreach_dma_channel! {
 }
 
 /// safety: must be called only once
-pub(crate) unsafe fn init() {
+pub(crate) unsafe fn init(_cs: critical_section::CriticalSection) {
     crate::_generated::init_dmamux();
 }
diff --git a/embassy-stm32/src/dma/gpdma.rs b/embassy-stm32/src/dma/gpdma.rs
index 97cc200d7..b811da1fb 100644
--- a/embassy-stm32/src/dma/gpdma.rs
+++ b/embassy-stm32/src/dma/gpdma.rs
@@ -53,10 +53,10 @@ impl State {
 static STATE: State = State::new();
 
 /// safety: must be called only once
-pub(crate) unsafe fn init(irq_priority: Priority) {
+pub(crate) unsafe fn init(cs: critical_section::CriticalSection, irq_priority: Priority) {
     foreach_interrupt! {
         ($peri:ident, gpdma, $block:ident, $signal_name:ident, $irq:ident) => {
-            crate::interrupt::typelevel::$irq::set_priority(irq_priority);
+            crate::interrupt::typelevel::$irq::set_priority_with_cs(cs, irq_priority);
             crate::interrupt::typelevel::$irq::enable();
         };
     }
diff --git a/embassy-stm32/src/dma/mod.rs b/embassy-stm32/src/dma/mod.rs
index 4f1a58ae2..29fced8fc 100644
--- a/embassy-stm32/src/dma/mod.rs
+++ b/embassy-stm32/src/dma/mod.rs
@@ -56,16 +56,17 @@ pub(crate) fn slice_ptr_parts_mut<T>(slice: *mut [T]) -> (usize, usize) {
 
 // safety: must be called only once at startup
 pub(crate) unsafe fn init(
+    cs: critical_section::CriticalSection,
     #[cfg(bdma)] bdma_priority: Priority,
     #[cfg(dma)] dma_priority: Priority,
     #[cfg(gpdma)] gpdma_priority: Priority,
 ) {
     #[cfg(bdma)]
-    bdma::init(bdma_priority);
+    bdma::init(cs, bdma_priority);
     #[cfg(dma)]
-    dma::init(dma_priority);
+    dma::init(cs, dma_priority);
     #[cfg(gpdma)]
-    gpdma::init(gpdma_priority);
+    gpdma::init(cs, gpdma_priority);
     #[cfg(dmamux)]
-    dmamux::init();
+    dmamux::init(cs);
 }
diff --git a/embassy-stm32/src/exti.rs b/embassy-stm32/src/exti.rs
index 62f321709..538791a51 100644
--- a/embassy-stm32/src/exti.rs
+++ b/embassy-stm32/src/exti.rs
@@ -367,7 +367,7 @@ macro_rules! enable_irq {
 }
 
 /// safety: must be called only once
-pub(crate) unsafe fn init() {
+pub(crate) unsafe fn init(_cs: critical_section::CriticalSection) {
     use crate::interrupt::typelevel::Interrupt;
 
     foreach_exti_irq!(enable_irq);
diff --git a/embassy-stm32/src/gpio.rs b/embassy-stm32/src/gpio.rs
index 37fedf8e1..e1702b008 100644
--- a/embassy-stm32/src/gpio.rs
+++ b/embassy-stm32/src/gpio.rs
@@ -1,6 +1,7 @@
 #![macro_use]
 use core::convert::Infallible;
 
+use critical_section::CriticalSection;
 use embassy_hal_internal::{impl_peripheral, into_ref, PeripheralRef};
 
 use crate::pac::gpio::{self, vals};
@@ -757,9 +758,9 @@ foreach_pin!(
     };
 );
 
-pub(crate) unsafe fn init() {
+pub(crate) unsafe fn init(_cs: CriticalSection) {
     #[cfg(afio)]
-    <crate::peripherals::AFIO as crate::rcc::sealed::RccPeripheral>::enable_and_reset();
+    <crate::peripherals::AFIO as crate::rcc::sealed::RccPeripheral>::enable_and_reset_with_cs(_cs);
 
     crate::_generated::init_gpio();
 }
diff --git a/embassy-stm32/src/lib.rs b/embassy-stm32/src/lib.rs
index b93e5ee87..372246f87 100644
--- a/embassy-stm32/src/lib.rs
+++ b/embassy-stm32/src/lib.rs
@@ -155,81 +155,82 @@ impl Default for Config {
 
 /// Initialize embassy.
 pub fn init(config: Config) -> Peripherals {
-    let p = Peripherals::take();
-
-    #[cfg(dbgmcu)]
-    if config.enable_debug_during_sleep {
-        crate::pac::DBGMCU.cr().modify(|cr| {
-            #[cfg(any(dbgmcu_f0, dbgmcu_c0, dbgmcu_g0, dbgmcu_u5, dbgmcu_wba))]
-            {
-                cr.set_dbg_stop(true);
-                cr.set_dbg_standby(true);
-            }
-            #[cfg(any(
-                dbgmcu_f1, dbgmcu_f2, dbgmcu_f3, dbgmcu_f4, dbgmcu_f7, dbgmcu_g4, dbgmcu_f7, dbgmcu_l0, dbgmcu_l1,
-                dbgmcu_l4, dbgmcu_wb, dbgmcu_wl
-            ))]
-            {
-                cr.set_dbg_sleep(true);
-                cr.set_dbg_stop(true);
-                cr.set_dbg_standby(true);
-            }
-            #[cfg(dbgmcu_h7)]
-            {
-                cr.set_d1dbgcken(true);
-                cr.set_d3dbgcken(true);
-                cr.set_dbgsleep_d1(true);
-                cr.set_dbgstby_d1(true);
-                cr.set_dbgstop_d1(true);
-            }
-        });
-    }
-
-    #[cfg(not(any(stm32f1, stm32wb, stm32wl)))]
-    peripherals::SYSCFG::enable_and_reset();
-    #[cfg(not(any(stm32h5, stm32h7, stm32wb, stm32wl)))]
-    peripherals::PWR::enable_and_reset();
-    #[cfg(not(any(stm32f2, stm32f4, stm32f7, stm32l0, stm32h5, stm32h7)))]
-    peripherals::FLASH::enable_and_reset();
-
-    unsafe {
-        #[cfg(feature = "_split-pins-enabled")]
-        crate::pac::SYSCFG.pmcr().modify(|pmcr| {
-            #[cfg(feature = "split-pa0")]
-            pmcr.set_pa0so(true);
-            #[cfg(feature = "split-pa1")]
-            pmcr.set_pa1so(true);
-            #[cfg(feature = "split-pc2")]
-            pmcr.set_pc2so(true);
-            #[cfg(feature = "split-pc3")]
-            pmcr.set_pc3so(true);
-        });
-
-        gpio::init();
-        dma::init(
-            #[cfg(bdma)]
-            config.bdma_interrupt_priority,
-            #[cfg(dma)]
-            config.dma_interrupt_priority,
-            #[cfg(gpdma)]
-            config.gpdma_interrupt_priority,
-        );
-        #[cfg(feature = "exti")]
-        exti::init();
-
-        rcc::init(config.rcc);
+    critical_section::with(|cs| {
+        let p = Peripherals::take_with_cs(cs);
+
+        #[cfg(dbgmcu)]
+        if config.enable_debug_during_sleep {
+            crate::pac::DBGMCU.cr().modify(|cr| {
+                #[cfg(any(dbgmcu_f0, dbgmcu_c0, dbgmcu_g0, dbgmcu_u5, dbgmcu_wba))]
+                {
+                    cr.set_dbg_stop(true);
+                    cr.set_dbg_standby(true);
+                }
+                #[cfg(any(
+                    dbgmcu_f1, dbgmcu_f2, dbgmcu_f3, dbgmcu_f4, dbgmcu_f7, dbgmcu_g4, dbgmcu_f7, dbgmcu_l0, dbgmcu_l1,
+                    dbgmcu_l4, dbgmcu_wb, dbgmcu_wl
+                ))]
+                {
+                    cr.set_dbg_sleep(true);
+                    cr.set_dbg_stop(true);
+                    cr.set_dbg_standby(true);
+                }
+                #[cfg(dbgmcu_h7)]
+                {
+                    cr.set_d1dbgcken(true);
+                    cr.set_d3dbgcken(true);
+                    cr.set_dbgsleep_d1(true);
+                    cr.set_dbgstby_d1(true);
+                    cr.set_dbgstop_d1(true);
+                }
+            });
+        }
 
-        // must be after rcc init
-        #[cfg(feature = "_time-driver")]
-        time_driver::init();
+        #[cfg(not(any(stm32f1, stm32wb, stm32wl)))]
+        peripherals::SYSCFG::enable_and_reset_with_cs(cs);
+        #[cfg(not(any(stm32h5, stm32h7, stm32wb, stm32wl)))]
+        peripherals::PWR::enable_and_reset_with_cs(cs);
+        #[cfg(not(any(stm32f2, stm32f4, stm32f7, stm32l0, stm32h5, stm32h7)))]
+        peripherals::FLASH::enable_and_reset_with_cs(cs);
+
+        unsafe {
+            #[cfg(feature = "_split-pins-enabled")]
+            crate::pac::SYSCFG.pmcr().modify(|pmcr| {
+                #[cfg(feature = "split-pa0")]
+                pmcr.set_pa0so(true);
+                #[cfg(feature = "split-pa1")]
+                pmcr.set_pa1so(true);
+                #[cfg(feature = "split-pc2")]
+                pmcr.set_pc2so(true);
+                #[cfg(feature = "split-pc3")]
+                pmcr.set_pc3so(true);
+            });
 
-        #[cfg(feature = "low-power")]
-        while !crate::rcc::low_power_ready() {
-            critical_section::with(|cs| {
+            gpio::init(cs);
+            dma::init(
+                cs,
+                #[cfg(bdma)]
+                config.bdma_interrupt_priority,
+                #[cfg(dma)]
+                config.dma_interrupt_priority,
+                #[cfg(gpdma)]
+                config.gpdma_interrupt_priority,
+            );
+            #[cfg(feature = "exti")]
+            exti::init(cs);
+
+            rcc::init(config.rcc);
+
+            // must be after rcc init
+            #[cfg(feature = "_time-driver")]
+            time_driver::init(cs);
+
+            #[cfg(feature = "low-power")]
+            while !crate::rcc::low_power_ready() {
                 crate::rcc::clock_refcount_sub(cs);
-            });
+            }
         }
-    }
 
-    p
+        p
+    })
 }
diff --git a/embassy-stm32/src/rcc/mod.rs b/embassy-stm32/src/rcc/mod.rs
index 0263c97aa..edbae30de 100644
--- a/embassy-stm32/src/rcc/mod.rs
+++ b/embassy-stm32/src/rcc/mod.rs
@@ -229,10 +229,19 @@ pub mod low_level {
 }
 
 pub(crate) mod sealed {
+    use critical_section::CriticalSection;
+
     pub trait RccPeripheral {
         fn frequency() -> crate::time::Hertz;
-        fn enable_and_reset();
-        fn disable();
+        fn enable_and_reset_with_cs(cs: CriticalSection);
+        fn disable_with_cs(cs: CriticalSection);
+
+        fn enable_and_reset() {
+            critical_section::with(|cs| Self::enable_and_reset_with_cs(cs))
+        }
+        fn disable() {
+            critical_section::with(|cs| Self::disable_with_cs(cs))
+        }
     }
 }
 
diff --git a/embassy-stm32/src/time_driver.rs b/embassy-stm32/src/time_driver.rs
index baea20aef..add8be831 100644
--- a/embassy-stm32/src/time_driver.rs
+++ b/embassy-stm32/src/time_driver.rs
@@ -152,45 +152,43 @@ embassy_time::time_driver_impl!(static DRIVER: RtcDriver = RtcDriver {
 });
 
 impl RtcDriver {
-    fn init(&'static self) {
+    fn init(&'static self, cs: critical_section::CriticalSection) {
         let r = T::regs_gp16();
 
-        <T as RccPeripheral>::enable_and_reset();
+        <T as RccPeripheral>::enable_and_reset_with_cs(cs);
 
         let timer_freq = T::frequency();
 
-        critical_section::with(|_| {
-            r.cr1().modify(|w| w.set_cen(false));
-            r.cnt().write(|w| w.set_cnt(0));
+        r.cr1().modify(|w| w.set_cen(false));
+        r.cnt().write(|w| w.set_cnt(0));
 
-            let psc = timer_freq.0 / TICK_HZ as u32 - 1;
-            let psc: u16 = match psc.try_into() {
-                Err(_) => panic!("psc division overflow: {}", psc),
-                Ok(n) => n,
-            };
+        let psc = timer_freq.0 / TICK_HZ as u32 - 1;
+        let psc: u16 = match psc.try_into() {
+            Err(_) => panic!("psc division overflow: {}", psc),
+            Ok(n) => n,
+        };
 
-            r.psc().write(|w| w.set_psc(psc));
-            r.arr().write(|w| w.set_arr(u16::MAX));
+        r.psc().write(|w| w.set_psc(psc));
+        r.arr().write(|w| w.set_arr(u16::MAX));
 
-            // Set URS, generate update and clear URS
-            r.cr1().modify(|w| w.set_urs(vals::Urs::COUNTERONLY));
-            r.egr().write(|w| w.set_ug(true));
-            r.cr1().modify(|w| w.set_urs(vals::Urs::ANYEVENT));
+        // Set URS, generate update and clear URS
+        r.cr1().modify(|w| w.set_urs(vals::Urs::COUNTERONLY));
+        r.egr().write(|w| w.set_ug(true));
+        r.cr1().modify(|w| w.set_urs(vals::Urs::ANYEVENT));
 
-            // Mid-way point
-            r.ccr(0).write(|w| w.set_ccr(0x8000));
+        // Mid-way point
+        r.ccr(0).write(|w| w.set_ccr(0x8000));
 
-            // Enable overflow and half-overflow interrupts
-            r.dier().write(|w| {
-                w.set_uie(true);
-                w.set_ccie(0, true);
-            });
+        // Enable overflow and half-overflow interrupts
+        r.dier().write(|w| {
+            w.set_uie(true);
+            w.set_ccie(0, true);
+        });
 
-            <T as BasicInstance>::Interrupt::unpend();
-            unsafe { <T as BasicInstance>::Interrupt::enable() };
+        <T as BasicInstance>::Interrupt::unpend();
+        unsafe { <T as BasicInstance>::Interrupt::enable() };
 
-            r.cr1().modify(|w| w.set_cen(true));
-        })
+        r.cr1().modify(|w| w.set_cen(true));
     }
 
     fn on_interrupt(&self) {
@@ -462,6 +460,6 @@ pub(crate) fn get_driver() -> &'static RtcDriver {
     &DRIVER
 }
 
-pub(crate) fn init() {
-    DRIVER.init()
+pub(crate) fn init(cs: CriticalSection) {
+    DRIVER.init(cs)
 }