2 files changed, 84 insertions, 3 deletions

diff --git a/embassy-embedded-hal/Cargo.toml b/embassy-embedded-hal/Cargo.toml
index 19d512585..81cece686 100644
--- a/embassy-embedded-hal/Cargo.toml
+++ b/embassy-embedded-hal/Cargo.toml
@@ -17,8 +17,11 @@ std = []
 nightly = ["embedded-hal-async", "embedded-storage-async"]
 
 [dependencies]
+embassy-futures = { version = "0.1.0", path = "../embassy-futures" }
 embassy-sync = { version = "0.2.0", path = "../embassy-sync" }
-embedded-hal-02 = { package = "embedded-hal", version = "0.2.6", features = ["unproven"] }
+embedded-hal-02 = { package = "embedded-hal", version = "0.2.6", features = [
+    "unproven",
+] }
 embedded-hal-1 = { package = "embedded-hal", version = "=1.0.0-alpha.10" }
 embedded-hal-async = { version = "=0.2.0-alpha.1", optional = true }
 embedded-storage = "0.3.0"
@@ -26,3 +29,6 @@ embedded-storage-async = { version = "0.4.0", optional = true }
 nb = "1.0.0"
 
 defmt = { version = "0.3", optional = true }
+
+[dev-dependencies]
+futures-test = "0.3.17"
diff --git a/embassy-embedded-hal/src/adapter.rs b/embassy-embedded-hal/src/adapter.rs
index 171ff6c9f..169aad5e3 100644
--- a/embassy-embedded-hal/src/adapter.rs
+++ b/embassy-embedded-hal/src/adapter.rs
@@ -1,5 +1,6 @@
 //! Adapters between embedded-hal traits.
 
+use embassy_futures::yield_now;
 use embedded_hal_02::{blocking, serial};
 
 /// Wrapper that implements async traits using blocking implementations.
@@ -150,11 +151,18 @@ where
     const ERASE_SIZE: usize = <T as NorFlash>::ERASE_SIZE;
 
     async fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
-        self.wrapped.write(offset, data)
+        self.wrapped.write(offset, data)?;
+        yield_now().await;
+        Ok(())
     }
 
     async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-        self.wrapped.erase(from, to)
+        for from in (from..to).step_by(T::ERASE_SIZE) {
+            let to = core::cmp::min(from + T::ERASE_SIZE as u32, to);
+            self.wrapped.erase(from, to)?;
+            yield_now().await;
+        }
+        Ok(())
     }
 }
 
@@ -171,3 +179,70 @@ where
         self.wrapped.capacity()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    extern crate std;
+
+    #[derive(Default)]
+    struct FakeFlash(Vec<(u32, u32)>);
+
+    impl embedded_storage::nor_flash::ErrorType for FakeFlash {
+        type Error = std::convert::Infallible;
+    }
+
+    impl embedded_storage::nor_flash::ReadNorFlash for FakeFlash {
+        const READ_SIZE: usize = 1;
+
+        fn read(&mut self, _offset: u32, _bytes: &mut [u8]) -> Result<(), Self::Error> {
+            unimplemented!()
+        }
+
+        fn capacity(&self) -> usize {
+            unimplemented!()
+        }
+    }
+
+    impl embedded_storage::nor_flash::NorFlash for FakeFlash {
+        const WRITE_SIZE: usize = 4;
+        const ERASE_SIZE: usize = 128;
+
+        fn write(&mut self, _offset: u32, _bytes: &[u8]) -> Result<(), Self::Error> {
+            unimplemented!()
+        }
+
+        fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+            self.0.push((from, to));
+            Ok(())
+        }
+    }
+
+    #[futures_test::test]
+    async fn can_erase() {
+        let fake = FakeFlash::default();
+        let mut yielding = BlockingAsync::new(fake);
+
+        yielding.erase(0, 256).await.unwrap();
+
+        let fake = yielding.wrapped;
+        assert_eq!(2, fake.0.len());
+        assert_eq!((0, 128), fake.0[0]);
+        assert_eq!((128, 256), fake.0[1]);
+    }
+
+    #[futures_test::test]
+    async fn can_erase_wrong_erase_size() {
+        let fake = FakeFlash::default();
+        let mut yielding = BlockingAsync::new(fake);
+
+        yielding.erase(0, 257).await.unwrap();
+
+        let fake = yielding.wrapped;
+        assert_eq!(3, fake.0.len());
+        assert_eq!((0, 128), fake.0[0]);
+        assert_eq!((128, 256), fake.0[1]);
+        assert_eq!((256, 257), fake.0[2]);
+    }
+}