Code cleanup

firmware/acid-firmware/src/flash.rs

@@ -0,0 +1,112 @@
+use core::fmt::Write;
+
+use alloc::{string::String, vec::Vec};
+use embassy_embedded_hal::{adapter::BlockingAsync, flash::partition::Partition};
+use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, mutex::Mutex};
+use esp_bootloader_esp_idf::partitions::PartitionTable;
+use esp_storage::FlashStorage;
+use indoc::writedoc;
+use log::info;
+use rmk::storage::async_flash_wrapper;
+use static_cell::StaticCell;
+
+use crate::PSRAM_ALLOCATOR;
+
+pub struct Partitions {
+    pub rmk: Partition<'static, CriticalSectionRawMutex, BlockingAsync<FlashStorage<'static>>>,
+    pub acid: Partition<'static, CriticalSectionRawMutex, BlockingAsync<FlashStorage<'static>>>,
+}
+
+/// Initialize the flash
+pub fn initialize(flash_peripheral: esp_hal::peripherals::FLASH<'static>) -> Partitions {
+    static PARTITION_TABLE_BUFFER: StaticCell<Vec<u8, &'static esp_alloc::EspHeap>> =
+        StaticCell::new();
+    let partition_table_buffer = PARTITION_TABLE_BUFFER.init_with(|| {
+        let mut buffer = Vec::<u8, _>::new_in(&PSRAM_ALLOCATOR);
+        buffer.resize(
+            esp_bootloader_esp_idf::partitions::PARTITION_TABLE_MAX_LEN,
+            0_u8,
+        );
+        buffer
+    });
+
+    static FLASH: StaticCell<(
+        Mutex<CriticalSectionRawMutex, BlockingAsync<FlashStorage>>,
+        PartitionTable<'static>,
+    )> = StaticCell::new();
+    let (flash, partition_table) = FLASH.init_with(|| {
+        let mut flash = FlashStorage::new(flash_peripheral)
+            // Flash memory may not be written to while another core is executing from it.
+            // By default, `FlashStorage` is configured to abort the operation and log an error message.
+            // However, it can also be configured to auto-park the other core, such that writing to
+            // flash succeeds.
+            // Alternatively, XiP from PSRAM could be used along with the `multicore_ignore` strategy,
+            // to avoid having to park the other core, which could result in better performance.
+            // Invalid configuration would then present itself as freezing/UB.
+            .multicore_auto_park();
+        let partition_table = {
+            esp_bootloader_esp_idf::partitions::read_partition_table(
+                &mut flash,
+                partition_table_buffer,
+            )
+            .expect("Failed to read the partition table.")
+        };
+
+        (
+            Mutex::<CriticalSectionRawMutex, _>::new(async_flash_wrapper(flash)),
+            partition_table,
+        )
+    });
+
+    {
+        let mut buffer = String::new();
+
+        writeln!(buffer, "Partition table:").unwrap();
+
+        for (index, partition) in partition_table.iter().enumerate() {
+            writedoc!(
+                buffer,
+                "
+                Partition #{index} {label:?}:
+                    offset: 0x{offset:x}
+                    length: 0x{len:x}
+                    type: 0x{type:?}
+                    read only: {read_only}
+                    encrypted: {encrypted}
+                    magic: {magic}
+                ",
+                label = partition.label_as_str(),
+                offset = partition.offset(),
+                len = partition.len(),
+                type = partition.partition_type(),
+                read_only = partition.is_read_only(),
+                encrypted = partition.is_encrypted(),
+                magic = partition.magic(),
+            )
+            .unwrap();
+        }
+
+        info!("{}", buffer);
+    }
+    let flash_part_info_rmk = partition_table
+        .iter()
+        .find(|partition| partition.label_as_str() == "rmk")
+        .expect("No \"rmk\" partition found. Make sure to use the custom partition-table.csv when flashing.");
+    let flash_part_info_acid = partition_table
+        .iter()
+        .find(|partition| partition.label_as_str() == "acid")
+        .expect("No \"acid\" partition found. Make sure to use the custom partition-table.csv when flashing.");
+
+    Partitions {
+        rmk: Partition::new(
+            flash,
+            flash_part_info_rmk.offset(),
+            flash_part_info_rmk.len(),
+        ),
+        acid: Partition::new(
+            flash,
+            flash_part_info_acid.offset(),
+            flash_part_info_acid.len(),
+        ),
+    }
+}

firmware/acid-firmware/src/main.rs

@@ -17,26 +17,20 @@
 extern crate alloc;
 
 use core::cell::RefCell;
-use core::fmt::Write;
 use core::sync::atomic::{AtomicBool, Ordering};
 
 use alloc::alloc::Global;
 use alloc::boxed::Box;
 use alloc::collections::vec_deque::VecDeque;
 use alloc::format;
-use alloc::string::String;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
-use embassy_embedded_hal::adapter::BlockingAsync;
-use embassy_embedded_hal::flash::partition::Partition;
 use embassy_executor::Spawner;
 use embassy_sync::blocking_mutex;
 use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
 use embassy_sync::channel::Channel;
-use embassy_sync::mutex::Mutex;
 use embassy_time::{Duration, Timer};
 use esp_alloc::{HeapRegion, MemoryCapability};
-use esp_bootloader_esp_idf::partitions::PartitionTable;
 use esp_hal::clock::CpuClock;
 use esp_hal::dma::{BurstConfig, ExternalBurstConfig, InternalBurstConfig};
 use esp_hal::efuse::Efuse;
@@ -64,7 +58,6 @@ use esp_hal_bounce_buffers::{
 use esp_rtos::embassy::{Executor, InterruptExecutor};
 use esp_storage::FlashStorage;
 use i_slint_core::software_renderer::TargetPixel;
-use indoc::writedoc;
 use itertools::Itertools;
 use log::{error, info, warn};
 use rmk::channel::{CONTROLLER_CHANNEL, ControllerSub};
@@ -76,7 +69,6 @@ use rmk::hid::Report;
 use rmk::input_device::Runnable;
 use rmk::join_all;
 use rmk::keyboard::Keyboard;
-use rmk::storage::async_flash_wrapper;
 use rmk::types::action::{Action, KeyAction};
 use rmk::{initialize_keymap_and_storage, run_devices, run_rmk};
 use slint::platform::software_renderer::Rgb565Pixel;
@@ -99,6 +91,7 @@ mod config;
 mod crypto;
 mod db;
 mod ffi;
+mod flash;
 mod logging;
 mod matrix;
 mod peripherals;
@@ -500,90 +493,7 @@ async fn main_task(peripherals: MainPeripherals) {
     };
 
     // Initialize the flash
-    static PARTITION_TABLE_BUFFER: StaticCell<Vec<u8, &'static esp_alloc::EspHeap>> =
+    let flash_partitions = flash::initialize(peripherals.FLASH);
-        StaticCell::new();
-    let partition_table_buffer = PARTITION_TABLE_BUFFER.init_with(|| {
-        let mut buffer = Vec::<u8, _>::new_in(&PSRAM_ALLOCATOR);
-        buffer.resize(1024, 0_u8);
-        buffer
-    });
-
-    static FLASH: StaticCell<(
-        Mutex<CriticalSectionRawMutex, BlockingAsync<FlashStorage>>,
-        PartitionTable<'static>,
-    )> = StaticCell::new();
-    let (flash, partition_table) = FLASH.init_with(|| {
-        let mut flash = FlashStorage::new(peripherals.FLASH)
-            // Flash memory may not be written to while another core is executing from it.
-            // By default, `FlashStorage` is configured to abort the operation and log an error message.
-            // However, it can also be configured to auto-park the other core, such that writing to
-            // flash succeeds.
-            // Alternatively, XiP from PSRAM could be used along with the `multicore_ignore` strategy,
-            // to avoid having to park the other core, which could result in better performance.
-            // Invalid configuration would then present itself as freezing/UB.
-            .multicore_auto_park();
-        let partition_table = {
-            esp_bootloader_esp_idf::partitions::read_partition_table(
-                &mut flash,
-                partition_table_buffer,
-            )
-            .expect("Failed to read the partition table.")
-        };
-
-        (
-            Mutex::<CriticalSectionRawMutex, _>::new(async_flash_wrapper(flash)),
-            partition_table,
-        )
-    });
-
-    {
-        let mut buffer = String::new();
-
-        writeln!(buffer, "Partition table:").unwrap();
-
-        for (index, partition) in partition_table.iter().enumerate() {
-            writedoc!(
-                buffer,
-                "
-                Partition #{index} {label:?}:
-                    offset: 0x{offset:x}
-                    length: 0x{len:x}
-                    type: 0x{type:?}
-                    read only: {read_only}
-                    encrypted: {encrypted}
-                    magic: {magic}
-                ",
-                label = partition.label_as_str(),
-                offset = partition.offset(),
-                len = partition.len(),
-                type = partition.partition_type(),
-                read_only = partition.is_read_only(),
-                encrypted = partition.is_encrypted(),
-                magic = partition.magic(),
-            )
-            .unwrap();
-        }
-
-        info!("{}", buffer);
-    }
-    let flash_part_info_rmk = partition_table
-        .iter()
-        .find(|partition| partition.label_as_str() == "rmk")
-        .expect("No \"rmk\" partition found. Make sure to use the custom partition-table.csv when flashing.");
-    let flash_part_info_acid = partition_table
-        .iter()
-        .find(|partition| partition.label_as_str() == "acid")
-        .expect("No \"acid\" partition found. Make sure to use the custom partition-table.csv when flashing.");
-    let flash_part_rmk = Partition::new(
-        flash,
-        flash_part_info_rmk.offset(),
-        flash_part_info_rmk.len(),
-    );
-    let flash_part_acid = Partition::new(
-        flash,
-        flash_part_info_acid.offset(),
-        flash_part_info_acid.len(),
-    );
 
     info!("Flash memory configured!");
 
@@ -641,12 +551,12 @@ async fn main_task(peripherals: MainPeripherals) {
         start_addr: 0,
         num_sectors: {
             assert!(
-                flash_part_info_rmk.len() % FlashStorage::SECTOR_SIZE == 0,
+                flash_partitions.rmk.size() % FlashStorage::SECTOR_SIZE == 0,
                 "The size of the RMK partition must be a multiple of {} bytes. Current size: {}",
                 FlashStorage::SECTOR_SIZE,
-                flash_part_info_rmk.len()
+                flash_partitions.rmk.size()
             );
-            (flash_part_info_rmk.len() / FlashStorage::SECTOR_SIZE) as u8
+            (flash_partitions.rmk.size() / FlashStorage::SECTOR_SIZE) as u8
         },
         ..Default::default()
     };
@@ -683,7 +593,7 @@ async fn main_task(peripherals: MainPeripherals) {
     let mut positional_config = config::get_positional_config();
     let (keymap, mut storage) = initialize_keymap_and_storage(
         &mut default_keymap,
-        flash_part_rmk,
+        flash_partitions.rmk,
         &storage_config,
         &mut behavior_config,
         &mut positional_config,
@@ -772,7 +682,7 @@ async fn main_task(peripherals: MainPeripherals) {
                     quit_event_loop: Default::default(),
                     events: Arc::new(blocking_mutex::Mutex::new(RefCell::new(VecDeque::new()))),
                 };
-                spawner.must_spawn(ui::run_renderer_task(slint_backend, flash_part_acid));
+                spawner.must_spawn(ui::run_renderer_task(slint_backend, flash_partitions.acid));
             });
         },
     );