Move as much logic from families to shared module as possible

6 files changed, 122 insertions, 180 deletions

diff --git a/embassy-stm32/src/flash/f3.rs b/embassy-stm32/src/flash/f3.rs
index 3da3962e8..7b339ccc5 100644
--- a/embassy-stm32/src/flash/f3.rs
+++ b/embassy-stm32/src/flash/f3.rs
@@ -41,33 +41,31 @@ pub(crate) unsafe fn blocking_write(start_address: u32, buf: &[u8; WRITE_SIZE])
     blocking_wait_ready()
 }
 
-pub(crate) unsafe fn blocking_erase(start_address: u32, end_address: u32) -> Result<(), Error> {
-    for page in (start_address..end_address).step_by(ERASE_SIZE) {
-        pac::FLASH.cr().modify(|w| {
-            w.set_per(true);
-        });
+pub(crate) unsafe fn blocking_erase_sector(sector: &FlashSector) -> Result<(), Error> {
+    pac::FLASH.cr().modify(|w| {
+        w.set_per(true);
+    });
 
-        pac::FLASH.ar().write(|w| w.set_far(page));
+    pac::FLASH.ar().write(|w| w.set_far(sector.first));
 
-        pac::FLASH.cr().modify(|w| {
-            w.set_strt(true);
-        });
+    pac::FLASH.cr().modify(|w| {
+        w.set_strt(true);
+    });
 
-        let mut ret: Result<(), Error> = blocking_wait_ready();
+    let mut ret: Result<(), Error> = blocking_wait_ready();
 
-        if !pac::FLASH.sr().read().eop() {
-            trace!("FLASH: EOP not set");
-            ret = Err(Error::Prog);
-        } else {
-            pac::FLASH.sr().write(|w| w.set_eop(true));
-        }
+    if !pac::FLASH.sr().read().eop() {
+        trace!("FLASH: EOP not set");
+        ret = Err(Error::Prog);
+    } else {
+        pac::FLASH.sr().write(|w| w.set_eop(true));
+    }
 
-        pac::FLASH.cr().modify(|w| w.set_per(false));
+    pac::FLASH.cr().modify(|w| w.set_per(false));
 
-        clear_all_err();
-        if ret.is_err() {
-            return ret;
-        }
+    clear_all_err();
+    if ret.is_err() {
+        return ret;
     }
     Ok(())
 }
diff --git a/embassy-stm32/src/flash/f4.rs b/embassy-stm32/src/flash/f4.rs
index 306359be1..cb420c69f 100644
--- a/embassy-stm32/src/flash/f4.rs
+++ b/embassy-stm32/src/flash/f4.rs
@@ -63,17 +63,8 @@ pub(crate) unsafe fn blocking_write(start_address: u32, buf: &[u8; WRITE_SIZE])
     blocking_wait_ready()
 }
 
-pub(crate) unsafe fn blocking_erase(start_address: u32, end_address: u32) -> Result<(), Error> {
-    let mut address = start_address;
-    while address < end_address {
-        let sector = get_sector(address);
-        erase_sector(sector.index)?;
-        address += sector.size;
-    }
-    Ok(())
-}
-
-unsafe fn erase_sector(sector: u8) -> Result<(), Error> {
+pub(crate) unsafe fn blocking_erase_sector(sector: &FlashSector) -> Result<(), Error> {
+    let sector = sector.index;
     let bank = sector / SECOND_BANK_SECTOR_OFFSET as u8;
     let snb = (bank << 4) + (sector % SECOND_BANK_SECTOR_OFFSET as u8);
 
@@ -132,7 +123,7 @@ unsafe fn blocking_wait_ready() -> Result<(), Error> {
 }
 
 pub(crate) fn get_sector(address: u32) -> FlashSector {
-    get_sector_inner(address, is_dual_bank(), FLASH_SIZE)
+    get_sector_inner(address, is_dual_bank(), FLASH_SIZE as u32)
 }
 
 fn get_sector_inner(address: u32, dual_bank: bool, flash_size: u32) -> FlashSector {
diff --git a/embassy-stm32/src/flash/f7.rs b/embassy-stm32/src/flash/f7.rs
index 588fc7a59..eba7df467 100644
--- a/embassy-stm32/src/flash/f7.rs
+++ b/embassy-stm32/src/flash/f7.rs
@@ -45,20 +45,10 @@ pub(crate) unsafe fn blocking_write(start_address: u32, buf: &[u8; WRITE_SIZE])
     blocking_wait_ready()
 }
 
-pub(crate) unsafe fn blocking_erase(start_address: u32, end_address: u32) -> Result<(), Error> {
-    let mut address = start_address;
-    while address < end_address {
-        let sector = get_sector(address);
-        erase_sector(sector.index)?;
-        address += sector.size;
-    }
-    Ok(())
-}
-
-unsafe fn erase_sector(sector: u8) -> Result<(), Error> {
+pub(crate) unsafe fn blocking_erase_sector(sector: &FlashSector) -> Result<(), Error> {
     pac::FLASH.cr().modify(|w| {
         w.set_ser(true);
-        w.set_snb(sector)
+        w.set_snb(sector.index)
     });
 
     pac::FLASH.cr().modify(|w| {
diff --git a/embassy-stm32/src/flash/h7.rs b/embassy-stm32/src/flash/h7.rs
index 732af8539..28999999d 100644
--- a/embassy-stm32/src/flash/h7.rs
+++ b/embassy-stm32/src/flash/h7.rs
@@ -78,20 +78,9 @@ pub(crate) unsafe fn blocking_write(start_address: u32, buf: &[u8; WRITE_SIZE])
     res.unwrap()
 }
 
-pub(crate) unsafe fn blocking_erase(start_address: u32, end_address: u32) -> Result<(), Error> {
-    let start_sector = (start_address - super::FLASH_BASE as u32) / ERASE_SIZE as u32;
-    let end_sector = (end_address - super::FLASH_BASE as u32) / ERASE_SIZE as u32;
-    for sector in start_sector..end_sector {
-        let bank = if sector >= 8 { 1 } else { 0 };
-        let ret = erase_sector(pac::FLASH.bank(bank), (sector % 8) as u8);
-        if ret.is_err() {
-            return ret;
-        }
-    }
-    Ok(())
-}
-
-unsafe fn erase_sector(bank: pac::flash::Bank, sector: u8) -> Result<(), Error> {
+pub(crate) unsafe fn blocking_erase_sector(sector: &FlashSector) -> Result<(), Error> {
+    let bank = pac::FLASH::bank(if sector.index >= 8 { 1 } else { 0 });
+    let sector = sector.index % 8;
     bank.cr().modify(|w| {
         w.set_ser(true);
         w.set_snb(sector)
diff --git a/embassy-stm32/src/flash/l.rs b/embassy-stm32/src/flash/l.rs
index 56a21885e..c8d060f0a 100644
--- a/embassy-stm32/src/flash/l.rs
+++ b/embassy-stm32/src/flash/l.rs
@@ -62,55 +62,50 @@ pub(crate) unsafe fn blocking_write(start_address: u32, buf: &[u8; WRITE_SIZE])
     blocking_wait_ready()
 }
 
-pub(crate) unsafe fn blocking_erase(start_address: u32, end_address: u32) -> Result<(), Error> {
-    for page in (start_address..end_address).step_by(ERASE_SIZE) {
-        #[cfg(any(flash_l0, flash_l1))]
-        {
-            pac::FLASH.pecr().modify(|w| {
-                w.set_erase(true);
-                w.set_prog(true);
-            });
-
-            write_volatile(page as *mut u32, 0xFFFFFFFF);
-        }
+pub(crate) unsafe fn blocking_erase_sector(sector: &FlashSector) -> Result<(), Error> {
+    #[cfg(any(flash_l0, flash_l1))]
+    {
+        pac::FLASH.pecr().modify(|w| {
+            w.set_erase(true);
+            w.set_prog(true);
+        });
 
-        #[cfg(any(flash_wl, flash_wb, flash_l4))]
-        {
-            let idx = (page - super::FLASH_BASE as u32) / ERASE_SIZE as u32;
-
-            #[cfg(flash_l4)]
-            let (idx, bank) = if idx > 255 { (idx - 256, true) } else { (idx, false) };
-
-            pac::FLASH.cr().modify(|w| {
-                w.set_per(true);
-                w.set_pnb(idx as u8);
-                #[cfg(any(flash_wl, flash_wb))]
-                w.set_strt(true);
-                #[cfg(any(flash_l4))]
-                w.set_start(true);
-                #[cfg(any(flash_l4))]
-                w.set_bker(bank);
-            });
-        }
+        write_volatile(sector.start as *mut u32, 0xFFFFFFFF);
+    }
+
+    #[cfg(any(flash_wl, flash_wb, flash_l4))]
+    {
+        let idx = (sector.start - super::FLASH_BASE as u32) / ERASE_SIZE as u32;
+
+        #[cfg(flash_l4)]
+        let (idx, bank) = if idx > 255 { (idx - 256, true) } else { (idx, false) };
+
+        pac::FLASH.cr().modify(|w| {
+            w.set_per(true);
+            w.set_pnb(idx as u8);
+            #[cfg(any(flash_wl, flash_wb))]
+            w.set_strt(true);
+            #[cfg(any(flash_l4))]
+            w.set_start(true);
+            #[cfg(any(flash_l4))]
+            w.set_bker(bank);
+        });
+    }
 
-        let ret: Result<(), Error> = blocking_wait_ready();
+    let ret: Result<(), Error> = blocking_wait_ready();
 
-        #[cfg(any(flash_wl, flash_wb, flash_l4))]
-        pac::FLASH.cr().modify(|w| w.set_per(false));
+    #[cfg(any(flash_wl, flash_wb, flash_l4))]
+    pac::FLASH.cr().modify(|w| w.set_per(false));
 
-        #[cfg(any(flash_l0, flash_l1))]
-        pac::FLASH.pecr().modify(|w| {
-            w.set_erase(false);
-            w.set_prog(false);
-        });
+    #[cfg(any(flash_l0, flash_l1))]
+    pac::FLASH.pecr().modify(|w| {
+        w.set_erase(false);
+        w.set_prog(false);
+    });
 
-        clear_all_err();
-        if ret.is_err() {
-            return ret;
-        }
-    }
+    clear_all_err();
 
-    Ok(())
+    ret
 }
 
 pub(crate) unsafe fn clear_all_err() {
diff --git a/embassy-stm32/src/flash/mod.rs b/embassy-stm32/src/flash/mod.rs
index 29cf3cc51..89fdabd4d 100644
--- a/embassy-stm32/src/flash/mod.rs
+++ b/embassy-stm32/src/flash/mod.rs
@@ -1,3 +1,4 @@
+use embassy_hal_common::drop::OnDrop;
 use embassy_hal_common::{into_ref, PeripheralRef};
 use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
 use embassy_sync::mutex::{Mutex, MutexGuard};
@@ -14,7 +15,6 @@ use crate::Peripheral;
 #[cfg_attr(flash_f7, path = "f7.rs")]
 #[cfg_attr(flash_h7, path = "h7.rs")]
 mod family;
-
 pub struct Flash<'d> {
     inner: PeripheralRef<'d, FLASH>,
 }
@@ -49,73 +49,93 @@ impl<'d> Flash<'d> {
     }
 
     pub fn blocking_read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Error> {
-        if offset as usize + bytes.len() > FLASH_SIZE {
+        Self::blocking_read_inner(FLASH_BASE as u32 + offset, bytes)
+    }
+
+    fn blocking_read_inner(start_address: u32, bytes: &mut [u8]) -> Result<(), Error> {
+        assert!(start_address >= FLASH_BASE as u32);
+        if start_address as usize + bytes.len() > FLASH_BASE + FLASH_SIZE {
             return Err(Error::Size);
         }
 
-        let first_address = FLASH_BASE as u32 + offset;
-        let flash_data = unsafe { core::slice::from_raw_parts(first_address as *const u8, bytes.len()) };
+        let flash_data = unsafe { core::slice::from_raw_parts(start_address as *const u8, bytes.len()) };
         bytes.copy_from_slice(flash_data);
         Ok(())
     }
 
     pub fn blocking_write(&mut self, offset: u32, buf: &[u8]) -> Result<(), Error> {
-        if offset as usize + buf.len() > FLASH_SIZE {
+        let start_address = FLASH_BASE as u32 + offset;
+
+        // No need to take lock here as we only have one mut flash reference.
+
+        unsafe { Flash::blocking_write_inner(start_address, buf) }
+    }
+
+    unsafe fn blocking_write_inner(start_address: u32, buf: &[u8]) -> Result<(), Error> {
+        assert!(start_address >= FLASH_BASE as u32);
+        if start_address as usize + buf.len() > FLASH_BASE + FLASH_SIZE {
             return Err(Error::Size);
         }
-        if offset as usize % WRITE_SIZE != 0 || buf.len() as usize % WRITE_SIZE != 0 {
+        if (start_address as usize - FLASH_BASE) % WRITE_SIZE != 0 || buf.len() as usize % WRITE_SIZE != 0 {
             return Err(Error::Unaligned);
         }
 
-        let start_address = FLASH_BASE as u32 + offset;
         trace!("Writing {} bytes at 0x{:x}", buf.len(), start_address);
 
-        // No need to take lock here as we only have one mut flash reference.
+        family::clear_all_err();
+        family::unlock();
+        family::begin_write();
 
-        unsafe {
-            family::clear_all_err();
-            family::unlock();
-            let res = Flash::blocking_write_all(start_address, buf);
+        let _ = OnDrop::new(|| {
+            family::end_write();
             family::lock();
-            res
-        }
-    }
+        });
 
-    unsafe fn blocking_write_all(start_address: u32, buf: &[u8]) -> Result<(), Error> {
-        family::begin_write();
         let mut address = start_address;
         for chunk in buf.chunks(WRITE_SIZE) {
-            let res = unsafe { family::blocking_write(address, chunk.try_into().unwrap()) };
-            if res.is_err() {
-                family::end_write();
-                return res;
-            }
+            unsafe { family::blocking_write(address, chunk.try_into().unwrap())? };
             address += WRITE_SIZE as u32;
         }
-
-        family::end_write();
         Ok(())
     }
 
     pub fn blocking_erase(&mut self, from: u32, to: u32) -> Result<(), Error> {
-        if to < from || to as usize > FLASH_SIZE {
-            return Err(Error::Size);
-        }
-
         let start_address = FLASH_BASE as u32 + from;
         let end_address = FLASH_BASE as u32 + to;
-        if !is_eraseable_range(start_address, end_address) {
+
+        unsafe { Flash::blocking_erase_inner(start_address, end_address) }
+    }
+
+    unsafe fn blocking_erase_inner(start_address: u32, end_address: u32) -> Result<(), Error> {
+        // Test if the address range is aligned at sector base addresses
+        let mut address = start_address;
+        while address < end_address {
+            let sector = family::get_sector(address);
+            if sector.start != address {
+                return Err(Error::Unaligned);
+            }
+            address += sector.size;
+        }
+        if address != end_address {
             return Err(Error::Unaligned);
         }
+
         trace!("Erasing from 0x{:x} to 0x{:x}", start_address, end_address);
 
-        unsafe {
-            family::clear_all_err();
-            family::unlock();
-            let res = family::blocking_erase(start_address, end_address);
+        family::clear_all_err();
+        family::unlock();
+
+        let _ = OnDrop::new(|| {
             family::lock();
-            res
+        });
+
+        let mut address = start_address;
+        while address < end_address {
+            let sector = family::get_sector(address);
+            family::blocking_erase_sector(&sector)?;
+            address += sector.size;
         }
+        Ok(())
     }
 }
 
@@ -135,74 +155,33 @@ pub trait FlashRegion {
     const SETTINGS: FlashRegionSettings;
 
     fn blocking_read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Error> {
-        if offset as usize + bytes.len() > Self::SETTINGS.size {
-            return Err(Error::Size);
-        }
-
-        let first_address = Self::SETTINGS.base as u32 + offset;
-        let flash_data = unsafe { core::slice::from_raw_parts(first_address as *const u8, bytes.len()) };
-        bytes.copy_from_slice(flash_data);
-        Ok(())
+        Flash::blocking_read_inner(Self::SETTINGS.base as u32 + offset, bytes)
     }
 
     fn blocking_write(&mut self, offset: u32, buf: &[u8]) -> Result<(), Error> {
-        if offset as usize + buf.len() > Self::SETTINGS.size {
-            return Err(Error::Size);
-        }
-        if offset as usize % Self::SETTINGS.write_size != 0 || buf.len() as usize % Self::SETTINGS.write_size != 0 {
-            return Err(Error::Unaligned);
-        }
-
         let start_address = Self::SETTINGS.base as u32 + offset;
-        trace!("Writing {} bytes from 0x{:x}", buf.len(), start_address);
 
         // Protect agains simultaneous write/erase to multiple regions.
         let _guard = take_lock_spin();
 
         unsafe {
             family::clear_all_err();
-            family::unlock();
-            let res = Flash::blocking_write_all(start_address, buf);
-            family::lock();
-            res
+            Flash::blocking_write_inner(start_address, buf)
         }
     }
 
     fn blocking_erase(&mut self, from: u32, to: u32) -> Result<(), Error> {
-        if to < from || to as usize > Self::SETTINGS.size {
-            return Err(Error::Size);
-        }
-        if (from as usize % Self::SETTINGS.erase_size) != 0 || (to as usize % Self::SETTINGS.erase_size) != 0 {
-            return Err(Error::Unaligned);
-        }
-
         let start_address = Self::SETTINGS.base as u32 + from;
         let end_address = Self::SETTINGS.base as u32 + to;
-        trace!("Erasing from 0x{:x} to 0x{:x}", start_address, end_address);
 
         // Protect agains simultaneous write/erase to multiple regions.
         let _guard = take_lock_spin();
 
         unsafe {
             family::clear_all_err();
-            family::unlock();
-            let res = family::blocking_erase(start_address, end_address);
-            family::lock();
-            res
-        }
-    }
-}
-
-fn is_eraseable_range(start_address: u32, end_address: u32) -> bool {
-    let mut address = start_address;
-    while address < end_address {
-        let sector = family::get_sector(address);
-        if sector.start != address {
-            return false;
+            Flash::blocking_erase_inner(start_address, end_address)
         }
-        address += sector.size;
     }
-    address == end_address
 }
 
 fn take_lock_spin() -> MutexGuard<'static, CriticalSectionRawMutex, ()> {