5 files changed, 313 insertions, 105 deletions

diff --git a/embassy-stm32/src/ipcc.rs b/embassy-stm32/src/ipcc.rs
index a24cba9f0..e100ca5cc 100644
--- a/embassy-stm32/src/ipcc.rs
+++ b/embassy-stm32/src/ipcc.rs
@@ -265,63 +265,9 @@ pub(crate) mod sealed {
 }
 
 fn _configure_pwr() {
-    // TODO: move this to RCC
-
-    let pwr = crate::pac::PWR;
+    // TODO: move the rest of this to rcc
     let rcc = crate::pac::RCC;
 
-    rcc.cfgr().modify(|w| w.set_stopwuck(true));
-
-    pwr.cr1().modify(|w| w.set_dbp(true));
-    pwr.cr1().modify(|w| w.set_dbp(true));
-
-    // configure LSE
-    rcc.bdcr().modify(|w| w.set_lseon(true));
-
-    // select system clock source = PLL
-    // set PLL coefficients
-    // m: 2,
-    // n: 12,
-    // r: 3,
-    // q: 4,
-    // p: 3,
-    let src_bits = 0b11;
-    let pllp = (3 - 1) & 0b11111;
-    let pllq = (4 - 1) & 0b111;
-    let pllr = (3 - 1) & 0b111;
-    let plln = 12 & 0b1111111;
-    let pllm = (2 - 1) & 0b111;
-    rcc.pllcfgr().modify(|w| {
-        w.set_pllsrc(src_bits);
-        w.set_pllm(pllm);
-        w.set_plln(plln);
-        w.set_pllr(pllr);
-        w.set_pllp(pllp);
-        w.set_pllpen(true);
-        w.set_pllq(pllq);
-        w.set_pllqen(true);
-    });
-    // enable PLL
-    rcc.cr().modify(|w| w.set_pllon(true));
-    rcc.cr().write(|w| w.set_hsion(false));
-    // while !rcc.cr().read().pllrdy() {}
-
-    // configure SYSCLK mux to use PLL clocl
-    rcc.cfgr().modify(|w| w.set_sw(0b11));
-
-    // configure CPU1 & CPU2 dividers
-    rcc.cfgr().modify(|w| w.set_hpre(0)); // not divided
-    rcc.extcfgr().modify(|w| {
-        w.set_c2hpre(0b1000); // div2
-        w.set_shdhpre(0); // not divided
-    });
-
-    // apply APB1 / APB2 values
-    rcc.cfgr().modify(|w| {
-        w.set_ppre1(0b000); // not divided
-        w.set_ppre2(0b000); // not divided
-    });
-
     // TODO: required
     // set RF wake-up clock = LSE
     rcc.csr().modify(|w| w.set_rfwkpsel(0b01));
diff --git a/embassy-stm32/src/rcc/mod.rs b/embassy-stm32/src/rcc/mod.rs
index 886fc0b93..4ae65d3e6 100644
--- a/embassy-stm32/src/rcc/mod.rs
+++ b/embassy-stm32/src/rcc/mod.rs
@@ -78,6 +78,14 @@ pub struct Clocks {
 /// The existence of this value indicates that the clock configuration can no longer be changed
 static mut CLOCK_FREQS: MaybeUninit<Clocks> = MaybeUninit::uninit();
 
+#[cfg(stm32wb)]
+/// RCC initialization function
+pub(crate) unsafe fn init(config: Config) {
+    set_freqs(compute_clocks(&config));
+
+    configure_clocks(&config);
+}
+
 /// Sets the clock frequencies
 ///
 /// Safety: Sets a mutable global.
diff --git a/embassy-stm32/src/rcc/wb.rs b/embassy-stm32/src/rcc/wb.rs
index e6123821a..4322b950a 100644
--- a/embassy-stm32/src/rcc/wb.rs
+++ b/embassy-stm32/src/rcc/wb.rs
@@ -1,6 +1,5 @@
-use crate::pac::RCC;
-use crate::rcc::{set_freqs, Clocks};
-use crate::time::Hertz;
+use crate::rcc::Clocks;
+use crate::time::{khz, mhz, Hertz};
 
 /// Most of clock setup is copied from stm32l0xx-hal, and adopted to the generated PAC,
 /// and with the addition of the init function to configure a system clock.
@@ -13,11 +12,94 @@ pub const HSI_FREQ: Hertz = Hertz(16_000_000);
 /// LSI speed
 pub const LSI_FREQ: Hertz = Hertz(32_000);
 
-/// System clock mux source
 #[derive(Clone, Copy)]
-pub enum ClockSrc {
-    HSE(Hertz),
-    HSI16,
+pub enum HsePrescaler {
+    NotDivided,
+    Div2,
+}
+
+impl From<HsePrescaler> for bool {
+    fn from(value: HsePrescaler) -> Self {
+        match value {
+            HsePrescaler::NotDivided => false,
+            HsePrescaler::Div2 => true,
+        }
+    }
+}
+
+pub struct Hse {
+    pub prediv: HsePrescaler,
+
+    pub frequency: Hertz,
+}
+
+/// System clock mux source
+#[derive(Clone, Copy, PartialEq)]
+pub enum Sysclk {
+    /// MSI selected as sysclk
+    MSI,
+    /// HSI selected as sysclk
+    HSI,
+    /// HSE selected as sysclk
+    HSE,
+    /// PLL selected as sysclk
+    Pll,
+}
+
+impl From<Sysclk> for u8 {
+    fn from(value: Sysclk) -> Self {
+        match value {
+            Sysclk::MSI => 0b00,
+            Sysclk::HSI => 0b01,
+            Sysclk::HSE => 0b10,
+            Sysclk::Pll => 0b11,
+        }
+    }
+}
+
+#[derive(Clone, Copy, PartialEq)]
+pub enum PllSource {
+    Hsi,
+    Msi,
+    Hse,
+}
+
+impl From<PllSource> for u8 {
+    fn from(value: PllSource) -> Self {
+        match value {
+            PllSource::Msi => 0b01,
+            PllSource::Hsi => 0b10,
+            PllSource::Hse => 0b11,
+        }
+    }
+}
+
+pub enum Pll48Source {
+    PllSai,
+    Pll,
+    Msi,
+    Hsi48,
+}
+
+pub struct PllMux {
+    /// Source clock selection.
+    pub source: PllSource,
+
+    /// PLL pre-divider (DIVM). Must be between 1 and 63.
+    pub prediv: u8,
+}
+
+pub struct Pll {
+    /// PLL multiplication factor. Must be between 4 and 512.
+    pub mul: u16,
+
+    /// PLL P division factor. If None, PLL P output is disabled. Must be between 1 and 128.
+    /// On PLL1, it must be even (in particular, it cannot be 1.)
+    pub divp: Option<u16>,
+    /// PLL Q division factor. If None, PLL Q output is disabled. Must be between 1 and 128.
+    pub divq: Option<u16>,
+    /// PLL R division factor. If None, PLL R output is disabled. Must be between 1 and 128.
+    pub divr: Option<u16>,
 }
 
 /// AHB prescaler
@@ -84,86 +166,250 @@ impl Into<u8> for AHBPrescaler {
 
 /// Clocks configutation
 pub struct Config {
-    pub mux: ClockSrc,
-    pub ahb_pre: AHBPrescaler,
+    pub hse: Option<Hse>,
+    pub lse: Option<Hertz>,
+    pub sys: Sysclk,
+    pub mux: Option<PllMux>,
+    pub pll48: Option<Pll48Source>,
+
+    pub pll: Option<Pll>,
+    pub pllsai: Option<Pll>,
+
+    pub ahb1_pre: AHBPrescaler,
+    pub ahb2_pre: AHBPrescaler,
+    pub ahb3_pre: AHBPrescaler,
     pub apb1_pre: APBPrescaler,
     pub apb2_pre: APBPrescaler,
 }
 
+pub const WPAN_DEFAULT: Config = Config {
+    hse: Some(Hse {
+        frequency: mhz(32),
+        prediv: HsePrescaler::NotDivided,
+    }),
+    lse: Some(khz(32)),
+    sys: Sysclk::Pll,
+    mux: Some(PllMux {
+        source: PllSource::Hse,
+        prediv: 2,
+    }),
+    pll48: None,
+
+    pll: Some(Pll {
+        mul: 12,
+        divp: Some(3),
+        divq: Some(4),
+        divr: Some(3),
+    }),
+    pllsai: None,
+
+    ahb1_pre: AHBPrescaler::NotDivided,
+    ahb2_pre: AHBPrescaler::Div2,
+    ahb3_pre: AHBPrescaler::NotDivided,
+    apb1_pre: APBPrescaler::NotDivided,
+    apb2_pre: APBPrescaler::NotDivided,
+};
+
 impl Default for Config {
     #[inline]
     fn default() -> Config {
         Config {
-            mux: ClockSrc::HSI16,
-            ahb_pre: AHBPrescaler::NotDivided,
+            hse: None,
+            lse: None,
+            sys: Sysclk::HSI,
+            mux: None,
+            pll48: None,
+            pll: None,
+            pllsai: None,
+
+            ahb1_pre: AHBPrescaler::NotDivided,
+            ahb2_pre: AHBPrescaler::NotDivided,
+            ahb3_pre: AHBPrescaler::NotDivided,
             apb1_pre: APBPrescaler::NotDivided,
             apb2_pre: APBPrescaler::NotDivided,
         }
     }
 }
 
-pub(crate) unsafe fn init(config: Config) {
-    let (sys_clk, sw) = match config.mux {
-        ClockSrc::HSI16 => {
-            // Enable HSI16
-            RCC.cr().write(|w| w.set_hsion(true));
-            while !RCC.cr().read().hsirdy() {}
+pub(crate) fn compute_clocks(config: &Config) -> Clocks {
+    let hse_clk = config.hse.as_ref().map(|hse| match hse.prediv {
+        HsePrescaler::NotDivided => hse.frequency,
+        HsePrescaler::Div2 => hse.frequency / 2u32,
+    });
 
-            (HSI_FREQ.0, 0x01)
+    let mux_clk = config.mux.as_ref().map(|pll_mux| {
+        (match pll_mux.source {
+            PllSource::Hse => hse_clk.unwrap(),
+            PllSource::Hsi => HSI_FREQ,
+            _ => unreachable!(),
+        } / pll_mux.prediv)
+    });
+
+    let (pll_r, _pll_q, _pll_p) = match &config.pll {
+        Some(pll) => {
+            let pll_vco = mux_clk.unwrap() * pll.mul as u32;
+
+            (
+                pll.divr.map(|divr| pll_vco / divr),
+                pll.divq.map(|divq| pll_vco / divq),
+                pll.divp.map(|divp| pll_vco / divp),
+            )
         }
-        ClockSrc::HSE(freq) => {
-            // Enable HSE
-            RCC.cr().write(|w| w.set_hseon(true));
-            while !RCC.cr().read().hserdy() {}
+        None => (None, None, None),
+    };
 
-            (freq.0, 0x02)
+    let sys_clk = match config.sys {
+        Sysclk::HSE => hse_clk.unwrap(),
+        Sysclk::HSI => HSI_FREQ,
+        Sysclk::Pll => pll_r.unwrap(),
+        _ => unreachable!(),
+    };
+
+    let ahb1_clk = match config.ahb1_pre {
+        AHBPrescaler::NotDivided => sys_clk,
+        pre => {
+            let pre: u8 = pre.into();
+            let pre = 1u32 << (pre as u32 - 7);
+            sys_clk / pre
         }
     };
 
-    RCC.cfgr().modify(|w| {
-        w.set_sw(sw.into());
-        w.set_hpre(config.ahb_pre.into());
-        w.set_ppre1(config.apb1_pre.into());
-        w.set_ppre2(config.apb2_pre.into());
-    });
+    let ahb2_clk = match config.ahb2_pre {
+        AHBPrescaler::NotDivided => sys_clk,
+        pre => {
+            let pre: u8 = pre.into();
+            let pre = 1u32 << (pre as u32 - 7);
+            sys_clk / pre
+        }
+    };
 
-    let ahb_freq: u32 = match config.ahb_pre {
+    let ahb3_clk = match config.ahb3_pre {
         AHBPrescaler::NotDivided => sys_clk,
         pre => {
             let pre: u8 = pre.into();
-            let pre = 1 << (pre as u32 - 7);
+            let pre = 1u32 << (pre as u32 - 7);
             sys_clk / pre
         }
     };
 
-    let (apb1_freq, apb1_tim_freq) = match config.apb1_pre {
-        APBPrescaler::NotDivided => (ahb_freq, ahb_freq),
+    let (apb1_clk, apb1_tim_clk) = match config.apb1_pre {
+        APBPrescaler::NotDivided => (ahb1_clk, ahb1_clk),
         pre => {
             let pre: u8 = pre.into();
             let pre: u8 = 1 << (pre - 3);
-            let freq = ahb_freq / pre as u32;
-            (freq, freq * 2)
+            let freq = ahb1_clk / pre as u32;
+            (freq, freq * 2u32)
         }
     };
 
-    let (apb2_freq, apb2_tim_freq) = match config.apb2_pre {
-        APBPrescaler::NotDivided => (ahb_freq, ahb_freq),
+    let (apb2_clk, apb2_tim_clk) = match config.apb2_pre {
+        APBPrescaler::NotDivided => (ahb1_clk, ahb1_clk),
         pre => {
             let pre: u8 = pre.into();
             let pre: u8 = 1 << (pre - 3);
-            let freq = ahb_freq / pre as u32;
-            (freq, freq * 2)
+            let freq = ahb1_clk / pre as u32;
+            (freq, freq * 2u32)
+        }
+    };
+
+    Clocks {
+        sys: sys_clk,
+        ahb1: ahb1_clk,
+        ahb2: ahb2_clk,
+        ahb3: ahb3_clk,
+        apb1: apb1_clk,
+        apb2: apb2_clk,
+        apb1_tim: apb1_tim_clk,
+        apb2_tim: apb2_tim_clk,
+    }
+}
+
+pub(crate) fn configure_clocks(config: &Config) {
+    let pwr = crate::pac::PWR;
+    let rcc = crate::pac::RCC;
+
+    let needs_hsi = if let Some(pll_mux) = &config.mux {
+        pll_mux.source == PllSource::Hsi
+    } else {
+        false
+    };
+
+    if needs_hsi || config.sys == Sysclk::HSI {
+        rcc.cr().modify(|w| {
+            w.set_hsion(true);
+        });
+
+        while !rcc.cr().read().hsirdy() {}
+    }
+
+    match &config.lse {
+        Some(_) => {
+            rcc.cfgr().modify(|w| w.set_stopwuck(true));
+
+            pwr.cr1().modify(|w| w.set_dbp(true));
+            pwr.cr1().modify(|w| w.set_dbp(true));
+
+            rcc.bdcr().modify(|w| w.set_lseon(true));
+        }
+        _ => {}
+    }
+
+    match &config.hse {
+        Some(hse) => {
+            rcc.cr().modify(|w| {
+                w.set_hsepre(hse.prediv.into());
+                w.set_hseon(true);
+            });
+
+            while !rcc.cr().read().hserdy() {}
+        }
+        _ => {}
+    }
+
+    match &config.mux {
+        Some(pll_mux) => {
+            rcc.pllcfgr().modify(|w| {
+                w.set_pllm(pll_mux.prediv);
+                w.set_pllsrc(pll_mux.source.into());
+            });
         }
+        _ => {}
     };
 
-    set_freqs(Clocks {
-        sys: Hertz(sys_clk),
-        ahb1: Hertz(ahb_freq),
-        ahb2: Hertz(ahb_freq),
-        ahb3: Hertz(ahb_freq),
-        apb1: Hertz(apb1_freq),
-        apb2: Hertz(apb2_freq),
-        apb1_tim: Hertz(apb1_tim_freq),
-        apb2_tim: Hertz(apb2_tim_freq),
+    match &config.pll {
+        Some(pll) => {
+            rcc.pllcfgr().modify(|w| {
+                w.set_plln(pll.mul as u8);
+                pll.divp.map(|divp| {
+                    w.set_pllpen(true);
+                    w.set_pllp((divp - 1) as u8)
+                });
+                pll.divq.map(|divq| {
+                    w.set_pllqen(true);
+                    w.set_pllq((divq - 1) as u8)
+                });
+                pll.divr.map(|divr| {
+                    // w.set_pllren(true);
+                    w.set_pllr((divr - 1) as u8);
+                });
+            });
+
+            rcc.cr().modify(|w| w.set_pllon(true));
+
+            while !rcc.cr().read().pllrdy() {}
+        }
+        _ => {}
+    }
+
+    rcc.cfgr().modify(|w| {
+        w.set_sw(config.sys.into());
+        w.set_hpre(config.ahb1_pre.into());
+        w.set_ppre1(config.apb1_pre.into());
+        w.set_ppre2(config.apb2_pre.into());
+    });
+
+    rcc.extcfgr().modify(|w| {
+        w.set_c2hpre(config.ahb2_pre.into());
+        w.set_shdhpre(config.ahb3_pre.into());
     });
 }
diff --git a/tests/stm32/src/bin/wpan_ble.rs b/tests/stm32/src/bin/wpan_ble.rs
index 3ad8aca4e..452da77a4 100644
--- a/tests/stm32/src/bin/wpan_ble.rs
+++ b/tests/stm32/src/bin/wpan_ble.rs
@@ -12,6 +12,7 @@ use common::*;
 use embassy_executor::Spawner;
 use embassy_stm32::bind_interrupts;
 use embassy_stm32::ipcc::{Config, ReceiveInterruptHandler, TransmitInterruptHandler};
+use embassy_stm32::rcc::WPAN_DEFAULT;
 use embassy_stm32_wpan::hci::host::uart::UartHci;
 use embassy_stm32_wpan::hci::host::{AdvertisingFilterPolicy, EncryptionKey, HostHci, OwnAddressType};
 use embassy_stm32_wpan::hci::types::AdvertisingType;
@@ -40,7 +41,10 @@ async fn run_mm_queue(memory_manager: mm::MemoryManager) {
 
 #[embassy_executor::main]
 async fn main(spawner: Spawner) {
-    let p = embassy_stm32::init(config());
+    let mut config = config();
+    config.rcc = WPAN_DEFAULT;
+
+    let p = embassy_stm32::init(config);
     info!("Hello World!");
 
     let config = Config::default();
diff --git a/tests/stm32/src/bin/wpan_mac.rs b/tests/stm32/src/bin/wpan_mac.rs
index b04a19ee9..7eab2fd38 100644
--- a/tests/stm32/src/bin/wpan_mac.rs
+++ b/tests/stm32/src/bin/wpan_mac.rs
@@ -10,6 +10,7 @@ use common::*;
 use embassy_executor::Spawner;
 use embassy_stm32::bind_interrupts;
 use embassy_stm32::ipcc::{Config, ReceiveInterruptHandler, TransmitInterruptHandler};
+use embassy_stm32::rcc::WPAN_DEFAULT;
 use embassy_stm32_wpan::mac::commands::{AssociateRequest, GetRequest, ResetRequest, SetRequest};
 use embassy_stm32_wpan::mac::event::MacEvent;
 use embassy_stm32_wpan::mac::typedefs::{
@@ -31,7 +32,10 @@ async fn run_mm_queue(memory_manager: mm::MemoryManager) {
 
 #[embassy_executor::main]
 async fn main(spawner: Spawner) {
-    let p = embassy_stm32::init(config());
+    let mut config = config();
+    config.rcc = WPAN_DEFAULT;
+
+    let p = embassy_stm32::init(config);
     info!("Hello World!");
 
     let config = Config::default();