From 68d053c66eb14b29ea16efba89cf01cb879eca3e Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Mon, 1 Jun 2020 21:24:48 +1000
Subject: [PATCH] stm32/modmachine: Allow changing AHB and APB bus frequencies
 on STM32WB.

For now SYSCLK cannot be changed and must remain at 64MHz.
---
 ports/stm32/modmachine.c |  7 ++++++-
 ports/stm32/powerctrl.c  | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/ports/stm32/modmachine.c b/ports/stm32/modmachine.c
index 0795d2ccf..6fd2488a5 100644
--- a/ports/stm32/modmachine.c
+++ b/ports/stm32/modmachine.c
@@ -307,7 +307,7 @@ STATIC mp_obj_t machine_freq(size_t n_args, const mp_obj_t *args) {
         return mp_obj_new_tuple(MP_ARRAY_SIZE(tuple), tuple);
     } else {
         // set
-        #if defined(STM32F0) || defined(STM32L0) || defined(STM32L4) || defined(STM32WB)
+        #if defined(STM32F0) || defined(STM32L0) || defined(STM32L4)
         mp_raise_NotImplementedError(MP_ERROR_TEXT("machine.freq set not supported yet"));
         #else
         mp_int_t sysclk = mp_obj_get_int(args[0]);
@@ -317,8 +317,13 @@ STATIC mp_obj_t machine_freq(size_t n_args, const mp_obj_t *args) {
             ahb /= 2;
         }
         #endif
+        #if defined(STM32WB)
+        mp_int_t apb1 = ahb;
+        mp_int_t apb2 = ahb;
+        #else
         mp_int_t apb1 = ahb / 4;
         mp_int_t apb2 = ahb / 2;
+        #endif
         if (n_args > 1) {
             ahb = mp_obj_get_int(args[1]);
             if (n_args > 2) {
diff --git a/ports/stm32/powerctrl.c b/ports/stm32/powerctrl.c
index e11de8b02..946185fba 100644
--- a/ports/stm32/powerctrl.c
+++ b/ports/stm32/powerctrl.c
@@ -201,7 +201,7 @@ int powerctrl_rcc_clock_config_pll(RCC_ClkInitTypeDef *rcc_init, uint32_t sysclk
 
 #endif
 
-#if !defined(STM32F0) && !defined(STM32L0) && !defined(STM32L4) && !defined(STM32WB)
+#if !defined(STM32F0) && !defined(STM32L0) && !defined(STM32L4)
 
 STATIC uint32_t calc_ahb_div(uint32_t wanted_div) {
     #if defined(STM32H7)
@@ -293,6 +293,8 @@ STATIC uint32_t calc_apb2_div(uint32_t wanted_div) {
     #endif
 }
 
+#if defined(STM32F4) || defined(STM32F7) || defined(STM32H7)
+
 int powerctrl_set_sysclk(uint32_t sysclk, uint32_t ahb, uint32_t apb1, uint32_t apb2) {
     // Return straightaway if the clocks are already at the desired frequency
     if (sysclk == HAL_RCC_GetSysClockFreq()
@@ -455,8 +457,36 @@ set_clk:
     return 0;
 }
 
+#elif defined(STM32WB)
+
+int powerctrl_set_sysclk(uint32_t sysclk, uint32_t ahb, uint32_t apb1, uint32_t apb2) {
+    // For now it's not supported to change SYSCLK (only bus dividers).
+    if (sysclk != HAL_RCC_GetSysClockFreq()) {
+        return -MP_EINVAL;
+    }
+
+    // Return straightaway if the clocks are already at the desired frequency.
+    if (ahb == HAL_RCC_GetHCLKFreq()
+        && apb1 == HAL_RCC_GetPCLK1Freq()
+        && apb2 == HAL_RCC_GetPCLK2Freq()) {
+        return 0;
+    }
+
+    // Calculate and configure the bus clock dividers.
+    uint32_t cfgr = RCC->CFGR;
+    cfgr &= ~(7 << RCC_CFGR_PPRE2_Pos | 7 << RCC_CFGR_PPRE1_Pos | 0xf << RCC_CFGR_HPRE_Pos);
+    cfgr |= calc_ahb_div(sysclk / ahb);
+    cfgr |= calc_apb1_div(ahb / apb1);
+    cfgr |= calc_apb2_div(ahb / apb2) << (RCC_CFGR_PPRE2_Pos - RCC_CFGR_PPRE1_Pos);
+    RCC->CFGR = cfgr;
+
+    return 0;
+}
+
 #endif
 
+#endif // !defined(STM32F0) && !defined(STM32L0) && !defined(STM32L4)
+
 void powerctrl_enter_stop_mode(void) {
     // Disable IRQs so that the IRQ that wakes the device from stop mode is not
     // executed until after the clocks are reconfigured