From 863bb6847c05d77eecda609068ee7757b208ecbd Mon Sep 17 00:00:00 2001
From: Wilfried Chauveau <wilfried.chauveau@ithinuel.me>
Date: Wed, 31 Jan 2024 05:44:55 +0000
Subject: [PATCH] i2c: Implement i2c-write-iter traits

---
 on-target-tests/Cargo.toml                    |  9 ++--
 rp2040-hal/Cargo.toml                         |  5 ++
 rp2040-hal/src/i2c/controller.rs              | 36 +++++++++++++
 rp2040-hal/src/i2c/controller/non_blocking.rs | 52 +++++++++++++++++++
 rp2040-hal/src/lib.rs                         |  2 +
 5 files changed, 101 insertions(+), 3 deletions(-)

diff --git a/on-target-tests/Cargo.toml b/on-target-tests/Cargo.toml
index 48b377b39..2b37315b8 100644
--- a/on-target-tests/Cargo.toml
+++ b/on-target-tests/Cargo.toml
@@ -51,9 +51,10 @@ defmt-test = "0.3.1"
 panic-probe = { version = "0.3", features = ["print-defmt"] }
 
 rp2040-hal = { path = "../rp2040-hal", features = [
-    "defmt",
-    "critical-section-impl",
-    "rt",
+  "critical-section-impl",
+  "defmt",
+  "rt",
+  "i2c-write-iter",
 ] }
 # Needed to set spi frequencies
 fugit = "0.3.6"
@@ -70,4 +71,6 @@ nostd_async = { version = "0.6.1", features = ["wfe"] }
 futures = { version = "0.3.30", default-features = false, features = [
   "async-await",
 ] }
+i2c-write-iter = "1.0.0"
 itertools = { version = "0.12.0", default-features = false }
+portable-atomic = { version = "1.6.0", features = ["critical-section"] }
diff --git a/rp2040-hal/Cargo.toml b/rp2040-hal/Cargo.toml
index bcd075b0c..4f413d8fb 100644
--- a/rp2040-hal/Cargo.toml
+++ b/rp2040-hal/Cargo.toml
@@ -48,6 +48,8 @@ frunk = { version = "0.4.1", default-features = false }
 
 bitfield = { version = "0.14.0" }
 
+i2c-write-iter = { version = "1.0.0", features = ["async"], optional = true }
+
 [dev-dependencies]
 cortex-m-rt = "0.7"
 cortex-m-rtic = "1.1.4"
@@ -98,6 +100,9 @@ defmt = ["dep:defmt"]
 # Implement `rtic_monotonic::Monotonic` based on the RP2040 timer peripheral
 rtic-monotonic = ["dep:rtic-monotonic"]
 
+# Implement `i2c-write-iter` traits
+i2c-write-iter = ["dep:i2c-write-iter"]
+
 [[example]]
 # irq example uses cortex-m-rt::interrupt, need rt feature for that
 name = "gpio_irq_example"
diff --git a/rp2040-hal/src/i2c/controller.rs b/rp2040-hal/src/i2c/controller.rs
index 5e0afb53f..fd9bb4b8f 100644
--- a/rp2040-hal/src/i2c/controller.rs
+++ b/rp2040-hal/src/i2c/controller.rs
@@ -386,6 +386,29 @@ impl<T: Deref<Target = Block>, PINS> I2C<T, PINS, Controller> {
         }
         Ok(())
     }
+
+    #[cfg(feature = "i2c-write-iter")]
+    fn transaction_iter<'op, A, O, B>(&mut self, address: A, operations: O) -> Result<(), Error>
+    where
+        A: ValidAddress,
+        O: IntoIterator<Item = i2c_write_iter::Operation<'op, B>>,
+        B: IntoIterator<Item = u8>,
+    {
+        use i2c_write_iter::Operation;
+        self.setup(address)?;
+
+        let mut first = true;
+        let mut operations = operations.into_iter().peekable();
+        while let Some(operation) = operations.next() {
+            let last = operations.peek().is_none();
+            match operation {
+                Operation::Read(buf) => self.read_internal(first, buf, last)?,
+                Operation::WriteIter(buf) => self.write_internal(first, buf, last)?,
+            }
+            first = false;
+        }
+        Ok(())
+    }
 }
 
 impl<A: ValidAddress, T: Deref<Target = Block>, PINS> Read<A> for I2C<T, PINS, Controller> {
@@ -458,3 +481,16 @@ impl<A: ValidAddress, T: Deref<Target = Block>, PINS> eh1::I2c<A> for I2C<T, PIN
         self.transaction(address, operations.iter_mut())
     }
 }
+
+#[cfg(feature = "i2c-write-iter")]
+impl<A: i2c_write_iter::AddressMode + ValidAddress, T: Deref<Target = Block>, PINS>
+    i2c_write_iter::I2cIter<A> for I2C<T, PINS, Controller>
+{
+    fn transaction_iter<'a, O, B>(&mut self, address: A, operations: O) -> Result<(), Self::Error>
+    where
+        O: IntoIterator<Item = i2c_write_iter::Operation<'a, B>>,
+        B: IntoIterator<Item = u8>,
+    {
+        self.transaction_iter(address, operations)
+    }
+}
diff --git a/rp2040-hal/src/i2c/controller/non_blocking.rs b/rp2040-hal/src/i2c/controller/non_blocking.rs
index 8edf5ef9f..b39dff46f 100644
--- a/rp2040-hal/src/i2c/controller/non_blocking.rs
+++ b/rp2040-hal/src/i2c/controller/non_blocking.rs
@@ -259,6 +259,38 @@ where
         self.non_blocking_write_internal(true, bytes, false).await?;
         self.non_blocking_read_internal(false, read, true).await
     }
+
+    /// Writes to the i2c bus taking operations from and iterator, writing from iterator of bytes,
+    /// reading to slices of bytes.
+    #[cfg(feature = "i2c-write-iter")]
+    pub async fn transaction_iter_async<'b, A, O, B>(
+        &mut self,
+        address: A,
+        operations: O,
+    ) -> Result<(), super::Error>
+    where
+        A: ValidAddress,
+        O: IntoIterator<Item = i2c_write_iter::Operation<'b, B>>,
+        B: IntoIterator<Item = u8>,
+    {
+        self.setup(address)?;
+
+        let mut first = true;
+        let mut operations = operations.into_iter().peekable();
+        while let Some(operation) = operations.next() {
+            let last = operations.peek().is_none();
+            match operation {
+                i2c_write_iter::Operation::Read(buf) => {
+                    self.non_blocking_read_internal(first, buf, last).await?
+                }
+                i2c_write_iter::Operation::WriteIter(buf) => {
+                    self.non_blocking_write_internal(first, buf, last).await?
+                }
+            }
+            first = false;
+        }
+        Ok(())
+    }
 }
 
 impl<T, PINS, A> embedded_hal_async::i2c::I2c<A> for I2C<T, PINS, Controller>
@@ -292,3 +324,23 @@ where
         Ok(())
     }
 }
+
+#[cfg(feature = "i2c-write-iter")]
+impl<T, PINS, A> i2c_write_iter::non_blocking::I2cIter<A> for I2C<T, PINS, Controller>
+where
+    Self: AsyncPeripheral,
+    A: 'static + ValidAddress + AddressMode,
+    T: Deref<Target = RegisterBlock>,
+{
+    async fn transaction_iter<'a, O, B>(
+        &mut self,
+        address: A,
+        operations: O,
+    ) -> Result<(), Self::Error>
+    where
+        O: IntoIterator<Item = i2c_write_iter::Operation<'a, B>>,
+        B: IntoIterator<Item = u8>,
+    {
+        self.transaction_iter_async(address, operations).await
+    }
+}
diff --git a/rp2040-hal/src/lib.rs b/rp2040-hal/src/lib.rs
index 0675485cd..bb4b8e26b 100644
--- a/rp2040-hal/src/lib.rs
+++ b/rp2040-hal/src/lib.rs
@@ -33,6 +33,8 @@
 //! * **rtic-monotonic** -
 //!   Implement
 //!   `rtic_monotonic::Monotonic` based on the RP2040 timer peripheral
+//! * **i2c-write-iter** -
+//!   Implement `i2c_write_iter` traits for `I2C<_, _, Controller>`.
 
 #![warn(missing_docs)]
 #![no_std]