make usb endpoint allocator methods accept an optional EndpointAddress

2 files changed, 51 insertions, 13 deletions

diff --git a/embassy-stm32/src/usb/otg.rs b/embassy-stm32/src/usb/otg.rs
index 590d1a427..3547ded00 100644
--- a/embassy-stm32/src/usb/otg.rs
+++ b/embassy-stm32/src/usb/otg.rs
@@ -231,19 +231,23 @@ impl<'d, T: Instance> embassy_usb_driver::Driver<'d> for Driver<'d, T> {
     fn alloc_endpoint_in(
         &mut self,
         ep_type: EndpointType,
+        ep_addr: Option<EndpointAddress>,
         max_packet_size: u16,
         interval_ms: u8,
     ) -> Result<Self::EndpointIn, EndpointAllocError> {
-        self.inner.alloc_endpoint_in(ep_type, max_packet_size, interval_ms)
+        self.inner
+            .alloc_endpoint_in(ep_type, ep_addr, max_packet_size, interval_ms)
     }
 
     fn alloc_endpoint_out(
         &mut self,
         ep_type: EndpointType,
+        ep_addr: Option<EndpointAddress>,
         max_packet_size: u16,
         interval_ms: u8,
     ) -> Result<Self::EndpointOut, EndpointAllocError> {
-        self.inner.alloc_endpoint_out(ep_type, max_packet_size, interval_ms)
+        self.inner
+            .alloc_endpoint_out(ep_type, ep_addr, max_packet_size, interval_ms)
     }
 
     fn start(self, control_max_packet_size: u16) -> (Self::Bus, Self::ControlPipe) {
diff --git a/embassy-stm32/src/usb/usb.rs b/embassy-stm32/src/usb/usb.rs
index 3e8e74a1f..05c28aceb 100644
--- a/embassy-stm32/src/usb/usb.rs
+++ b/embassy-stm32/src/usb/usb.rs
@@ -362,6 +362,7 @@ impl<'d, T: Instance> Driver<'d, T> {
     fn alloc_endpoint<D: Dir>(
         &mut self,
         ep_type: EndpointType,
+        ep_addr: Option<EndpointAddress>,
         max_packet_size: u16,
         interval_ms: u8,
     ) -> Result<Endpoint<'d, T, D>, driver::EndpointAllocError> {
@@ -373,25 +374,56 @@ impl<'d, T: Instance> Driver<'d, T> {
             D::dir()
         );
 
-        let index = self.alloc.iter_mut().enumerate().find(|(i, ep)| {
-            if *i == 0 && ep_type != EndpointType::Control {
-                return false; // reserved for control pipe
+        let index = if let Some(addr) = ep_addr {
+            // Use the specified endpoint address
+            let requested_index = addr.index();
+            if requested_index >= EP_COUNT {
+                return Err(EndpointAllocError);
             }
+            if requested_index == 0 && ep_type != EndpointType::Control {
+                return Err(EndpointAllocError); // EP0 is reserved for control
+            }
+
+            let ep = &self.alloc[requested_index];
             let used = ep.used_out || ep.used_in;
             if used && (ep.ep_type == EndpointType::Isochronous) {
                 // Isochronous endpoints are always double-buffered.
                 // Their corresponding endpoint/channel registers are forced to be unidirectional.
                 // Do not reuse this index.
-                // FIXME: Bulk endpoints can be double buffered, but are not in the current implementation.
-                return false;
+                return Err(EndpointAllocError);
             }
 
             let used_dir = match D::dir() {
                 Direction::Out => ep.used_out,
                 Direction::In => ep.used_in,
             };
-            !used || (ep.ep_type == ep_type && !used_dir)
-        });
+            if used && (ep.ep_type != ep_type || used_dir) {
+                return Err(EndpointAllocError);
+            }
+
+            Some((requested_index, &mut self.alloc[requested_index]))
+        } else {
+            // Find any available endpoint
+            self.alloc.iter_mut().enumerate().find(|(i, ep)| {
+                if *i == 0 && ep_type != EndpointType::Control {
+                    return false; // reserved for control pipe
+                }
+                let used = ep.used_out || ep.used_in;
+                if used && (ep.ep_type == EndpointType::Isochronous) {
+                    // Isochronous endpoints are always double-buffered.
+                    // Their corresponding endpoint/channel registers are forced to be unidirectional.
+                    // Do not reuse this index.
+                    // FIXME: Bulk endpoints can be double buffered, but are not in the current implementation.
+                    return false;
+                }
+
+                let used_dir = match D::dir() {
+                    Direction::Out => ep.used_out,
+                    Direction::In => ep.used_in,
+                };
+                !used || (ep.ep_type == ep_type && !used_dir)
+            })
+        };
 
         let (index, ep) = match index {
             Some(x) => x,
@@ -479,27 +511,29 @@ impl<'d, T: Instance> driver::Driver<'d> for Driver<'d, T> {
     fn alloc_endpoint_in(
         &mut self,
         ep_type: EndpointType,
+        ep_addr: Option<EndpointAddress>,
         max_packet_size: u16,
         interval_ms: u8,
     ) -> Result<Self::EndpointIn, driver::EndpointAllocError> {
-        self.alloc_endpoint(ep_type, max_packet_size, interval_ms)
+        self.alloc_endpoint(ep_type, ep_addr, max_packet_size, interval_ms)
     }
 
     fn alloc_endpoint_out(
         &mut self,
         ep_type: EndpointType,
+        ep_addr: Option<EndpointAddress>,
         max_packet_size: u16,
         interval_ms: u8,
     ) -> Result<Self::EndpointOut, driver::EndpointAllocError> {
-        self.alloc_endpoint(ep_type, max_packet_size, interval_ms)
+        self.alloc_endpoint(ep_type, ep_addr, max_packet_size, interval_ms)
     }
 
     fn start(mut self, control_max_packet_size: u16) -> (Self::Bus, Self::ControlPipe) {
         let ep_out = self
-            .alloc_endpoint(EndpointType::Control, control_max_packet_size, 0)
+            .alloc_endpoint(EndpointType::Control, None, control_max_packet_size, 0)
             .unwrap();
         let ep_in = self
-            .alloc_endpoint(EndpointType::Control, control_max_packet_size, 0)
+            .alloc_endpoint(EndpointType::Control, None, control_max_packet_size, 0)
             .unwrap();
         assert_eq!(ep_out.info.addr.index(), 0);
         assert_eq!(ep_in.info.addr.index(), 0);