1 files changed, 89 insertions, 74 deletions

diff --git a/embassy-net/src/device.rs b/embassy-net/src/device.rs
index 4bdfd7720..5d86ca91e 100644
--- a/embassy-net/src/device.rs
+++ b/embassy-net/src/device.rs
@@ -1,10 +1,7 @@
-use core::task::Waker;
+use core::task::Context;
 
-use smoltcp::phy::{Device as SmolDevice, DeviceCapabilities};
-use smoltcp::time::Instant as SmolInstant;
-
-use crate::packet_pool::PacketBoxExt;
-use crate::{Packet, PacketBox, PacketBuf};
+use smoltcp::phy;
+pub use smoltcp::phy::{Checksum, ChecksumCapabilities, DeviceCapabilities, Medium};
 
 #[derive(PartialEq, Eq, Clone, Copy)]
 pub enum LinkState {
@@ -13,115 +10,133 @@ pub enum LinkState {
 }
 
 pub trait Device {
-    fn is_transmit_ready(&mut self) -> bool;
-    fn transmit(&mut self, pkt: PacketBuf);
-    fn receive(&mut self) -> Option<PacketBuf>;
+    type RxToken<'a>: RxToken
+    where
+        Self: 'a;
+    type TxToken<'a>: TxToken
+    where
+        Self: 'a;
 
-    fn register_waker(&mut self, waker: &Waker);
-    fn capabilities(&self) -> DeviceCapabilities;
-    fn link_state(&mut self) -> LinkState;
+    fn receive(&mut self, cx: &mut Context) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)>;
+    fn transmit(&mut self, cx: &mut Context) -> Option<Self::TxToken<'_>>;
+    fn link_state(&mut self, cx: &mut Context) -> LinkState;
+
+    fn capabilities(&self) -> phy::DeviceCapabilities;
     fn ethernet_address(&self) -> [u8; 6];
 }
 
 impl<T: ?Sized + Device> Device for &mut T {
-    fn is_transmit_ready(&mut self) -> bool {
-        T::is_transmit_ready(self)
-    }
-    fn transmit(&mut self, pkt: PacketBuf) {
-        T::transmit(self, pkt)
-    }
-    fn receive(&mut self) -> Option<PacketBuf> {
-        T::receive(self)
+    type RxToken<'a> = T::RxToken<'a>
+    where
+        Self: 'a;
+    type TxToken<'a> = T::TxToken<'a>
+    where
+        Self: 'a;
+
+    fn transmit(&mut self, cx: &mut Context) -> Option<Self::TxToken<'_>> {
+        T::transmit(self, cx)
     }
-    fn register_waker(&mut self, waker: &Waker) {
-        T::register_waker(self, waker)
+    fn receive(&mut self, cx: &mut Context) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
+        T::receive(self, cx)
     }
-    fn capabilities(&self) -> DeviceCapabilities {
+    fn capabilities(&self) -> phy::DeviceCapabilities {
         T::capabilities(self)
     }
-    fn link_state(&mut self) -> LinkState {
-        T::link_state(self)
+    fn link_state(&mut self, cx: &mut Context) -> LinkState {
+        T::link_state(self, cx)
     }
     fn ethernet_address(&self) -> [u8; 6] {
         T::ethernet_address(self)
     }
 }
 
-pub struct DeviceAdapter<D: Device> {
-    pub device: D,
-    caps: DeviceCapabilities,
+/// A token to receive a single network packet.
+pub trait RxToken {
+    /// Consumes the token to receive a single network packet.
+    ///
+    /// This method receives a packet and then calls the given closure `f` with the raw
+    /// packet bytes as argument.
+    fn consume<R, F>(self, f: F) -> R
+    where
+        F: FnOnce(&mut [u8]) -> R;
 }
 
-impl<D: Device> DeviceAdapter<D> {
-    pub(crate) fn new(device: D) -> Self {
-        Self {
-            caps: device.capabilities(),
-            device,
-        }
-    }
+/// A token to transmit a single network packet.
+pub trait TxToken {
+    /// Consumes the token to send a single network packet.
+    ///
+    /// This method constructs a transmit buffer of size `len` and calls the passed
+    /// closure `f` with a mutable reference to that buffer. The closure should construct
+    /// a valid network packet (e.g. an ethernet packet) in the buffer. When the closure
+    /// returns, the transmit buffer is sent out.
+    fn consume<R, F>(self, len: usize, f: F) -> R
+    where
+        F: FnOnce(&mut [u8]) -> R;
+}
+
+///////////////////////////
+
+pub(crate) struct DeviceAdapter<'d, 'c, T>
+where
+    T: Device,
+{
+    // must be Some when actually using this to rx/tx
+    pub cx: Option<&'d mut Context<'c>>,
+    pub inner: &'d mut T,
 }
 
-impl<D: Device> SmolDevice for DeviceAdapter<D> {
-    type RxToken<'a> = RxToken where Self: 'a;
-    type TxToken<'a> = TxToken<'a, D> where Self: 'a;
+impl<'d, 'c, T> phy::Device for DeviceAdapter<'d, 'c, T>
+where
+    T: Device,
+{
+    type RxToken<'a> = RxTokenAdapter<T::RxToken<'a>> where Self: 'a;
+    type TxToken<'a> = TxTokenAdapter<T::TxToken<'a>> where Self: 'a;
 
     fn receive(&mut self) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
-        let tx_pkt = PacketBox::new(Packet::new())?;
-        let rx_pkt = self.device.receive()?;
-        let rx_token = RxToken { pkt: rx_pkt };
-        let tx_token = TxToken {
-            device: &mut self.device,
-            pkt: tx_pkt,
-        };
-
-        Some((rx_token, tx_token))
+        self.inner
+            .receive(self.cx.as_deref_mut().unwrap())
+            .map(|(rx, tx)| (RxTokenAdapter(rx), TxTokenAdapter(tx)))
     }
 
     /// Construct a transmit token.
     fn transmit(&mut self) -> Option<Self::TxToken<'_>> {
-        if !self.device.is_transmit_ready() {
-            return None;
-        }
-
-        let tx_pkt = PacketBox::new(Packet::new())?;
-        Some(TxToken {
-            device: &mut self.device,
-            pkt: tx_pkt,
-        })
+        self.inner.transmit(self.cx.as_deref_mut().unwrap()).map(TxTokenAdapter)
     }
 
     /// Get a description of device capabilities.
-    fn capabilities(&self) -> DeviceCapabilities {
-        self.caps.clone()
+    fn capabilities(&self) -> phy::DeviceCapabilities {
+        self.inner.capabilities()
     }
 }
 
-pub struct RxToken {
-    pkt: PacketBuf,
-}
+pub(crate) struct RxTokenAdapter<T>(T)
+where
+    T: RxToken;
 
-impl smoltcp::phy::RxToken for RxToken {
-    fn consume<R, F>(mut self, _timestamp: SmolInstant, f: F) -> smoltcp::Result<R>
+impl<T> phy::RxToken for RxTokenAdapter<T>
+where
+    T: RxToken,
+{
+    fn consume<R, F>(self, _timestamp: smoltcp::time::Instant, f: F) -> smoltcp::Result<R>
     where
         F: FnOnce(&mut [u8]) -> smoltcp::Result<R>,
     {
-        f(&mut self.pkt)
+        self.0.consume(|buf| f(buf))
     }
 }
 
-pub struct TxToken<'a, D: Device> {
-    device: &'a mut D,
-    pkt: PacketBox,
-}
+pub(crate) struct TxTokenAdapter<T>(T)
+where
+    T: TxToken;
 
-impl<'a, D: Device> smoltcp::phy::TxToken for TxToken<'a, D> {
-    fn consume<R, F>(self, _timestamp: SmolInstant, len: usize, f: F) -> smoltcp::Result<R>
+impl<T> phy::TxToken for TxTokenAdapter<T>
+where
+    T: TxToken,
+{
+    fn consume<R, F>(self, _timestamp: smoltcp::time::Instant, len: usize, f: F) -> smoltcp::Result<R>
     where
         F: FnOnce(&mut [u8]) -> smoltcp::Result<R>,
     {
-        let mut buf = self.pkt.slice(0..len);
-        let r = f(&mut buf)?;
-        self.device.transmit(buf);
-        Ok(r)
+        self.0.consume(len, |buf| f(buf))
     }
 }