5 files changed, 105 insertions, 197 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))
     }
 }
diff --git a/embassy-net/src/lib.rs b/embassy-net/src/lib.rs
index edb969842..e18e819cb 100644
--- a/embassy-net/src/lib.rs
+++ b/embassy-net/src/lib.rs
@@ -8,12 +8,9 @@
 // This mod MUST go first, so that the others see its macros.
 pub(crate) mod fmt;
 
-mod device;
-mod packet_pool;
+pub mod device;
 mod stack;
 
-pub use device::{Device, LinkState};
-pub use packet_pool::{Packet, PacketBox, PacketBoxExt, PacketBuf, MTU};
 pub use stack::{Config, ConfigStrategy, Stack, StackResources};
 
 #[cfg(feature = "tcp")]
@@ -23,7 +20,6 @@ pub mod tcp;
 pub mod udp;
 
 // smoltcp reexports
-pub use smoltcp::phy::{DeviceCapabilities, Medium};
 pub use smoltcp::time::{Duration as SmolDuration, Instant as SmolInstant};
 #[cfg(feature = "medium-ethernet")]
 pub use smoltcp::wire::{EthernetAddress, HardwareAddress};
diff --git a/embassy-net/src/packet_pool.rs b/embassy-net/src/packet_pool.rs
deleted file mode 100644
index cb8a1316c..000000000
--- a/embassy-net/src/packet_pool.rs
+++ /dev/null
@@ -1,107 +0,0 @@
-use core::ops::{Deref, DerefMut, Range};
-
-use as_slice::{AsMutSlice, AsSlice};
-use atomic_pool::{pool, Box};
-
-pub const MTU: usize = 1516;
-
-#[cfg(feature = "pool-4")]
-pub const PACKET_POOL_SIZE: usize = 4;
-
-#[cfg(feature = "pool-8")]
-pub const PACKET_POOL_SIZE: usize = 8;
-
-#[cfg(feature = "pool-16")]
-pub const PACKET_POOL_SIZE: usize = 16;
-
-#[cfg(feature = "pool-32")]
-pub const PACKET_POOL_SIZE: usize = 32;
-
-#[cfg(feature = "pool-64")]
-pub const PACKET_POOL_SIZE: usize = 64;
-
-#[cfg(feature = "pool-128")]
-pub const PACKET_POOL_SIZE: usize = 128;
-
-pool!(pub PacketPool: [Packet; PACKET_POOL_SIZE]);
-pub type PacketBox = Box<PacketPool>;
-
-#[repr(align(4))]
-pub struct Packet(pub [u8; MTU]);
-
-impl Packet {
-    pub const fn new() -> Self {
-        Self([0; MTU])
-    }
-}
-
-pub trait PacketBoxExt {
-    fn slice(self, range: Range<usize>) -> PacketBuf;
-}
-
-impl PacketBoxExt for PacketBox {
-    fn slice(self, range: Range<usize>) -> PacketBuf {
-        PacketBuf { packet: self, range }
-    }
-}
-
-impl AsSlice for Packet {
-    type Element = u8;
-
-    fn as_slice(&self) -> &[Self::Element] {
-        &self.deref()[..]
-    }
-}
-
-impl AsMutSlice for Packet {
-    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
-        &mut self.deref_mut()[..]
-    }
-}
-
-impl Deref for Packet {
-    type Target = [u8; MTU];
-
-    fn deref(&self) -> &[u8; MTU] {
-        &self.0
-    }
-}
-
-impl DerefMut for Packet {
-    fn deref_mut(&mut self) -> &mut [u8; MTU] {
-        &mut self.0
-    }
-}
-
-pub struct PacketBuf {
-    packet: PacketBox,
-    range: Range<usize>,
-}
-
-impl AsSlice for PacketBuf {
-    type Element = u8;
-
-    fn as_slice(&self) -> &[Self::Element] {
-        &self.packet[self.range.clone()]
-    }
-}
-
-impl AsMutSlice for PacketBuf {
-    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
-        &mut self.packet[self.range.clone()]
-    }
-}
-
-impl Deref for PacketBuf {
-    type Target = [u8];
-
-    fn deref(&self) -> &[u8] {
-        &self.packet[self.range.clone()]
-    }
-}
-
-impl DerefMut for PacketBuf {
-    fn deref_mut(&mut self) -> &mut [u8] {
-        &mut self.packet[self.range.clone()]
-    }
-}
diff --git a/embassy-net/src/stack.rs b/embassy-net/src/stack.rs
index 5c4fb0442..21316e485 100644
--- a/embassy-net/src/stack.rs
+++ b/embassy-net/src/stack.rs
@@ -12,7 +12,7 @@ use smoltcp::iface::{Interface, InterfaceBuilder, SocketSet, SocketStorage};
 #[cfg(feature = "medium-ethernet")]
 use smoltcp::iface::{Neighbor, NeighborCache, Route, Routes};
 #[cfg(feature = "medium-ethernet")]
-use smoltcp::phy::{Device as _, Medium};
+use smoltcp::phy::Medium;
 #[cfg(feature = "dhcpv4")]
 use smoltcp::socket::dhcpv4;
 use smoltcp::time::Instant as SmolInstant;
@@ -67,7 +67,7 @@ pub struct Stack<D: Device> {
 }
 
 struct Inner<D: Device> {
-    device: DeviceAdapter<D>,
+    device: D,
     link_up: bool,
     config: Option<Config>,
     #[cfg(feature = "dhcpv4")]
@@ -83,7 +83,7 @@ pub(crate) struct SocketStack {
 
 impl<D: Device + 'static> Stack<D> {
     pub fn new<const ADDR: usize, const SOCK: usize, const NEIGH: usize>(
-        device: D,
+        mut device: D,
         config: ConfigStrategy,
         resources: &'static mut StackResources<ADDR, SOCK, NEIGH>,
         random_seed: u64,
@@ -98,8 +98,6 @@ impl<D: Device + 'static> Stack<D> {
             [0, 0, 0, 0, 0, 0]
         };
 
-        let mut device = DeviceAdapter::new(device);
-
         let mut b = InterfaceBuilder::new();
         b = b.ip_addrs(&mut resources.addresses[..]);
         b = b.random_seed(random_seed);
@@ -111,7 +109,10 @@ impl<D: Device + 'static> Stack<D> {
             b = b.routes(Routes::new(&mut resources.routes[..]));
         }
 
-        let iface = b.finalize(&mut device);
+        let iface = b.finalize(&mut DeviceAdapter {
+            inner: &mut device,
+            cx: None,
+        });
 
         let sockets = SocketSet::new(&mut resources.sockets[..]);
 
@@ -155,7 +156,7 @@ impl<D: Device + 'static> Stack<D> {
     }
 
     pub fn ethernet_address(&self) -> [u8; 6] {
-        self.with(|_s, i| i.device.device.ethernet_address())
+        self.with(|_s, i| i.device.ethernet_address())
     }
 
     pub fn is_link_up(&self) -> bool {
@@ -238,11 +239,14 @@ impl<D: Device + 'static> Inner<D> {
     }
 
     fn poll(&mut self, cx: &mut Context<'_>, s: &mut SocketStack) {
-        self.device.device.register_waker(cx.waker());
         s.waker.register(cx.waker());
 
         let timestamp = instant_to_smoltcp(Instant::now());
-        if s.iface.poll(timestamp, &mut self.device, &mut s.sockets).is_err() {
+        let mut smoldev = DeviceAdapter {
+            cx: Some(cx),
+            inner: &mut self.device,
+        };
+        if s.iface.poll(timestamp, &mut smoldev, &mut s.sockets).is_err() {
             // If poll() returns error, it may not be done yet, so poll again later.
             cx.waker().wake_by_ref();
             return;
@@ -250,7 +254,7 @@ impl<D: Device + 'static> Inner<D> {
 
         // Update link up
         let old_link_up = self.link_up;
-        self.link_up = self.device.device.link_state() == LinkState::Up;
+        self.link_up = self.device.link_state(cx) == LinkState::Up;
 
         // Print when changed
         if old_link_up != self.link_up {
diff --git a/embassy-net/src/tcp.rs b/embassy-net/src/tcp.rs
index 73cf2d4e4..0ed4b81e2 100644
--- a/embassy-net/src/tcp.rs
+++ b/embassy-net/src/tcp.rs
@@ -9,8 +9,8 @@ use smoltcp::time::Duration;
 use smoltcp::wire::{IpEndpoint, IpListenEndpoint};
 
 use super::stack::Stack;
+use crate::device::Device;
 use crate::stack::SocketStack;
-use crate::Device;
 
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]