rp235x flash support.

The 2350 doesn't have a boot2 like the 2040, but it does have the
concept of a xip setup function that could be customized. By default the
bootrom searches for the attached flash chip and provides an xip setup
func at the base of the bootram. That bootram is not executable, so it
still needs to be copied to ram like boot2 would be.

Currently does not use inline assembly.

Also switch to picotool, as elf2uf2 has not been patched to support the
2350.

4 files changed, 68 insertions, 89 deletions

diff --git a/embassy-rp/src/flash.rs b/embassy-rp/src/flash.rs
index 5f7922f8e..dab99b4e2 100644
--- a/embassy-rp/src/flash.rs
+++ b/embassy-rp/src/flash.rs
@@ -17,9 +17,13 @@ use crate::peripherals::FLASH;
 /// Flash base address.
 pub const FLASH_BASE: *const u32 = 0x10000000 as _;
 
+/// Address for xip setup function set up by the 235x bootrom.
+#[cfg(feature = "_rp235x")]
+pub const BOOTROM_BASE: *const u32 = 0x400e0000 as _;
+
 /// If running from RAM, we might have no boot2. Use bootrom `flash_enter_cmd_xip` instead.
 // TODO: when run-from-ram is set, completely skip the "pause cores and jumpp to RAM" dance.
-pub const USE_BOOT2: bool = !cfg!(feature = "run-from-ram");
+pub const USE_BOOT2: bool = !cfg!(feature = "run-from-ram") | cfg!(feature = "_rp235x");
 
 // **NOTE**:
 //
@@ -97,7 +101,10 @@ impl<'a, 'd, T: Instance, const FLASH_SIZE: usize> Drop for BackgroundRead<'a, '
         // Errata RP2040-E8: Perform an uncached read to make sure there's not a transfer in
         // flight that might effect an address written to start a new transfer.  This stalls
         // until after any transfer is complete, so the address will not change anymore.
+        #[cfg(feature = "rp2040")]
         const XIP_NOCACHE_NOALLOC_BASE: *const u32 = 0x13000000 as *const _;
+        #[cfg(feature = "_rp235x")]
+        const XIP_NOCACHE_NOALLOC_BASE: *const u32 = 0x14000000 as *const _;
         unsafe {
             core::ptr::read_volatile(XIP_NOCACHE_NOALLOC_BASE);
         }
@@ -225,12 +232,14 @@ impl<'d, T: Instance, M: Mode, const FLASH_SIZE: usize> Flash<'d, T, M, FLASH_SI
     }
 
     /// Read SPI flash unique ID
+    #[cfg(feature = "rp2040")]
     pub fn blocking_unique_id(&mut self, uid: &mut [u8]) -> Result<(), Error> {
         unsafe { in_ram(|| ram_helpers::flash_unique_id(uid))? };
         Ok(())
     }
 
     /// Read SPI flash JEDEC ID
+    #[cfg(feature = "rp2040")]
     pub fn blocking_jedec_id(&mut self) -> Result<u32, Error> {
         let mut jedec = None;
         unsafe {
@@ -301,7 +310,10 @@ impl<'d, T: Instance, const FLASH_SIZE: usize> Flash<'d, T, Async, FLASH_SIZE> {
         // Use the XIP AUX bus port, rather than the FIFO register access (e.x.
         // pac::XIP_CTRL.stream_fifo().as_ptr()) to avoid DMA stalling on
         // general XIP access.
+        #[cfg(feature = "rp2040")]
         const XIP_AUX_BASE: *const u32 = 0x50400000 as *const _;
+        #[cfg(feature = "_rp235x")]
+        const XIP_AUX_BASE: *const u32 = 0x50500000 as *const _;
         let transfer = unsafe {
             crate::dma::read(
                 self.dma.as_mut().unwrap(),
@@ -510,19 +522,14 @@ mod ram_helpers {
     ///
     /// `addr` and `len` parameters must be valid and are not checked.
     pub unsafe fn flash_range_erase(addr: u32, len: u32) {
-        #[cfg(feature = "rp2040")]
         let mut boot2 = [0u32; 256 / 4];
-        let ptrs = {
+        let ptrs = if USE_BOOT2 {
             #[cfg(feature = "rp2040")]
-            {
-                if USE_BOOT2 {
-                    rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
-                    flash_function_pointers_with_boot2(true, false, &boot2)
-                } else {
-                    flash_function_pointers(true, false)
-                }
-            }
+            rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
             #[cfg(feature = "_rp235x")]
+            core::ptr::copy_nonoverlapping(BOOTROM_BASE as *const u8, boot2.as_mut_ptr() as *mut u8, 256);
+            flash_function_pointers_with_boot2(true, false, &boot2)
+        } else {
             flash_function_pointers(true, false)
         };
 
@@ -548,19 +555,14 @@ mod ram_helpers {
     ///
     /// `addr` and `len` parameters must be valid and are not checked.
     pub unsafe fn flash_range_erase_and_program(addr: u32, data: &[u8]) {
-        #[cfg(feature = "rp2040")]
         let mut boot2 = [0u32; 256 / 4];
-        let ptrs = {
+        let ptrs = if USE_BOOT2 {
             #[cfg(feature = "rp2040")]
-            {
-                if USE_BOOT2 {
-                    rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
-                    flash_function_pointers_with_boot2(true, true, &boot2)
-                } else {
-                    flash_function_pointers(true, true)
-                }
-            }
+            rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
             #[cfg(feature = "_rp235x")]
+            core::ptr::copy_nonoverlapping(BOOTROM_BASE as *const u8, (boot2).as_mut_ptr() as *mut u8, 256);
+            flash_function_pointers_with_boot2(true, true, &boot2)
+        } else {
             flash_function_pointers(true, true)
         };
 
@@ -591,19 +593,14 @@ mod ram_helpers {
     ///
     /// `addr` and `len` parameters must be valid and are not checked.
     pub unsafe fn flash_range_program(addr: u32, data: &[u8]) {
-        #[cfg(feature = "rp2040")]
         let mut boot2 = [0u32; 256 / 4];
-        let ptrs = {
+        let ptrs = if USE_BOOT2 {
             #[cfg(feature = "rp2040")]
-            {
-                if USE_BOOT2 {
-                    rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
-                    flash_function_pointers_with_boot2(false, true, &boot2)
-                } else {
-                    flash_function_pointers(false, true)
-                }
-            }
+            rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
             #[cfg(feature = "_rp235x")]
+            core::ptr::copy_nonoverlapping(BOOTROM_BASE as *const u8, boot2.as_mut_ptr() as *mut u8, 256);
+            flash_function_pointers_with_boot2(false, true, &boot2)
+        } else {
             flash_function_pointers(false, true)
         };
 
@@ -630,16 +627,7 @@ mod ram_helpers {
     #[link_section = ".data.ram_func"]
     #[cfg(feature = "rp2040")]
     unsafe fn write_flash_inner(addr: u32, len: u32, data: Option<&[u8]>, ptrs: *const FlashFunctionPointers) {
-        /*
-         Should be equivalent to:
-            rom_data::connect_internal_flash();
-            rom_data::flash_exit_xip();
-            rom_data::flash_range_erase(addr, len, 1 << 31, 0); // if selected
-            rom_data::flash_range_program(addr, data as *const _, len); // if selected
-            rom_data::flash_flush_cache();
-            rom_data::flash_enter_cmd_xip();
-        */
-        #[cfg(target_arch = "arm")]
+        //#[cfg(target_arch = "arm")]
         core::arch::asm!(
             "mov r8, r0",
             "mov r9, r2",
@@ -691,11 +679,30 @@ mod ram_helpers {
         );
     }
 
+    /// # Safety
+    ///
+    /// Nothing must access flash while this is running.
+    /// Usually this means:
+    ///   - interrupts must be disabled
+    ///   - 2nd core must be running code from RAM or ROM with interrupts disabled
+    ///   - DMA must not access flash memory
+    /// Length of data must be a multiple of 4096
+    /// addr must be aligned to 4096
     #[inline(never)]
     #[link_section = ".data.ram_func"]
     #[cfg(feature = "_rp235x")]
-    unsafe fn write_flash_inner(_addr: u32, _len: u32, _data: Option<&[u8]>, _ptrs: *const FlashFunctionPointers) {
-        todo!();
+    unsafe fn write_flash_inner(addr: u32, len: u32, data: Option<&[u8]>, ptrs: *const FlashFunctionPointers) {
+        let data = data.map(|d| d.as_ptr()).unwrap_or(core::ptr::null());
+        ((*ptrs).connect_internal_flash)();
+        ((*ptrs).flash_exit_xip)();
+        if (*ptrs).flash_range_erase.is_some() {
+            ((*ptrs).flash_range_erase.unwrap())(addr, len as usize, 1 << 31, 0);
+        }
+        if (*ptrs).flash_range_program.is_some() {
+            ((*ptrs).flash_range_program.unwrap())(addr, data as *const _, len as usize);
+        }
+        ((*ptrs).flash_flush_cache)();
+        ((*ptrs).flash_enter_cmd_xip)();
     }
 
     #[repr(C)]
@@ -731,20 +738,13 @@ mod ram_helpers {
     ///   - DMA must not access flash memory
     ///
     /// Credit: taken from `rp2040-flash` (also licensed Apache+MIT)
+    #[cfg(feature = "rp2040")]
     pub unsafe fn flash_unique_id(out: &mut [u8]) {
-        #[cfg(feature = "rp2040")]
         let mut boot2 = [0u32; 256 / 4];
-        let ptrs = {
-            #[cfg(feature = "rp2040")]
-            {
-                if USE_BOOT2 {
-                    rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
-                    flash_function_pointers_with_boot2(false, false, &boot2)
-                } else {
-                    flash_function_pointers(false, false)
-                }
-            }
-            #[cfg(feature = "_rp235x")]
+        let ptrs = if USE_BOOT2 {
+            rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
+            flash_function_pointers_with_boot2(false, false, &boot2)
+        } else {
             flash_function_pointers(false, false)
         };
 
@@ -768,20 +768,13 @@ mod ram_helpers {
     ///   - DMA must not access flash memory
     ///
     /// Credit: taken from `rp2040-flash` (also licensed Apache+MIT)
+    #[cfg(feature = "rp2040")]
     pub unsafe fn flash_jedec_id() -> u32 {
-        #[cfg(feature = "rp2040")]
         let mut boot2 = [0u32; 256 / 4];
-        let ptrs = {
-            #[cfg(feature = "rp2040")]
-            {
-                if USE_BOOT2 {
-                    rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
-                    flash_function_pointers_with_boot2(false, false, &boot2)
-                } else {
-                    flash_function_pointers(false, false)
-                }
-            }
-            #[cfg(feature = "_rp235x")]
+        let ptrs = if USE_BOOT2 {
+            rom_data::memcpy44(&mut boot2 as *mut _, FLASH_BASE, 256);
+            flash_function_pointers_with_boot2(false, false, &boot2)
+        } else {
             flash_function_pointers(false, false)
         };
 
@@ -792,6 +785,7 @@ mod ram_helpers {
         u32::from_be_bytes(id)
     }
 
+    #[cfg(feature = "rp2040")]
     unsafe fn read_flash(cmd_addr: &[u8], dummy_len: u32, out: &mut [u8], ptrs: *const FlashFunctionPointers) {
         read_flash_inner(
             FlashCommand {
@@ -932,13 +926,6 @@ mod ram_helpers {
             clobber_abi("C"),
         );
     }
-
-    #[inline(never)]
-    #[link_section = ".data.ram_func"]
-    #[cfg(feature = "_rp235x")]
-    unsafe fn read_flash_inner(_cmd: FlashCommand, _ptrs: *const FlashFunctionPointers) {
-        todo!();
-    }
 }
 
 /// Make sure to uphold the contract points with rp2040-flash.
diff --git a/embassy-rp/src/lib.rs b/embassy-rp/src/lib.rs
index c2c57eaa0..f8fcfe52b 100644
--- a/embassy-rp/src/lib.rs
+++ b/embassy-rp/src/lib.rs
@@ -402,7 +402,7 @@ embassy_hal_internal::peripherals! {
     BOOTSEL,
 }
 
-#[cfg(not(feature = "boot2-none"))]
+#[cfg(all(not(feature = "boot2-none"), feature = "rp2040"))]
 macro_rules! select_bootloader {
     ( $( $feature:literal => $loader:ident, )+ default => $default:ident ) => {
         $(
@@ -419,7 +419,7 @@ macro_rules! select_bootloader {
     }
 }
 
-#[cfg(not(feature = "boot2-none"))]
+#[cfg(all(not(feature = "boot2-none"), feature = "rp2040"))]
 select_bootloader! {
     "boot2-at25sf128a" => BOOT_LOADER_AT25SF128A,
     "boot2-gd25q64cs" => BOOT_LOADER_GD25Q64CS,
diff --git a/examples/rp23/.cargo/config.toml b/examples/rp23/.cargo/config.toml
index f77e004dc..9a92b1ce2 100644
--- a/examples/rp23/.cargo/config.toml
+++ b/examples/rp23/.cargo/config.toml
@@ -1,6 +1,7 @@
 [target.'cfg(all(target_arch = "arm", target_os = "none"))']
 #runner = "probe-rs run --chip RP2040"
-runner = "elf2uf2-rs -d"
+#runner = "elf2uf2-rs -d"
+runner = "picotool load -u -v -x -t elf"
 
 [build]
 target = "thumbv8m.main-none-eabihf"
diff --git a/examples/rp23/src/bin/flash.rs b/examples/rp23/src/bin/flash.rs
index 811561f26..84011e394 100644
--- a/examples/rp23/src/bin/flash.rs
+++ b/examples/rp23/src/bin/flash.rs
@@ -21,7 +21,7 @@ pub static IMAGE_DEF: ImageDef = ImageDef::secure_exe();
 pub static PICOTOOL_ENTRIES: [embassy_rp::binary_info::EntryAddr; 4] = [
     embassy_rp::binary_info::rp_program_name!(c"example"),
     embassy_rp::binary_info::rp_cargo_version!(),
-    embassy_rp::binary_info::rp_program_description!(c"Blinky"),
+    embassy_rp::binary_info::rp_program_description!(c"Flash"),
     embassy_rp::binary_info::rp_program_build_attribute!(),
 ];
 
@@ -41,22 +41,13 @@ async fn main(_spawner: Spawner) {
 
     let mut flash = embassy_rp::flash::Flash::<_, Async, FLASH_SIZE>::new(p.FLASH, p.DMA_CH0);
 
-    // Get JEDEC id
-    let jedec = flash.blocking_jedec_id().unwrap();
-    info!("jedec id: 0x{:x}", jedec);
-
-    // Get unique id
-    let mut uid = [0; 8];
-    flash.blocking_unique_id(&mut uid).unwrap();
-    info!("unique id: {:?}", uid);
-
     erase_write_sector(&mut flash, 0x00);
 
     multiwrite_bytes(&mut flash, ERASE_SIZE as u32);
 
     background_read(&mut flash, (ERASE_SIZE * 2) as u32).await;
 
-    loop {}
+    info!("Flash Works!");
 }
 
 fn multiwrite_bytes(flash: &mut embassy_rp::flash::Flash<'_, FLASH, Async, FLASH_SIZE>, offset: u32) {
@@ -82,7 +73,7 @@ fn multiwrite_bytes(flash: &mut embassy_rp::flash::Flash<'_, FLASH, Async, FLASH
 
     defmt::unwrap!(flash.blocking_read(ADDR_OFFSET + offset, &mut read_buf));
     info!("Contents after write starts with {=[u8]}", read_buf[0..4]);
-    if &read_buf[0..4] != &[0x01, 0x02, 0x03, 0x04] {
+    if read_buf[0..4] != [0x01, 0x02, 0x03, 0x04] {
         defmt::panic!("unexpected");
     }
 }