3 files changed, 55 insertions, 40 deletions

diff --git a/embassy-stm32/src/flash/f4.rs b/embassy-stm32/src/flash/f4.rs
index f6dc7e955..b8327ce4e 100644
--- a/embassy-stm32/src/flash/f4.rs
+++ b/embassy-stm32/src/flash/f4.rs
@@ -7,8 +7,6 @@ use super::{ERASE_SIZE, FLASH_BASE, FLASH_SIZE};
 use crate::flash::Error;
 use crate::pac;
 
-// Only available on some devices
-const SECOND_BANK_OFFSET: usize = FLASH_SIZE / 2;
 const SECOND_BANK_SECTOR_START: u32 = 12;
 
 unsafe fn is_dual_bank() -> bool {
@@ -68,44 +66,52 @@ pub(crate) unsafe fn blocking_write(offset: u32, buf: &[u8]) -> Result<(), Error
     ret
 }
 
-unsafe fn get_sector(addr: u32) -> u8 {
+struct FlashSector {
+    index: u8,
+    size: u32,
+}
+
+fn get_sector(addr: u32, dual_bank: bool) -> FlashSector {
     let offset = addr - FLASH_BASE as u32;
 
-    let sector = if is_dual_bank() {
-        let bank = offset / SECOND_BANK_OFFSET as u32;
-        let offset_in_bank = offset % SECOND_BANK_OFFSET as u32;
+    let bank_size = match dual_bank {
+        true => FLASH_SIZE / 2,
+        false => FLASH_SIZE,
+    } as u32;
 
-        let sector_in_bank = if offset_in_bank >= ERASE_SIZE as u32 / 2 {
-            4 + offset_in_bank / ERASE_SIZE as u32
-        } else {
-            offset_in_bank / (ERASE_SIZE as u32 / 8)
-        };
+    let bank = offset / bank_size;
+    let offset_in_bank = offset % bank_size;
 
-        if bank == 1 {
-            SECOND_BANK_SECTOR_START + sector_in_bank
-        } else {
-            sector_in_bank
-        }
+    let index_in_bank = if offset_in_bank >= ERASE_SIZE as u32 / 2 {
+        4 + offset_in_bank / ERASE_SIZE as u32
     } else {
-        if offset >= ERASE_SIZE as u32 / 2 {
-            4 + offset / ERASE_SIZE as u32
-        } else {
-            offset / (ERASE_SIZE as u32 / 8)
-        }
+        offset_in_bank / (ERASE_SIZE as u32 / 8)
+    };
+
+    // First 4 sectors are 16KB, then one 64KB, and rest are 128KB
+    let size = match index_in_bank {
+        0..=3 => 16 * 1024,
+        4 => 64 * 1024,
+        _ => 128 * 1024,
     };
 
-    sector as u8
+    let index = if bank == 1 {
+        SECOND_BANK_SECTOR_START + index_in_bank
+    } else {
+        index_in_bank
+    } as u8;
+
+    FlashSector { index, size }
 }
 
 pub(crate) unsafe fn blocking_erase(from: u32, to: u32) -> Result<(), Error> {
-    let start_sector = get_sector(from);
-    let end_sector = get_sector(to);
+    let mut addr = from;
+    let dual_bank = is_dual_bank();
 
-    for sector in start_sector..end_sector {
-        let ret = erase_sector(sector as u8);
-        if ret.is_err() {
-            return ret;
-        }
+    while addr < to {
+        let sector = get_sector(addr, dual_bank);
+        erase_sector(sector.index)?;
+        addr += sector.size;
     }
 
     Ok(())
@@ -115,6 +121,8 @@ unsafe fn erase_sector(sector: u8) -> Result<(), Error> {
     let bank = sector / SECOND_BANK_SECTOR_START as u8;
     let snb = (bank << 4) + (sector % SECOND_BANK_SECTOR_START as u8);
 
+    trace!("Erasing sector: {}", sector);
+
     pac::FLASH.cr().modify(|w| {
         w.set_ser(true);
         w.set_snb(snb)
diff --git a/embassy-stm32/src/flash/mod.rs b/embassy-stm32/src/flash/mod.rs
index 59ca59f65..2047f70e1 100644
--- a/embassy-stm32/src/flash/mod.rs
+++ b/embassy-stm32/src/flash/mod.rs
@@ -72,7 +72,7 @@ impl<'d> Flash<'d> {
         if to < from || to as usize > FLASH_END {
             return Err(Error::Size);
         }
-        if ((to - from) as usize % ERASE_SIZE) != 0 {
+        if (from as usize % ERASE_SIZE) != 0 || (to as usize % ERASE_SIZE) != 0 {
             return Err(Error::Unaligned);
         }
 
diff --git a/examples/stm32f4/src/bin/flash.rs b/examples/stm32f4/src/bin/flash.rs
index b531d6f13..13fd2c90b 100644
--- a/examples/stm32f4/src/bin/flash.rs
+++ b/examples/stm32f4/src/bin/flash.rs
@@ -4,7 +4,6 @@
 
 use defmt::{info, unwrap};
 use embassy::executor::Spawner;
-use embassy::time::{Duration, Timer};
 use embassy_stm32::flash::Flash;
 use embassy_stm32::Peripherals;
 use embedded_storage::nor_flash::{NorFlash, ReadNorFlash};
@@ -14,29 +13,37 @@ use {defmt_rtt as _, panic_probe as _};
 async fn main(_spawner: Spawner, p: Peripherals) {
     info!("Hello Flash!");
 
-    const ADDR: u32 = 0x10_0000;
+    let mut f = Flash::unlock(p.FLASH);
 
-    // wait a bit before accessing the flash
-    Timer::after(Duration::from_millis(300)).await;
+    // Sector 5
+    test_flash(&mut f, 128 * 1024, 128 * 1024);
 
-    let mut f = Flash::unlock(p.FLASH);
+    // Sectors 11..=16, across banks (128K, 16K, 16K, 16K, 16K, 64K)
+    test_flash(&mut f, (1024 - 128) * 1024, 256 * 1024);
+
+    // Sectors 23, last in bank 2
+    test_flash(&mut f, (2048 - 128) * 1024, 128 * 1024);
+}
+
+fn test_flash(f: &mut Flash, offset: u32, size: u32) {
+    info!("Testing offset: {=u32:#X}, size: {=u32:#X}", offset, size);
 
     info!("Reading...");
     let mut buf = [0u8; 32];
-    unwrap!(f.read(ADDR, &mut buf));
+    unwrap!(f.read(offset, &mut buf));
     info!("Read: {=[u8]:x}", buf);
 
     info!("Erasing...");
-    unwrap!(f.erase(ADDR, ADDR + 128 * 1024));
+    unwrap!(f.erase(offset, offset + size));
 
     info!("Reading...");
     let mut buf = [0u8; 32];
-    unwrap!(f.read(ADDR, &mut buf));
+    unwrap!(f.read(offset, &mut buf));
     info!("Read after erase: {=[u8]:x}", buf);
 
     info!("Writing...");
     unwrap!(f.write(
-        ADDR,
+        offset,
         &[
             1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
             30, 31, 32
@@ -45,7 +52,7 @@ async fn main(_spawner: Spawner, p: Peripherals) {
 
     info!("Reading...");
     let mut buf = [0u8; 32];
-    unwrap!(f.read(ADDR, &mut buf));
+    unwrap!(f.read(offset, &mut buf));
     info!("Read: {=[u8]:x}", buf);
     assert_eq!(
         &buf[..],