5 files changed, 71 insertions, 12 deletions

diff --git a/embassy-nrf/src/chips/nrf5340_app.rs b/embassy-nrf/src/chips/nrf5340_app.rs
index 39fb9a9bf..4575f09ff 100644
--- a/embassy-nrf/src/chips/nrf5340_app.rs
+++ b/embassy-nrf/src/chips/nrf5340_app.rs
@@ -213,6 +213,8 @@ pub mod pac {
 pub const EASY_DMA_SIZE: usize = (1 << 16) - 1;
 pub const FORCE_COPY_BUFFER_SIZE: usize = 1024;
 
+pub const FLASH_SIZE: usize = 1024 * 1024;
+
 embassy_hal_common::peripherals! {
     // USB
     USBD,
@@ -224,6 +226,9 @@ embassy_hal_common::peripherals! {
     // WDT
     WDT,
 
+    // NVMC
+    NVMC,
+
     // UARTE, TWI & SPI
     UARTETWISPI0,
     UARTETWISPI1,
diff --git a/embassy-nrf/src/chips/nrf5340_net.rs b/embassy-nrf/src/chips/nrf5340_net.rs
index b8ee7745b..54827238a 100644
--- a/embassy-nrf/src/chips/nrf5340_net.rs
+++ b/embassy-nrf/src/chips/nrf5340_net.rs
@@ -104,6 +104,8 @@ pub mod pac {
 pub const EASY_DMA_SIZE: usize = (1 << 16) - 1;
 pub const FORCE_COPY_BUFFER_SIZE: usize = 1024;
 
+pub const FLASH_SIZE: usize = 256 * 1024;
+
 embassy_hal_common::peripherals! {
     // RTC
     RTC0,
@@ -112,6 +114,9 @@ embassy_hal_common::peripherals! {
     // WDT
     WDT,
 
+    // NVMC
+    NVMC,
+
     // UARTE, TWI & SPI
     UARTETWISPI0,
     UARTETWISPI1,
diff --git a/embassy-nrf/src/chips/nrf9160.rs b/embassy-nrf/src/chips/nrf9160.rs
index a47c0a397..472ee6772 100644
--- a/embassy-nrf/src/chips/nrf9160.rs
+++ b/embassy-nrf/src/chips/nrf9160.rs
@@ -164,6 +164,8 @@ pub mod pac {
 pub const EASY_DMA_SIZE: usize = (1 << 13) - 1;
 pub const FORCE_COPY_BUFFER_SIZE: usize = 1024;
 
+pub const FLASH_SIZE: usize = 1024 * 1024;
+
 embassy_hal_common::peripherals! {
     // RTC
     RTC0,
@@ -172,6 +174,9 @@ embassy_hal_common::peripherals! {
     // WDT
     WDT,
 
+    // NVMC
+    NVMC,
+
     // UARTE, TWI & SPI
     UARTETWISPI0,
     UARTETWISPI1,
diff --git a/embassy-nrf/src/lib.rs b/embassy-nrf/src/lib.rs
index 4832e143f..1dd0e7905 100644
--- a/embassy-nrf/src/lib.rs
+++ b/embassy-nrf/src/lib.rs
@@ -80,7 +80,6 @@ pub mod gpio;
 pub mod gpiote;
 #[cfg(any(feature = "nrf52832", feature = "nrf52833", feature = "nrf52840"))]
 pub mod i2s;
-#[cfg(not(any(feature = "_nrf5340", feature = "_nrf9160")))]
 pub mod nvmc;
 #[cfg(any(
     feature = "nrf52810",
diff --git a/embassy-nrf/src/nvmc.rs b/embassy-nrf/src/nvmc.rs
index 6f66f7a78..405ea3171 100644
--- a/embassy-nrf/src/nvmc.rs
+++ b/embassy-nrf/src/nvmc.rs
@@ -10,8 +10,12 @@ use embedded_storage::nor_flash::{
 use crate::peripherals::NVMC;
 use crate::{pac, Peripheral};
 
+#[cfg(not(feature = "_nrf5340-net"))]
 /// Erase size of NVMC flash in bytes.
 pub const PAGE_SIZE: usize = 4096;
+#[cfg(feature = "_nrf5340-net")]
+/// Erase size of NVMC flash in bytes.
+pub const PAGE_SIZE: usize = 2048;
 
 /// Size of NVMC flash in bytes.
 pub const FLASH_SIZE: usize = crate::chip::FLASH_SIZE;
@@ -55,6 +59,51 @@ impl<'d> Nvmc<'d> {
         let p = Self::regs();
         while p.ready.read().ready().is_busy() {}
     }
+
+    #[cfg(not(any(feature = "_nrf9160", feature = "_nrf5340")))]
+    fn wait_ready_write(&mut self) {
+        self.wait_ready();
+    }
+
+    #[cfg(any(feature = "_nrf9160", feature = "_nrf5340"))]
+    fn wait_ready_write(&mut self) {
+        let p = Self::regs();
+        while p.readynext.read().readynext().is_busy() {}
+    }
+
+    #[cfg(not(any(feature = "_nrf9160", feature = "_nrf5340")))]
+    fn erase_page(&mut self, page_addr: u32) {
+        Self::regs().erasepage().write(|w| unsafe { w.bits(page_addr) });
+    }
+
+    #[cfg(any(feature = "_nrf9160", feature = "_nrf5340"))]
+    fn erase_page(&mut self, page_addr: u32) {
+        let first_page_word = page_addr as *mut u32;
+        unsafe {
+            first_page_word.write_volatile(0xFFFF_FFFF);
+        }
+    }
+
+    fn enable_erase(&self) {
+        #[cfg(not(any(feature = "nrf5340-app-ns", feature = "nrf9160-ns")))]
+        Self::regs().config.write(|w| w.wen().een());
+        #[cfg(any(feature = "nrf5340-app-ns", feature = "nrf9160-ns"))]
+        Self::regs().configns.write(|w| w.wen().een());
+    }
+
+    fn enable_read(&self) {
+        #[cfg(not(any(feature = "nrf5340-app-ns", feature = "nrf9160-ns")))]
+        Self::regs().config.write(|w| w.wen().ren());
+        #[cfg(any(feature = "nrf5340-app-ns", feature = "nrf9160-ns"))]
+        Self::regs().configns.write(|w| w.wen().ren());
+    }
+
+    fn enable_write(&self) {
+        #[cfg(not(any(feature = "nrf5340-app-ns", feature = "nrf9160-ns")))]
+        Self::regs().config.write(|w| w.wen().wen());
+        #[cfg(any(feature = "nrf5340-app-ns", feature = "nrf9160-ns"))]
+        Self::regs().configns.write(|w| w.wen().wen());
+    }
 }
 
 impl<'d> MultiwriteNorFlash for Nvmc<'d> {}
@@ -93,17 +142,15 @@ impl<'d> NorFlash for Nvmc<'d> {
             return Err(Error::Unaligned);
         }
 
-        let p = Self::regs();
-
-        p.config.write(|w| w.wen().een());
+        self.enable_erase();
         self.wait_ready();
 
-        for page in (from..to).step_by(PAGE_SIZE) {
-            p.erasepage().write(|w| unsafe { w.bits(page) });
+        for page_addr in (from..to).step_by(PAGE_SIZE) {
+            self.erase_page(page_addr);
             self.wait_ready();
         }
 
-        p.config.reset();
+        self.enable_read();
         self.wait_ready();
 
         Ok(())
@@ -117,9 +164,7 @@ impl<'d> NorFlash for Nvmc<'d> {
             return Err(Error::Unaligned);
         }
 
-        let p = Self::regs();
-
-        p.config.write(|w| w.wen().wen());
+        self.enable_write();
         self.wait_ready();
 
         unsafe {
@@ -129,11 +174,11 @@ impl<'d> NorFlash for Nvmc<'d> {
             for i in 0..words {
                 let w = ptr::read_unaligned(p_src.add(i));
                 ptr::write_volatile(p_dst.add(i), w);
-                self.wait_ready();
+                self.wait_ready_write();
             }
         }
 
-        p.config.reset();
+        self.enable_read();
         self.wait_ready();
 
         Ok(())