Merge pull request #1096 from embassy-rs/net-refactor2

net: remove packet pool

6 files changed, 395 insertions, 498 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..afe0d6da0 100644
--- a/embassy-net/src/lib.rs
+++ b/embassy-net/src/lib.rs
@@ -8,22 +8,30 @@
 // This mod MUST go first, so that the others see its macros.
 pub(crate) mod fmt;
 
-mod device;
-mod packet_pool;
-mod stack;
-
-pub use device::{Device, LinkState};
-pub use packet_pool::{Packet, PacketBox, PacketBoxExt, PacketBuf, MTU};
-pub use stack::{Config, ConfigStrategy, Stack, StackResources};
-
+pub mod device;
 #[cfg(feature = "tcp")]
 pub mod tcp;
-
 #[cfg(feature = "udp")]
 pub mod udp;
 
+use core::cell::RefCell;
+use core::future::{poll_fn, Future};
+use core::task::{Context, Poll};
+
+use embassy_sync::waitqueue::WakerRegistration;
+use embassy_time::{Instant, Timer};
+use futures::pin_mut;
+use heapless::Vec;
+#[cfg(feature = "dhcpv4")]
+use smoltcp::iface::SocketHandle;
+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::Medium;
+#[cfg(feature = "dhcpv4")]
+use smoltcp::socket::dhcpv4;
 // 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};
@@ -32,3 +40,288 @@ pub use smoltcp::wire::{IpAddress, IpCidr, Ipv4Address, Ipv4Cidr};
 pub use smoltcp::wire::{Ipv6Address, Ipv6Cidr};
 #[cfg(feature = "udp")]
 pub use smoltcp::{socket::udp::PacketMetadata, wire::IpListenEndpoint};
+
+use crate::device::{Device, DeviceAdapter, LinkState};
+
+const LOCAL_PORT_MIN: u16 = 1025;
+const LOCAL_PORT_MAX: u16 = 65535;
+
+pub struct StackResources<const ADDR: usize, const SOCK: usize, const NEIGHBOR: usize> {
+    addresses: [IpCidr; ADDR],
+    sockets: [SocketStorage<'static>; SOCK],
+
+    #[cfg(feature = "medium-ethernet")]
+    routes: [Option<(IpCidr, Route)>; 1],
+    #[cfg(feature = "medium-ethernet")]
+    neighbor_cache: [Option<(IpAddress, Neighbor)>; NEIGHBOR],
+}
+
+impl<const ADDR: usize, const SOCK: usize, const NEIGHBOR: usize> StackResources<ADDR, SOCK, NEIGHBOR> {
+    pub fn new() -> Self {
+        Self {
+            addresses: [IpCidr::new(Ipv4Address::UNSPECIFIED.into(), 32); ADDR],
+            sockets: [SocketStorage::EMPTY; SOCK],
+            #[cfg(feature = "medium-ethernet")]
+            routes: [None; 1],
+            #[cfg(feature = "medium-ethernet")]
+            neighbor_cache: [None; NEIGHBOR],
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Config {
+    pub address: Ipv4Cidr,
+    pub gateway: Option<Ipv4Address>,
+    pub dns_servers: Vec<Ipv4Address, 3>,
+}
+
+pub enum ConfigStrategy {
+    Static(Config),
+    #[cfg(feature = "dhcpv4")]
+    Dhcp,
+}
+
+pub struct Stack<D: Device> {
+    pub(crate) socket: RefCell<SocketStack>,
+    inner: RefCell<Inner<D>>,
+}
+
+struct Inner<D: Device> {
+    device: D,
+    link_up: bool,
+    config: Option<Config>,
+    #[cfg(feature = "dhcpv4")]
+    dhcp_socket: Option<SocketHandle>,
+}
+
+pub(crate) struct SocketStack {
+    pub(crate) sockets: SocketSet<'static>,
+    pub(crate) iface: Interface<'static>,
+    pub(crate) waker: WakerRegistration,
+    next_local_port: u16,
+}
+
+impl<D: Device + 'static> Stack<D> {
+    pub fn new<const ADDR: usize, const SOCK: usize, const NEIGH: usize>(
+        mut device: D,
+        config: ConfigStrategy,
+        resources: &'static mut StackResources<ADDR, SOCK, NEIGH>,
+        random_seed: u64,
+    ) -> Self {
+        #[cfg(feature = "medium-ethernet")]
+        let medium = device.capabilities().medium;
+
+        #[cfg(feature = "medium-ethernet")]
+        let ethernet_addr = if medium == Medium::Ethernet {
+            device.ethernet_address()
+        } else {
+            [0, 0, 0, 0, 0, 0]
+        };
+
+        let mut b = InterfaceBuilder::new();
+        b = b.ip_addrs(&mut resources.addresses[..]);
+        b = b.random_seed(random_seed);
+
+        #[cfg(feature = "medium-ethernet")]
+        if medium == Medium::Ethernet {
+            b = b.hardware_addr(HardwareAddress::Ethernet(EthernetAddress(ethernet_addr)));
+            b = b.neighbor_cache(NeighborCache::new(&mut resources.neighbor_cache[..]));
+            b = b.routes(Routes::new(&mut resources.routes[..]));
+        }
+
+        let iface = b.finalize(&mut DeviceAdapter {
+            inner: &mut device,
+            cx: None,
+        });
+
+        let sockets = SocketSet::new(&mut resources.sockets[..]);
+
+        let next_local_port = (random_seed % (LOCAL_PORT_MAX - LOCAL_PORT_MIN) as u64) as u16 + LOCAL_PORT_MIN;
+
+        let mut inner = Inner {
+            device,
+            link_up: false,
+            config: None,
+            #[cfg(feature = "dhcpv4")]
+            dhcp_socket: None,
+        };
+        let mut socket = SocketStack {
+            sockets,
+            iface,
+            waker: WakerRegistration::new(),
+            next_local_port,
+        };
+
+        match config {
+            ConfigStrategy::Static(config) => inner.apply_config(&mut socket, config),
+            #[cfg(feature = "dhcpv4")]
+            ConfigStrategy::Dhcp => {
+                let handle = socket.sockets.add(smoltcp::socket::dhcpv4::Socket::new());
+                inner.dhcp_socket = Some(handle);
+            }
+        }
+
+        Self {
+            socket: RefCell::new(socket),
+            inner: RefCell::new(inner),
+        }
+    }
+
+    fn with<R>(&self, f: impl FnOnce(&SocketStack, &Inner<D>) -> R) -> R {
+        f(&*self.socket.borrow(), &*self.inner.borrow())
+    }
+
+    fn with_mut<R>(&self, f: impl FnOnce(&mut SocketStack, &mut Inner<D>) -> R) -> R {
+        f(&mut *self.socket.borrow_mut(), &mut *self.inner.borrow_mut())
+    }
+
+    pub fn ethernet_address(&self) -> [u8; 6] {
+        self.with(|_s, i| i.device.ethernet_address())
+    }
+
+    pub fn is_link_up(&self) -> bool {
+        self.with(|_s, i| i.link_up)
+    }
+
+    pub fn is_config_up(&self) -> bool {
+        self.with(|_s, i| i.config.is_some())
+    }
+
+    pub fn config(&self) -> Option<Config> {
+        self.with(|_s, i| i.config.clone())
+    }
+
+    pub async fn run(&self) -> ! {
+        poll_fn(|cx| {
+            self.with_mut(|s, i| i.poll(cx, s));
+            Poll::<()>::Pending
+        })
+        .await;
+        unreachable!()
+    }
+}
+
+impl SocketStack {
+    #[allow(clippy::absurd_extreme_comparisons)]
+    pub fn get_local_port(&mut self) -> u16 {
+        let res = self.next_local_port;
+        self.next_local_port = if res >= LOCAL_PORT_MAX { LOCAL_PORT_MIN } else { res + 1 };
+        res
+    }
+}
+
+impl<D: Device + 'static> Inner<D> {
+    fn apply_config(&mut self, s: &mut SocketStack, config: Config) {
+        #[cfg(feature = "medium-ethernet")]
+        let medium = self.device.capabilities().medium;
+
+        debug!("Acquired IP configuration:");
+
+        debug!("   IP address:      {}", config.address);
+        self.set_ipv4_addr(s, config.address);
+
+        #[cfg(feature = "medium-ethernet")]
+        if medium == Medium::Ethernet {
+            if let Some(gateway) = config.gateway {
+                debug!("   Default gateway: {}", gateway);
+                s.iface.routes_mut().add_default_ipv4_route(gateway).unwrap();
+            } else {
+                debug!("   Default gateway: None");
+                s.iface.routes_mut().remove_default_ipv4_route();
+            }
+        }
+        for (i, s) in config.dns_servers.iter().enumerate() {
+            debug!("   DNS server {}:    {}", i, s);
+        }
+
+        self.config = Some(config)
+    }
+
+    #[allow(unused)] // used only with dhcp
+    fn unapply_config(&mut self, s: &mut SocketStack) {
+        #[cfg(feature = "medium-ethernet")]
+        let medium = self.device.capabilities().medium;
+
+        debug!("Lost IP configuration");
+        self.set_ipv4_addr(s, Ipv4Cidr::new(Ipv4Address::UNSPECIFIED, 0));
+        #[cfg(feature = "medium-ethernet")]
+        if medium == Medium::Ethernet {
+            s.iface.routes_mut().remove_default_ipv4_route();
+        }
+        self.config = None
+    }
+
+    fn set_ipv4_addr(&mut self, s: &mut SocketStack, cidr: Ipv4Cidr) {
+        s.iface.update_ip_addrs(|addrs| {
+            let dest = addrs.iter_mut().next().unwrap();
+            *dest = IpCidr::Ipv4(cidr);
+        });
+    }
+
+    fn poll(&mut self, cx: &mut Context<'_>, s: &mut SocketStack) {
+        s.waker.register(cx.waker());
+
+        let timestamp = instant_to_smoltcp(Instant::now());
+        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;
+        }
+
+        // Update link up
+        let old_link_up = self.link_up;
+        self.link_up = self.device.link_state(cx) == LinkState::Up;
+
+        // Print when changed
+        if old_link_up != self.link_up {
+            info!("link_up = {:?}", self.link_up);
+        }
+
+        #[cfg(feature = "dhcpv4")]
+        if let Some(dhcp_handle) = self.dhcp_socket {
+            let socket = s.sockets.get_mut::<dhcpv4::Socket>(dhcp_handle);
+
+            if self.link_up {
+                match socket.poll() {
+                    None => {}
+                    Some(dhcpv4::Event::Deconfigured) => self.unapply_config(s),
+                    Some(dhcpv4::Event::Configured(config)) => {
+                        let config = Config {
+                            address: config.address,
+                            gateway: config.router,
+                            dns_servers: config.dns_servers,
+                        };
+                        self.apply_config(s, config)
+                    }
+                }
+            } else if old_link_up {
+                socket.reset();
+                self.unapply_config(s);
+            }
+        }
+        //if old_link_up || self.link_up {
+        //    self.poll_configurator(timestamp)
+        //}
+
+        if let Some(poll_at) = s.iface.poll_at(timestamp, &mut s.sockets) {
+            let t = Timer::at(instant_from_smoltcp(poll_at));
+            pin_mut!(t);
+            if t.poll(cx).is_ready() {
+                cx.waker().wake_by_ref();
+            }
+        }
+    }
+}
+
+fn instant_to_smoltcp(instant: Instant) -> SmolInstant {
+    SmolInstant::from_millis(instant.as_millis() as i64)
+}
+
+fn instant_from_smoltcp(instant: SmolInstant) -> Instant {
+    Instant::from_millis(instant.total_millis() as u64)
+}
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
deleted file mode 100644
index 5c4fb0442..000000000
--- a/embassy-net/src/stack.rs
+++ /dev/null
@@ -1,302 +0,0 @@
-use core::cell::RefCell;
-use core::future::{poll_fn, Future};
-use core::task::{Context, Poll};
-
-use embassy_sync::waitqueue::WakerRegistration;
-use embassy_time::{Instant, Timer};
-use futures::pin_mut;
-use heapless::Vec;
-#[cfg(feature = "dhcpv4")]
-use smoltcp::iface::SocketHandle;
-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};
-#[cfg(feature = "dhcpv4")]
-use smoltcp::socket::dhcpv4;
-use smoltcp::time::Instant as SmolInstant;
-#[cfg(feature = "medium-ethernet")]
-use smoltcp::wire::{EthernetAddress, HardwareAddress, IpAddress};
-use smoltcp::wire::{IpCidr, Ipv4Address, Ipv4Cidr};
-
-use crate::device::{Device, DeviceAdapter, LinkState};
-
-const LOCAL_PORT_MIN: u16 = 1025;
-const LOCAL_PORT_MAX: u16 = 65535;
-
-pub struct StackResources<const ADDR: usize, const SOCK: usize, const NEIGHBOR: usize> {
-    addresses: [IpCidr; ADDR],
-    sockets: [SocketStorage<'static>; SOCK],
-
-    #[cfg(feature = "medium-ethernet")]
-    routes: [Option<(IpCidr, Route)>; 1],
-    #[cfg(feature = "medium-ethernet")]
-    neighbor_cache: [Option<(IpAddress, Neighbor)>; NEIGHBOR],
-}
-
-impl<const ADDR: usize, const SOCK: usize, const NEIGHBOR: usize> StackResources<ADDR, SOCK, NEIGHBOR> {
-    pub fn new() -> Self {
-        Self {
-            addresses: [IpCidr::new(Ipv4Address::UNSPECIFIED.into(), 32); ADDR],
-            sockets: [SocketStorage::EMPTY; SOCK],
-            #[cfg(feature = "medium-ethernet")]
-            routes: [None; 1],
-            #[cfg(feature = "medium-ethernet")]
-            neighbor_cache: [None; NEIGHBOR],
-        }
-    }
-}
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct Config {
-    pub address: Ipv4Cidr,
-    pub gateway: Option<Ipv4Address>,
-    pub dns_servers: Vec<Ipv4Address, 3>,
-}
-
-pub enum ConfigStrategy {
-    Static(Config),
-    #[cfg(feature = "dhcpv4")]
-    Dhcp,
-}
-
-pub struct Stack<D: Device> {
-    pub(crate) socket: RefCell<SocketStack>,
-    inner: RefCell<Inner<D>>,
-}
-
-struct Inner<D: Device> {
-    device: DeviceAdapter<D>,
-    link_up: bool,
-    config: Option<Config>,
-    #[cfg(feature = "dhcpv4")]
-    dhcp_socket: Option<SocketHandle>,
-}
-
-pub(crate) struct SocketStack {
-    pub(crate) sockets: SocketSet<'static>,
-    pub(crate) iface: Interface<'static>,
-    pub(crate) waker: WakerRegistration,
-    next_local_port: u16,
-}
-
-impl<D: Device + 'static> Stack<D> {
-    pub fn new<const ADDR: usize, const SOCK: usize, const NEIGH: usize>(
-        device: D,
-        config: ConfigStrategy,
-        resources: &'static mut StackResources<ADDR, SOCK, NEIGH>,
-        random_seed: u64,
-    ) -> Self {
-        #[cfg(feature = "medium-ethernet")]
-        let medium = device.capabilities().medium;
-
-        #[cfg(feature = "medium-ethernet")]
-        let ethernet_addr = if medium == Medium::Ethernet {
-            device.ethernet_address()
-        } else {
-            [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);
-
-        #[cfg(feature = "medium-ethernet")]
-        if medium == Medium::Ethernet {
-            b = b.hardware_addr(HardwareAddress::Ethernet(EthernetAddress(ethernet_addr)));
-            b = b.neighbor_cache(NeighborCache::new(&mut resources.neighbor_cache[..]));
-            b = b.routes(Routes::new(&mut resources.routes[..]));
-        }
-
-        let iface = b.finalize(&mut device);
-
-        let sockets = SocketSet::new(&mut resources.sockets[..]);
-
-        let next_local_port = (random_seed % (LOCAL_PORT_MAX - LOCAL_PORT_MIN) as u64) as u16 + LOCAL_PORT_MIN;
-
-        let mut inner = Inner {
-            device,
-            link_up: false,
-            config: None,
-            #[cfg(feature = "dhcpv4")]
-            dhcp_socket: None,
-        };
-        let mut socket = SocketStack {
-            sockets,
-            iface,
-            waker: WakerRegistration::new(),
-            next_local_port,
-        };
-
-        match config {
-            ConfigStrategy::Static(config) => inner.apply_config(&mut socket, config),
-            #[cfg(feature = "dhcpv4")]
-            ConfigStrategy::Dhcp => {
-                let handle = socket.sockets.add(smoltcp::socket::dhcpv4::Socket::new());
-                inner.dhcp_socket = Some(handle);
-            }
-        }
-
-        Self {
-            socket: RefCell::new(socket),
-            inner: RefCell::new(inner),
-        }
-    }
-
-    fn with<R>(&self, f: impl FnOnce(&SocketStack, &Inner<D>) -> R) -> R {
-        f(&*self.socket.borrow(), &*self.inner.borrow())
-    }
-
-    fn with_mut<R>(&self, f: impl FnOnce(&mut SocketStack, &mut Inner<D>) -> R) -> R {
-        f(&mut *self.socket.borrow_mut(), &mut *self.inner.borrow_mut())
-    }
-
-    pub fn ethernet_address(&self) -> [u8; 6] {
-        self.with(|_s, i| i.device.device.ethernet_address())
-    }
-
-    pub fn is_link_up(&self) -> bool {
-        self.with(|_s, i| i.link_up)
-    }
-
-    pub fn is_config_up(&self) -> bool {
-        self.with(|_s, i| i.config.is_some())
-    }
-
-    pub fn config(&self) -> Option<Config> {
-        self.with(|_s, i| i.config.clone())
-    }
-
-    pub async fn run(&self) -> ! {
-        poll_fn(|cx| {
-            self.with_mut(|s, i| i.poll(cx, s));
-            Poll::<()>::Pending
-        })
-        .await;
-        unreachable!()
-    }
-}
-
-impl SocketStack {
-    #[allow(clippy::absurd_extreme_comparisons)]
-    pub fn get_local_port(&mut self) -> u16 {
-        let res = self.next_local_port;
-        self.next_local_port = if res >= LOCAL_PORT_MAX { LOCAL_PORT_MIN } else { res + 1 };
-        res
-    }
-}
-
-impl<D: Device + 'static> Inner<D> {
-    fn apply_config(&mut self, s: &mut SocketStack, config: Config) {
-        #[cfg(feature = "medium-ethernet")]
-        let medium = self.device.capabilities().medium;
-
-        debug!("Acquired IP configuration:");
-
-        debug!("   IP address:      {}", config.address);
-        self.set_ipv4_addr(s, config.address);
-
-        #[cfg(feature = "medium-ethernet")]
-        if medium == Medium::Ethernet {
-            if let Some(gateway) = config.gateway {
-                debug!("   Default gateway: {}", gateway);
-                s.iface.routes_mut().add_default_ipv4_route(gateway).unwrap();
-            } else {
-                debug!("   Default gateway: None");
-                s.iface.routes_mut().remove_default_ipv4_route();
-            }
-        }
-        for (i, s) in config.dns_servers.iter().enumerate() {
-            debug!("   DNS server {}:    {}", i, s);
-        }
-
-        self.config = Some(config)
-    }
-
-    #[allow(unused)] // used only with dhcp
-    fn unapply_config(&mut self, s: &mut SocketStack) {
-        #[cfg(feature = "medium-ethernet")]
-        let medium = self.device.capabilities().medium;
-
-        debug!("Lost IP configuration");
-        self.set_ipv4_addr(s, Ipv4Cidr::new(Ipv4Address::UNSPECIFIED, 0));
-        #[cfg(feature = "medium-ethernet")]
-        if medium == Medium::Ethernet {
-            s.iface.routes_mut().remove_default_ipv4_route();
-        }
-        self.config = None
-    }
-
-    fn set_ipv4_addr(&mut self, s: &mut SocketStack, cidr: Ipv4Cidr) {
-        s.iface.update_ip_addrs(|addrs| {
-            let dest = addrs.iter_mut().next().unwrap();
-            *dest = IpCidr::Ipv4(cidr);
-        });
-    }
-
-    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() {
-            // If poll() returns error, it may not be done yet, so poll again later.
-            cx.waker().wake_by_ref();
-            return;
-        }
-
-        // Update link up
-        let old_link_up = self.link_up;
-        self.link_up = self.device.device.link_state() == LinkState::Up;
-
-        // Print when changed
-        if old_link_up != self.link_up {
-            info!("link_up = {:?}", self.link_up);
-        }
-
-        #[cfg(feature = "dhcpv4")]
-        if let Some(dhcp_handle) = self.dhcp_socket {
-            let socket = s.sockets.get_mut::<dhcpv4::Socket>(dhcp_handle);
-
-            if self.link_up {
-                match socket.poll() {
-                    None => {}
-                    Some(dhcpv4::Event::Deconfigured) => self.unapply_config(s),
-                    Some(dhcpv4::Event::Configured(config)) => {
-                        let config = Config {
-                            address: config.address,
-                            gateway: config.router,
-                            dns_servers: config.dns_servers,
-                        };
-                        self.apply_config(s, config)
-                    }
-                }
-            } else if old_link_up {
-                socket.reset();
-                self.unapply_config(s);
-            }
-        }
-        //if old_link_up || self.link_up {
-        //    self.poll_configurator(timestamp)
-        //}
-
-        if let Some(poll_at) = s.iface.poll_at(timestamp, &mut s.sockets) {
-            let t = Timer::at(instant_from_smoltcp(poll_at));
-            pin_mut!(t);
-            if t.poll(cx).is_ready() {
-                cx.waker().wake_by_ref();
-            }
-        }
-    }
-}
-
-fn instant_to_smoltcp(instant: Instant) -> SmolInstant {
-    SmolInstant::from_millis(instant.as_millis() as i64)
-}
-
-fn instant_from_smoltcp(instant: SmolInstant) -> Instant {
-    Instant::from_millis(instant.total_millis() as u64)
-}
diff --git a/embassy-net/src/tcp.rs b/embassy-net/src/tcp.rs
index 73cf2d4e4..0dc8da73a 100644
--- a/embassy-net/src/tcp.rs
+++ b/embassy-net/src/tcp.rs
@@ -8,9 +8,8 @@ use smoltcp::socket::tcp;
 use smoltcp::time::Duration;
 use smoltcp::wire::{IpEndpoint, IpListenEndpoint};
 
-use super::stack::Stack;
-use crate::stack::SocketStack;
-use crate::Device;
+use crate::device::Device;
+use crate::{SocketStack, Stack};
 
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
diff --git a/embassy-net/src/udp.rs b/embassy-net/src/udp.rs
index 4ddad77d4..2f5334df3 100644
--- a/embassy-net/src/udp.rs
+++ b/embassy-net/src/udp.rs
@@ -7,8 +7,7 @@ use smoltcp::iface::{Interface, SocketHandle};
 use smoltcp::socket::udp::{self, PacketMetadata};
 use smoltcp::wire::{IpEndpoint, IpListenEndpoint};
 
-use super::stack::SocketStack;
-use crate::{Device, Stack};
+use crate::{Device, SocketStack, Stack};
 
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]