fix: all warining, clean up code, made it more readable

Author: Bordi00

src/assembler.rs

@@ -1,47 +1,43 @@
+use std::collections::hash_map::Entry::Vacant;
 use std::collections::HashMap;
 
 use wg_internal::{network::NodeId, packet::Fragment};
 
 
-#[derive(Debug)]
-pub struct FragmentAssembler {
-    pub fragments: HashMap<(u64, NodeId), Vec<Fragment>>, // session_id -> data buffer
+#[derive(Debug, Default)]
+pub struct FragmentAssembler<'a> {
+    pub fragments: HashMap<(u64, NodeId), Vec<&'a Fragment>>, // session_id -> data buffer
     pub expected_fragments: HashMap<(u64, NodeId), u64>, // session_id -> total_fragments
     pub received_fragments: HashMap<(u64, NodeId), Vec<bool>>, // session_id -> received status
 }
 
-impl FragmentAssembler {
-    pub fn new() -> Self {
-        Self {
-            fragments: HashMap::new(),
-            expected_fragments: HashMap::new(),
-            received_fragments: HashMap::new()
-        }
-    }
-
-    pub fn add_fragment(&mut self, fragment: Fragment, session_id: u64, sender: NodeId) -> Option<Vec<u8>> {
+impl<'a> FragmentAssembler<'a> {
+    pub fn add_fragment(&mut self, fragment: &'a Fragment, session_id: u64, sender: NodeId) -> Option<Vec<u8>> {
         let communication_id = ( session_id, sender );
-        if !self.fragments.contains_key(&communication_id) {
-            self.fragments.insert(communication_id, vec![fragment.clone()]);
+        #[allow(clippy::cast_possible_truncation)]
+        let index = fragment.fragment_index as usize;
+
+        if let Vacant(entry) = self.fragments.entry(communication_id) {
+            entry.insert(vec![fragment]);
             self.expected_fragments.insert(communication_id, fragment.total_n_fragments);
-            self.received_fragments.insert(communication_id, vec![false; fragment.total_n_fragments as usize]);
+            self.received_fragments.insert(communication_id, vec![false; index]);
         }
 
         {
-            let received = self.received_fragments.get_mut(&communication_id).unwrap();
-            received[fragment.fragment_index as usize] = true;
+            let received = self.received_fragments.get_mut(&communication_id)?;
+            received[index] = true;
         }
 
-        let expected = self.expected_fragments.get(&communication_id).unwrap();
-        let received = self.received_fragments.get(&communication_id).unwrap();
-        let fragments = self.fragments.get(&communication_id).unwrap();
+        let expected = self.expected_fragments.get(&communication_id)?;
+        let received = self.received_fragments.get(&communication_id)?;
+        let fragments = self.fragments.get(&communication_id)?;
 
         // check if all fragments has been received
         if fragments.len() as u64 == *expected && received.iter().all(|f| *f){
-            let fragments = self.fragments.get_mut(&communication_id).unwrap();
+            let fragments = self.fragments.get_mut(&communication_id)?;
             fragments.sort_by(|t, n| t.fragment_index.cmp(&n.fragment_index));
             let mut data = vec![];
-            for f in fragments.iter() {
+            for f in fragments {
                 data.copy_from_slice(&f.data);
             }
             let _ = self.fragments.remove(&communication_id);

src/network.rs

@@ -19,10 +19,10 @@ impl Display for NetworkError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::TopologyError => write!(f, "Topology error"),
-            Self::PathNotFound(id) => write!(f, "Path not found for node {}", id),
-            Self::NodeNotFound(id) => write!(f, "Node {} not found", id),
-            Self::NodeIsNotANeighbor(id) => write!(f, "Node {} is not a neighbor", id),
-            Self::SendError(msg) => write!(f, "Send error: {}", msg),
+            Self::PathNotFound(id) => write!(f, "Path not found for node {id}"),
+            Self::NodeNotFound(id) => write!(f, "Node {id} not found"),
+            Self::NodeIsNotANeighbor(id) => write!(f, "Node {id} is not a neighbor"),
+            Self::SendError(msg) => write!(f, "Send error: {msg}"),
             Self::ControllerDisconnected => write!(f, "Controller disconnected"),
             Self::NoDestination => write!(f, "Packet has no destination specified"),
             Self::NoNeighborAssigned => write!(f, "No neighbor assigned"),
@@ -35,32 +35,34 @@ impl std::error::Error for NetworkError {}
 
 impl<T: Send + std::fmt::Debug> From<SendError<T>> for NetworkError {
     fn from(value: SendError<T>) -> Self {
-        NetworkError::SendError(format!("{:?}", value))
+        NetworkError::SendError(format!("{value:?}"))
     }
 }
 
 
 #[derive(Clone)]
-pub struct Node {
-    id: NodeId,
-    node_type: NodeType,
+pub(crate) struct Node {
+    pub id: NodeId,
+    kind: NodeType,
     adjacents: Vec<NodeId>
 }
 
 
 impl Node {
-    pub fn new(id: NodeId, node_type: NodeType, adjacents: Vec<NodeId>) -> Self {
-        Self { id, node_type, adjacents }
+    #[must_use]
+    pub fn new(id: NodeId, kind: NodeType, adjacents: Vec<NodeId>) -> Self {
+        Self { id, kind, adjacents }
     }
 
     pub fn get_id(&self) -> NodeId {
         self.id
     }
 
     pub fn get_node_type(&self) -> NodeType {
-        self.node_type.clone()
+        self.kind
     }
 
+    #[must_use]
     pub fn get_adjacents(&self) -> &Vec<NodeId> {
         &self.adjacents
     }
@@ -70,8 +72,9 @@ impl Node {
     }
 
     pub fn remove_adjacent(&mut self, adj: NodeId) {
-        let index_to_remove = self.adjacents.iter().position(|i| *i == adj).expect(&format!("Node with id {} not found in {} adjacents", adj, self.id));
-        let _ = self.adjacents.remove(index_to_remove);
+        if let Some(index_to_remove) = self.adjacents.iter().position(|i| *i == adj) {
+            let _ = self.adjacents.remove(index_to_remove);
+        }
     }
 }
 
@@ -83,19 +86,20 @@ impl std::fmt::Debug for Node{
 
 #[derive(Debug, Clone)]
 pub struct Network {
-    pub nodes: Vec<Node>
+    pub(crate) nodes: Vec<Node>
 }
 
 impl Network {
-    pub fn new(root: Node) -> Self {
-        let mut nodes = vec![];
-        nodes.push(root);
+    #[must_use]
+    pub(crate) fn new(root: Node) -> Self {
+        let nodes = vec![root];
         Self { nodes }
     }
 
-    pub fn add_node(&mut self, new_node: Node) {
+
+    pub(crate) fn add_node(&mut self, new_node: Node) {
         for adj in new_node.get_adjacents() {
-            if let Some(node) = self.nodes.iter_mut().find(|n| n.get_id() == *adj) {
+            if let Some(node) = self.nodes.iter_mut().find(|n| n.id == *adj) {
                 match (new_node.get_node_type(), node.get_node_type()) {
                     (_, NodeType::Drone) | (NodeType::Drone, _) => {
                         node.add_adjacent(*adj);
@@ -108,19 +112,22 @@ impl Network {
         self.nodes.push(new_node);
     }
 
-    pub fn remove_node(&mut self, node_id: NodeId) {
-        for n in self.nodes.iter_mut() {
+    pub(crate) fn remove_node(&mut self, node_id: NodeId) {
+        for n in &mut self.nodes{
             if n.get_adjacents().contains(&node_id){
                 n.remove_adjacent(node_id);
             }
         }
-        if let Some(index_to_remove) = self.nodes.iter().position(|n| n.get_id() == node_id) {
+        if let Some(index_to_remove) = self.nodes.iter().position(|n| n.id == node_id) {
             let _ = self.nodes.remove(index_to_remove);
         }
     }
 
-    pub fn update_node(&mut self, node_id: NodeId, adjacents: Vec<NodeId>) -> Result<(), NetworkError> {
-        if let Some(node) = self.nodes.iter_mut().find(|n| n.get_id() == node_id) {
+    /// Updates the node's adjacents with the provided list.
+    /// # Errors
+    /// If the node is not found, returns an error.
+    pub(crate) fn update_node(&mut self, node_id: NodeId, adjacents: Vec<NodeId>) -> Result<(), NetworkError> {
+        if let Some(node) = self.nodes.iter_mut().find(|n| n.id == node_id) {
             for adj in adjacents {
                 if !node.get_adjacents().contains(&adj) {
                     node.add_adjacent(adj);
@@ -131,20 +138,18 @@ impl Network {
             // automatically by the protocol
 
             return Ok(());
-        } else {
-            return Err(NetworkError::NodeNotFound(node_id));
         }
+        Err(NetworkError::NodeNotFound(node_id))
     }
 
-    pub fn change_node_type(&mut self, id: NodeId, new_type: NodeType) {
+    pub(crate) fn change_node_type(&mut self, id: NodeId, new_type: NodeType) {
         if let Some(node) = self.nodes.iter_mut().find(|n| n.get_id() == id) {
-            if node.get_node_type() != new_type {
-                node.node_type = new_type;
-            }
+                node.kind = new_type;
         }
     }
 
-    pub fn find_path(&self, destination: NodeId) -> Option<Vec<NodeId>> {
+    #[must_use]
+    pub(crate) fn find_path(&self, destination: NodeId) -> Option<Vec<NodeId>> {
         let start = self.nodes[0].id;
         let mut visited = HashSet::new();
         let mut queue = VecDeque::new();
@@ -165,8 +170,8 @@ impl Network {
                 return Some(path);
             }
 
-            if let Some(node) = self.nodes.iter().find(|n| n.get_id() == current) {
-                for neighbor in node.get_adjacents().iter() {
+            if let Some(node) = self.nodes.iter().find(|n| n.id == current) {
+                for neighbor in node.get_adjacents() {
                     if !visited.contains(neighbor) {
                         visited.insert(*neighbor);
                         parent_map.insert(neighbor, current);
@@ -192,7 +197,7 @@ mod tests {
         network.add_node(new_node);
 
         assert_eq!(network.nodes.len(), 2);
-        assert!(network.nodes.iter().any(|n| n.get_id() == 2));
+        assert!(network.nodes.iter().any(|n| n.id == 2));
     }
 
     #[test]
@@ -206,7 +211,7 @@ mod tests {
         network.remove_node(2);
 
         assert_eq!(network.nodes.len(), 1);
-        assert!(!network.nodes.iter().any(|n| n.get_id() == 2));
+        assert!(!network.nodes.iter().any(|n| n.id == 2));
     }
 
     #[test]
@@ -260,7 +265,6 @@ mod tests {
 
         let result = network.find_path(3);
 
-        assert!(result.is_err());
-        assert_eq!(result.unwrap_err().to_string(), "Path not found for node 3");
+        assert!(result.is_none());
     }
 }