From 883e28a0fb80f9139c0236ce44597c23a0917e4d Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Mon, 30 May 2022 00:08:28 +0200
Subject: [PATCH] usb: add first, last params to ControlPipe data_in, data_out.

---
 embassy-nrf/src/usb.rs    | 17 +++++++++++++----
 embassy-usb/src/driver.rs | 10 ++++++++--
 embassy-usb/src/lib.rs    | 23 ++++++++++++++++++-----
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/embassy-nrf/src/usb.rs b/embassy-nrf/src/usb.rs
index 23401c837..9155eb37a 100644
--- a/embassy-nrf/src/usb.rs
+++ b/embassy-nrf/src/usb.rs
@@ -657,7 +657,12 @@ impl<'d, T: Instance> driver::ControlPipe for ControlPipe<'d, T> {
         }
     }
 
-    fn data_out<'a>(&'a mut self, buf: &'a mut [u8]) -> Self::DataOutFuture<'a> {
+    fn data_out<'a>(
+        &'a mut self,
+        buf: &'a mut [u8],
+        _first: bool,
+        _last: bool,
+    ) -> Self::DataOutFuture<'a> {
         async move {
             let regs = T::regs();
 
@@ -694,13 +699,17 @@ impl<'d, T: Instance> driver::ControlPipe for ControlPipe<'d, T> {
         }
     }
 
-    fn data_in<'a>(&'a mut self, buf: &'a [u8], last_packet: bool) -> Self::DataInFuture<'a> {
+    fn data_in<'a>(
+        &'a mut self,
+        buf: &'a [u8],
+        _first: bool,
+        last: bool,
+    ) -> Self::DataInFuture<'a> {
         async move {
             let regs = T::regs();
             regs.events_ep0datadone.reset();
 
-            regs.shorts
-                .write(|w| w.ep0datadone_ep0status().bit(last_packet));
+            regs.shorts.write(|w| w.ep0datadone_ep0status().bit(last));
 
             // This starts a TX on EP0. events_ep0datadone notifies when done.
             unsafe { write_dma::<T>(0, buf) }
diff --git a/embassy-usb/src/driver.rs b/embassy-usb/src/driver.rs
index acd2e298d..9cd4afdab 100644
--- a/embassy-usb/src/driver.rs
+++ b/embassy-usb/src/driver.rs
@@ -155,12 +155,18 @@ pub trait ControlPipe {
     ///
     /// Must be called after `setup()` for requests with `direction` of `Out`
     /// and `length` greater than zero.
-    fn data_out<'a>(&'a mut self, buf: &'a mut [u8]) -> Self::DataOutFuture<'a>;
+    fn data_out<'a>(
+        &'a mut self,
+        buf: &'a mut [u8],
+        first: bool,
+        last: bool,
+    ) -> Self::DataOutFuture<'a>;
 
     /// Sends a DATA IN packet with `data` in response to a control read request.
     ///
     /// If `last_packet` is true, the STATUS packet will be ACKed following the transfer of `data`.
-    fn data_in<'a>(&'a mut self, data: &'a [u8], last_packet: bool) -> Self::DataInFuture<'a>;
+    fn data_in<'a>(&'a mut self, data: &'a [u8], first: bool, last: bool)
+        -> Self::DataInFuture<'a>;
 
     /// Accepts a control request.
     ///
diff --git a/embassy-usb/src/lib.rs b/embassy-usb/src/lib.rs
index a2879bdbc..9101d81bd 100644
--- a/embassy-usb/src/lib.rs
+++ b/embassy-usb/src/lib.rs
@@ -292,12 +292,12 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
                 let len = data.len().min(resp_length);
                 let need_zlp = len != resp_length && (len % usize::from(max_packet_size)) == 0;
 
-                let mut chunks = data[0..len]
+                let chunks = data[0..len]
                     .chunks(max_packet_size)
                     .chain(need_zlp.then(|| -> &[u8] { &[] }));
 
-                while let Some(chunk) = chunks.next() {
-                    match self.control.data_in(chunk, chunks.size_hint().0 == 0).await {
+                for (first, last, chunk) in first_last(chunks) {
+                    match self.control.data_in(chunk, first, last).await {
                         Ok(()) => {}
                         Err(e) => {
                             warn!("control accept_in failed: {:?}", e);
@@ -315,8 +315,9 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         let max_packet_size = self.control.max_packet_size();
         let mut total = 0;
 
-        for chunk in self.control_buf[..req_length].chunks_mut(max_packet_size) {
-            let size = match self.control.data_out(chunk).await {
+        let chunks = self.control_buf[..req_length].chunks_mut(max_packet_size);
+        for (first, last, chunk) in first_last(chunks) {
+            let size = match self.control.data_out(chunk, first, last).await {
                 Ok(x) => x,
                 Err(e) => {
                     warn!("usb: failed to read CONTROL OUT data stage: {:?}", e);
@@ -668,3 +669,15 @@ impl<'d, D: Driver<'d>> Inner<'d, D> {
         }
     }
 }
+
+fn first_last<T: Iterator>(iter: T) -> impl Iterator<Item = (bool, bool, T::Item)> {
+    let mut iter = iter.peekable();
+    let mut first = true;
+    core::iter::from_fn(move || {
+        let val = iter.next()?;
+        let is_first = first;
+        first = false;
+        let is_last = iter.peek().is_none();
+        Some((is_first, is_last, val))
+    })
+}