Refactoring device state management

src/device.rs

@@ -1,4 +1,4 @@
-use bit_field::BitArray;
+use bit_field::BitField;
 
 use crate::bus::{Bus, FourWire, ThreeWire};
 use crate::host::Host;
@@ -23,35 +23,93 @@ impl<E> From<E> for ResetError<E> {
     }
 }
 
+pub trait State {
+    type Host: Host;
+
+    fn socket(&mut self) -> Option<Socket>;
+    fn release_socket(&mut self, socket: Socket);
+    fn any_allocated(&self) -> bool;
+    fn host(&self) -> &Self::Host;
+}
+
 #[derive(Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub struct DeviceState<HostImpl: Host> {
     host: HostImpl,
-    sockets: [u8; 1],
+    sockets: u8,
 }
 
-pub struct Device<SpiBus: Bus, HostImpl: Host> {
+impl<HostImpl: Host> DeviceState<HostImpl> {
-    pub(crate) bus: SpiBus,
+    pub fn new(host: HostImpl) -> Self {
-    state: DeviceState<HostImpl>,
+        Self {
-}
+            sockets: 0xFF,
-
+            host,
-impl<SpiBus: Bus, HostImpl: Host> Device<SpiBus, HostImpl> {
-    pub(crate) fn new(bus: SpiBus, host: HostImpl) -> Self {
-        Device {
-            bus,
-            state: DeviceState {
-                host,
-                sockets: [0b11111111],
-            },
         }
     }
+}
 
-    pub fn get_state(&self) -> &DeviceState<HostImpl> {
+impl<HostImpl: Host> State for DeviceState<HostImpl> {
+    type Host = HostImpl;
+
+    fn socket(&mut self) -> Option<Socket> {
+        for index in 0..8 {
+            if self.sockets.get_bit(index) {
+                self.sockets.set_bit(index, false);
+                return Some(Socket::new(index as u8));
+            }
+        }
+        None
+    }
+
+    fn release_socket(&mut self, socket: Socket) {
+        self.sockets.set_bit(socket.index.into(), true);
+    }
+
+    fn host(&self) -> &Self::Host {
+        &self.host
+    }
+
+    fn any_allocated(&self) -> bool {
+        self.sockets != 0xFF
+    }
+}
+
+impl<'a, T: State> State for &'a mut T {
+    type Host = T::Host;
+
+    fn socket(&mut self) -> Option<Socket> {
+        T::socket(self)
+    }
+
+    fn release_socket(&mut self, socket: Socket) {
+        T::release_socket(self, socket)
+    }
+
+    fn host(&self) -> &Self::Host {
+        T::host(self)
+    }
+
+    fn any_allocated(&self) -> bool {
+        T::any_allocated(self)
+    }
+}
+
+pub struct Device<SpiBus: Bus, StateImpl: State> {
+    pub(crate) bus: SpiBus,
+    state: StateImpl,
+}
+
+impl<SpiBus: Bus, StateImpl: State> Device<SpiBus, StateImpl> {
+    pub(crate) fn new(bus: SpiBus, state: StateImpl) -> Self {
+        Device { bus, state }
+    }
+
+    pub fn get_state(&self) -> &StateImpl {
         &self.state
     }
 
     pub fn reset(mut self) -> Result<UninitializedDevice<SpiBus>, ResetError<SpiBus::Error>> {
-        if self.state.sockets != [0b11111111] {
+        if self.state.any_allocated() {
             Err(ResetError::SocketsNotReleased)
         } else {
             self.reset_device()?;
@@ -59,23 +117,17 @@ impl<SpiBus: Bus, HostImpl: Host> Device<SpiBus, HostImpl> {
         }
     }
 
-    pub fn release(self) -> (SpiBus, HostImpl) {
+    pub fn release(self) -> (SpiBus, StateImpl) {
-        (self.bus, self.state.host)
+        (self.bus, self.state)
     }
 
     pub fn take_socket(&mut self) -> Option<Socket> {
         // TODO maybe return Future that resolves when release_socket invoked
-        for index in 0..8 {
+        self.state.socket()
-            if self.state.sockets.get_bit(index) {
-                self.state.sockets.set_bit(index, false);
-                return Some(Socket::new(index as u8));
-            }
-        }
-        None
     }
 
     pub fn release_socket(&mut self, socket: Socket) {
-        self.state.sockets.set_bit(socket.index.into(), true);
+        self.state.release_socket(socket)
     }
 
     pub fn gateway(&mut self) -> Result<Ipv4Addr, SpiBus::Error> {
@@ -230,18 +282,25 @@ impl<SpiBus: Bus, HostImpl: Host> Device<SpiBus, HostImpl> {
         Ok(retry_count_register[0])
     }
 
-    pub fn deactivate(self) -> (SpiBus, InactiveDevice<HostImpl>) {
+    pub fn deactivate(self) -> (SpiBus, InactiveDevice<StateImpl>) {
         (self.bus, InactiveDevice(self.state))
     }
 }
 
 #[derive(Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
-pub struct InactiveDevice<HostImpl: Host>(DeviceState<HostImpl>);
+pub struct InactiveDevice<StateImpl: State>(StateImpl);
 
-impl<HostImpl: Host> InactiveDevice<HostImpl> {
+impl<StateImpl: State> InactiveDevice<StateImpl> {
     /// Activates the device by taking ownership
-    pub fn activate<SpiBus: Bus>(self, bus: SpiBus) -> Device<SpiBus, HostImpl> {
+    pub fn activate<SpiBus: Bus>(self, bus: SpiBus) -> Device<SpiBus, StateImpl> {
         Device { bus, state: self.0 }
     }
+
+    pub fn activate_ref<SpiBus: Bus>(&mut self, bus: SpiBus) -> Device<SpiBus, &mut StateImpl> {
+        Device {
+            bus,
+            state: &mut self.0,
+        }
+    }
 }

src/tcp.rs

@@ -1,4 +1,9 @@
-use crate::{bus::Bus, device::Device, host::Host, register::socketn, socket::Socket};
+use crate::{
+    bus::Bus,
+    device::{Device, State},
+    register::socketn,
+    socket::Socket,
+};
 
 use embedded_nal::{nb, IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4, TcpClientStack};
 
@@ -200,7 +205,7 @@ impl TcpSocket {
     }
 }
 
-impl<SpiBus: Bus, HostImpl: Host> TcpClientStack for Device<SpiBus, HostImpl> {
+impl<SpiBus: Bus, StateImpl: State> TcpClientStack for Device<SpiBus, StateImpl> {
     type TcpSocket = TcpSocket;
     type Error = TcpSocketError<SpiBus::Error>;
 

src/udp.rs

@@ -4,8 +4,7 @@ use embedded_nal::{nb, IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4, UdpClientStac
 
 use crate::{
     bus::Bus,
-    device::Device,
+    device::{Device, State},
-    host::Host,
     register::socketn::{self, Status},
     socket::Socket,
 };
@@ -470,10 +469,10 @@ impl<E: Debug> From<NbError<E>> for nb::Error<E> {
     }
 }
 
-impl<SpiBus, HostImpl> UdpClientStack for Device<SpiBus, HostImpl>
+impl<SpiBus, StateImpl> UdpClientStack for Device<SpiBus, StateImpl>
 where
     SpiBus: Bus,
-    HostImpl: Host,
+    StateImpl: State,
 {
     type UdpSocket = UdpSocket;
     type Error = UdpSocketError<SpiBus::Error>;
@@ -521,10 +520,10 @@ where
     }
 }
 
-impl<SpiBus, HostImpl> UdpFullStack for Device<SpiBus, HostImpl>
+impl<SpiBus, StateImpl> UdpFullStack for Device<SpiBus, StateImpl>
 where
     SpiBus: Bus,
-    HostImpl: Host,
+    StateImpl: State,
 {
     fn bind(&mut self, socket: &mut Self::UdpSocket, local_port: u16) -> Result<(), Self::Error> {
         socket.set_port(&mut self.bus, local_port)?;

src/uninitialized_device.rs

@@ -2,7 +2,7 @@ use embedded_hal::spi::SpiDevice;
 use embedded_nal::Ipv4Addr;
 
 use crate::bus::{Bus, FourWire, ThreeWire};
-use crate::device::Device;
+use crate::device::{Device, DeviceState};
 use crate::host::{Dhcp, Host, Manual};
 use crate::raw_device::RawDevice;
 use crate::{
@@ -54,7 +54,7 @@ impl<SpiBus: Bus> UninitializedDevice<SpiBus> {
         self,
         mac: MacAddress,
         mode_options: Mode,
-    ) -> Result<Device<SpiBus, Dhcp>, InitializeError<SpiBus::Error>> {
+    ) -> Result<Device<SpiBus, DeviceState<Dhcp>>, InitializeError<SpiBus::Error>> {
         let host = Dhcp::new(mac);
         self.initialize_with_host(host, mode_options)
     }
@@ -67,7 +67,7 @@ impl<SpiBus: Bus> UninitializedDevice<SpiBus> {
         mac: MacAddress,
         ip: Ipv4Addr,
         mode_options: Mode,
-    ) -> Result<Device<SpiBus, Manual>, InitializeError<SpiBus::Error>> {
+    ) -> Result<Device<SpiBus, DeviceState<Manual>>, InitializeError<SpiBus::Error>> {
         let mut ip_bytes = ip.octets();
         ip_bytes[3] = 1;
         let gateway = Ipv4Addr::from(ip_bytes);
@@ -82,7 +82,7 @@ impl<SpiBus: Bus> UninitializedDevice<SpiBus> {
         gateway: Ipv4Addr,
         subnet: Ipv4Addr,
         mode_options: Mode,
-    ) -> Result<Device<SpiBus, Manual>, InitializeError<SpiBus::Error>> {
+    ) -> Result<Device<SpiBus, DeviceState<Manual>>, InitializeError<SpiBus::Error>> {
         let host = Manual::new(mac, ip, gateway, subnet);
         self.initialize_with_host(host, mode_options)
     }
@@ -91,7 +91,7 @@ impl<SpiBus: Bus> UninitializedDevice<SpiBus> {
         mut self,
         mut host: HostImpl,
         mode_options: Mode,
-    ) -> Result<Device<SpiBus, HostImpl>, InitializeError<SpiBus::Error>> {
+    ) -> Result<Device<SpiBus, DeviceState<HostImpl>>, InitializeError<SpiBus::Error>> {
         #[cfg(not(feature = "no-chip-version-assertion"))]
         self.assert_chip_version(0x4)?;
 
@@ -100,7 +100,7 @@ impl<SpiBus: Bus> UninitializedDevice<SpiBus> {
 
         self.set_mode(mode_options)?;
         host.refresh(&mut self.bus)?;
-        Ok(Device::new(self.bus, host))
+        Ok(Device::new(self.bus, DeviceState::new(host)))
     }
 
     pub fn initialize_macraw(