Convert struct Defragmenter into Defragment trait.

Author: neumann-paulmann

src/defragmentation.rs

@@ -4,20 +4,17 @@ use log::{error, trace, warn};
 
 pub use self::defragmentation_error::DefragmentationError;
 pub use self::defragmented_message::DefragmentedMessage;
-use self::transaction::Transaction;
+pub use self::transaction::Transaction;
 use crate::parameters::messaging::handler::IncomingMessage;
 
 mod defragmentation_error;
 mod defragmented_message;
 mod transaction;
 
-/// Defragments potentially fragmented EZSP frames.
-#[derive(Clone, Debug, Default, Eq, PartialEq, Hash)]
-pub struct Defragmenter<const BACKPRESSURE_THRESHOLD: usize> {
-    inner: BTreeMap<u8, Transaction>,
-}
+const BACKPRESSURE_THRESHOLD: usize = 2;
 
-impl<const BACKPRESSURE_THRESHOLD: usize> Defragmenter<BACKPRESSURE_THRESHOLD> {
+/// Defragment potentially fragmented EZSP frames.
+pub trait Defragment {
     /// Attempt to defragment the given frame.
     ///
     /// # Returns
@@ -30,14 +27,21 @@ impl<const BACKPRESSURE_THRESHOLD: usize> Defragmenter<BACKPRESSURE_THRESHOLD> {
     ///
     /// - [`DefragmentationError::StrayFragment`] if a follow-up fragment is received without an initial fragment.
     /// - [`DefragmentationError::DuplicateFragment`] if a duplicate fragment is received.
-    pub fn defragment(
+    fn defragment(
+        &mut self,
+        incoming_message: IncomingMessage,
+    ) -> Result<Option<DefragmentedMessage>, DefragmentationError>;
+}
+
+impl Defragment for BTreeMap<u8, Transaction> {
+    fn defragment(
         &mut self,
         incoming_message: IncomingMessage,
     ) -> Result<Option<DefragmentedMessage>, DefragmentationError> {
-        if self.inner.len() > BACKPRESSURE_THRESHOLD {
+        if self.len() > BACKPRESSURE_THRESHOLD {
             warn!(
                 "Defragmenter backpressure threshold exceeded: {} transactions in progress.",
-                self.inner.len()
+                self.len()
             );
         }
 
@@ -48,8 +52,7 @@ impl<const BACKPRESSURE_THRESHOLD: usize> Defragmenter<BACKPRESSURE_THRESHOLD> {
                 trace!("Received initial fragment.");
 
                 if size > 1 {
-                    self.inner
-                        .insert(seq, Transaction::new(size, index, incoming_message));
+                    self.insert(seq, Transaction::new(size, index, incoming_message));
                     return Ok(None);
                 }
 
@@ -59,13 +62,13 @@ impl<const BACKPRESSURE_THRESHOLD: usize> Defragmenter<BACKPRESSURE_THRESHOLD> {
             Some((index, None)) => {
                 trace!("Received follow-up fragment.");
 
-                let Some(transaction) = self.inner.get_mut(&seq) else {
+                let Some(transaction) = self.get_mut(&seq) else {
                     return Err(DefragmentationError::StrayFragment { seq, index });
                 };
 
                 if let Some(message) = transaction.update(index, incoming_message) {
                     error!("Duplicate fragment received. Previous was: {message:?}");
-                    self.inner.remove(&seq);
+                    self.remove(&seq);
                     return Err(DefragmentationError::DuplicateFragment { seq, index });
                 }
 
@@ -77,7 +80,7 @@ impl<const BACKPRESSURE_THRESHOLD: usize> Defragmenter<BACKPRESSURE_THRESHOLD> {
                     return Ok(None);
                 };
 
-                self.inner.remove(&seq);
+                self.remove(&seq);
                 Ok(Some(defragmented_message))
             }
             None => {

src/lib.rs

@@ -12,7 +12,7 @@ pub use self::commands::{
     Networking, ProxyTable, Security, SinkTable, TokenInterface, TrustCenter, Utilities, Wwah, Zll,
 };
 pub use self::constants::{MAX_HEADER_SIZE, MAX_PARAMETER_SIZE, MIN_NON_LEGACY_VERSION};
-pub use self::defragmentation::{DefragmentationError, DefragmentedMessage, Defragmenter};
+pub use self::defragmentation::{Defragment, DefragmentationError, DefragmentedMessage};
 pub use self::error::{Error, ValueError};
 pub use self::extensions::{ConfigurationExt, Displayable, PolicyExt};
 pub use self::frame::{

src/zigbee/network_manager/message_handler.rs

@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use std::sync::Arc;
 
 use log::{debug, error, trace, warn};
@@ -7,7 +8,7 @@ use tokio::sync::mpsc::{Receiver, Sender, channel};
 use zigbee::Endpoint;
 use zigbee_nwk::Event;
 
-use crate::defragmentation::Defragmenter;
+use crate::defragmentation::{Defragment, Transaction};
 use crate::ember::device::Update;
 use crate::frame::parameters::networking::handler::Handler as Networking;
 use crate::parameters::messaging::handler::{Handler as Messaging, IncomingMessage, MessageSent};
@@ -23,7 +24,7 @@ pub type Handlers = Arc<Mutex<Vec<Sender<Callback>>>>;
 pub struct MessageHandler {
     handlers: Handlers,
     outgoing: Sender<Event>,
-    transactions: Defragmenter<2>,
+    transactions: BTreeMap<u8, Transaction>,
 }
 
 impl MessageHandler {
@@ -35,7 +36,7 @@ impl MessageHandler {
             Self {
                 handlers: handlers.clone(),
                 outgoing,
-                transactions: Defragmenter::default(),
+                transactions: BTreeMap::new(),
             },
             handlers,
             rx,