Finished some reorganizing

src/display/components/display_bandwidth.rs

@@ -1,22 +1,21 @@
 use ::std::fmt;
 
-// Should we make a `DisplayMode` enum with `Rate` and `Total`?
-// Instead of the ambiguous bool?
-
-// This might be better as a traditional struct now?
-pub struct DisplayBandwidth(pub f64, pub bool);
+pub struct DisplayBandwidth {
+    pub bandwidth: f64,
+    pub as_rate: bool,
+}
 
 impl fmt::Display for DisplayBandwidth {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let suffix = if self.1 { "" } else { "ps" }; // Should this be the other way around?
-        if self.0 > 999_999_999.0 {
-            write!(f, "{:.2}GB{}", self.0 / 1_000_000_000.0, suffix)
-        } else if self.0 > 999_999.0 {
-            write!(f, "{:.2}MB{}", self.0 / 1_000_000.0, suffix)
-        } else if self.0 > 999.0 {
-            write!(f, "{:.2}KB{}", self.0 / 1000.0, suffix)
+        let suffix = if self.as_rate { "" } else { "ps" };
+        if self.bandwidth > 999_999_999.0 {
+            write!(f, "{:.2}GB{}", self.bandwidth / 1_000_000_000.0, suffix)
+        } else if self.bandwidth > 999_999.0 {
+            write!(f, "{:.2}MB{}", self.bandwidth / 1_000_000.0, suffix)
+        } else if self.bandwidth > 999.0 {
+            write!(f, "{:.2}KB{}", self.bandwidth / 1000.0, suffix)
         } else {
-            write!(f, "{}B{}", self.0, suffix)
+            write!(f, "{}B{}", self.bandwidth, suffix)
         }
     }
 }

src/display/components/table.rs

@@ -9,46 +9,22 @@ use ::tui::widgets::{Block, Borders, Row, Widget};
 use crate::display::{Bandwidth, DisplayBandwidth, UIState};
 use crate::network::{display_connection_string, display_ip_or_host};
 
-use ::std::iter::FromIterator;
 use ::std::net::IpAddr;
 
 fn display_upload_and_download(bandwidth: &impl Bandwidth, total: bool) -> String {
     format!(
         "{} / {}",
-        DisplayBandwidth(bandwidth.get_total_bytes_uploaded() as f64, total),
-        DisplayBandwidth(bandwidth.get_total_bytes_downloaded() as f64, total)
+        DisplayBandwidth{
+            bandwidth: bandwidth.get_total_bytes_uploaded() as f64,
+            as_rate: !total,
+        },
+        DisplayBandwidth{
+            bandwidth: bandwidth.get_total_bytes_downloaded() as f64,
+            as_rate: !total,
+        },
     )
 }
 
-// Why is this returning the sorted list when it mutates things?
-fn sort_by_bandwidth<'a, T>(
-    list: &'a mut Vec<(T, &impl Bandwidth)>,
-) -> &'a Vec<(T, &'a impl Bandwidth)> {
-    list.sort_by(|(_, a), (_, b)| {
-        let a_highest = if a.get_total_bytes_downloaded() > a.get_total_bytes_uploaded() {
-            a.get_total_bytes_downloaded()
-        } else {
-            a.get_total_bytes_uploaded()
-        };
-        let b_highest = if b.get_total_bytes_downloaded() > b.get_total_bytes_uploaded() {
-            b.get_total_bytes_downloaded()
-        } else {
-            b.get_total_bytes_uploaded()
-        };
-        b_highest.cmp(&a_highest)
-    });
-    list
-}
-// I'd suggest using this:
-// fn sort_by_bandwidth<T>(list: &mut Vec<(T, impl Bandwidth)>) {
-//     list.sort_by_key(|(_, b)| {
-//         cmp::Reverse(cmp::max(
-//             b.get_total_bytes_downloaded(),
-//             b.get_total_bytes_uploaded(),
-//         ))
-//     });
-// }
-
 pub enum ColumnCount {
     Two,
     Three,
@@ -87,9 +63,8 @@ fn truncate_middle(row: &str, max_length: u16) -> String {
 
 impl<'a> Table<'a> {
     pub fn create_connections_table(state: &UIState, ip_to_host: &HashMap<IpAddr, String>) -> Self {
-        let mut connections_list = Vec::from_iter(&state.connections);
-        sort_by_bandwidth(&mut connections_list);
-        let connections_rows = connections_list
+        let connections_rows = state
+            .connections
             .iter()
             .map(|(connection, connection_data)| {
                 vec![
@@ -99,7 +74,7 @@ impl<'a> Table<'a> {
                         &connection_data.interface_name,
                     ),
                     connection_data.process_name.to_string(),
-                    display_upload_and_download(*connection_data, state.cumulative_mode),
+                    display_upload_and_download(connection_data, state.cumulative_mode),
                 ]
             })
             .collect();
@@ -142,15 +117,14 @@ impl<'a> Table<'a> {
         }
     }
     pub fn create_processes_table(state: &UIState) -> Self {
-        let mut processes_list = Vec::from_iter(&state.processes);
-        sort_by_bandwidth(&mut processes_list);
-        let processes_rows = processes_list
+        let processes_rows = state
+            .processes
             .iter()
             .map(|(process_name, data_for_process)| {
                 vec![
                     (*process_name).to_string(),
                     data_for_process.connection_count.to_string(),
-                    display_upload_and_download(*data_for_process, state.cumulative_mode),
+                    display_upload_and_download(data_for_process, state.cumulative_mode),
                 ]
             })
             .collect();
@@ -196,16 +170,15 @@ impl<'a> Table<'a> {
         state: &UIState,
         ip_to_host: &HashMap<IpAddr, String>,
     ) -> Self {
-        let mut remote_addresses_list = Vec::from_iter(&state.remote_addresses);
-        sort_by_bandwidth(&mut remote_addresses_list);
-        let remote_addresses_rows = remote_addresses_list
+        let remote_addresses_rows = state
+            .remote_addresses
             .iter()
             .map(|(remote_address, data_for_remote_address)| {
-                let remote_address = display_ip_or_host(**remote_address, &ip_to_host);
+                let remote_address = display_ip_or_host(*remote_address, &ip_to_host);
                 vec![
                     remote_address,
                     data_for_remote_address.connection_count.to_string(),
-                    display_upload_and_download(*data_for_remote_address, state.cumulative_mode),
+                    display_upload_and_download(data_for_remote_address, state.cumulative_mode),
                 ]
             })
             .collect();

src/display/components/total_bandwidth.rs

@@ -25,8 +25,14 @@ impl<'a> TotalBandwidth<'a> {
             [Text::styled(
                 format!(
                     " Total Up / Down: {} / {} {}",
-                    DisplayBandwidth(self.state.total_bytes_uploaded as f64, c_mode),
-                    DisplayBandwidth(self.state.total_bytes_downloaded as f64, c_mode),
+                    DisplayBandwidth{
+                        bandwidth: self.state.total_bytes_uploaded as f64,
+                        as_rate: !c_mode,
+                    },
+                    DisplayBandwidth{
+                        bandwidth: self.state.total_bytes_downloaded as f64,
+                        as_rate: !c_mode,
+                    },
                     paused_str
                 ),
                 Style::default().fg(color).modifier(Modifier::BOLD),

src/display/ui.rs

@@ -30,12 +30,11 @@ where
         let mut terminal = Terminal::new(terminal_backend).unwrap();
         terminal.clear().unwrap();
         terminal.hide_cursor().unwrap();
+        let mut state: UIState = Default::default();
+        state.cumulative_mode = opts.total_utilization;
         Ui {
             terminal,
-            state: UIState {
-                cumulative_mode: opts.total_utilization,
-                ..Default::default()
-            },
+            state,
             ip_to_host: Default::default(),
             opts,
         }

src/display/ui_state.rs

@@ -1,5 +1,6 @@
 use ::std::cmp;
-use ::std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
+use ::std::collections::{HashMap, HashSet, VecDeque};
+use ::std::hash::Hash;
 use ::std::iter::FromIterator;
 use ::std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 
@@ -11,7 +12,8 @@ static MAX_BANDWIDTH_ITEMS: usize = 1000;
 pub trait Bandwidth {
     fn get_total_bytes_downloaded(&self) -> u128;
     fn get_total_bytes_uploaded(&self) -> u128;
-    fn combine_bandwidth(&mut self, other: &impl Bandwidth);
+    fn combine_bandwidth(&mut self, other: &Self);
+    fn divide_by(&mut self, amount: u128);
 }
 
 #[derive(Clone, Default)]
@@ -29,44 +31,39 @@ pub struct ConnectionData {
     pub interface_name: String,
 }
 
-impl NetworkData {
-    pub fn divide_by(&mut self, amount: u128) {
-        self.total_bytes_downloaded /= amount;
-        self.total_bytes_uploaded /= amount;
-    }
-}
-
-impl ConnectionData {
-    pub fn divide_by(&mut self, amount: u128) {
-        self.total_bytes_downloaded /= amount;
-        self.total_bytes_uploaded /= amount;
-    }
-}
-
-impl Bandwidth for ConnectionData {
+impl Bandwidth for NetworkData {
     fn get_total_bytes_downloaded(&self) -> u128 {
         self.total_bytes_downloaded
     }
     fn get_total_bytes_uploaded(&self) -> u128 {
         self.total_bytes_uploaded
     }
-    fn combine_bandwidth(&mut self, other: &impl Bandwidth) {
+    fn combine_bandwidth(&mut self, other: &NetworkData) {
         self.total_bytes_downloaded += other.get_total_bytes_downloaded();
         self.total_bytes_uploaded += other.get_total_bytes_uploaded();
+        self.connection_count = other.connection_count;
+    }
+    fn divide_by(&mut self, amount: u128) {
+        self.total_bytes_downloaded /= amount;
+        self.total_bytes_uploaded /= amount;
     }
 }
 
-impl Bandwidth for NetworkData {
+impl Bandwidth for ConnectionData {
     fn get_total_bytes_downloaded(&self) -> u128 {
         self.total_bytes_downloaded
     }
     fn get_total_bytes_uploaded(&self) -> u128 {
         self.total_bytes_uploaded
     }
-    fn combine_bandwidth(&mut self, other: &impl Bandwidth) {
+    fn combine_bandwidth(&mut self, other: &ConnectionData) {
         self.total_bytes_downloaded += other.get_total_bytes_downloaded();
         self.total_bytes_uploaded += other.get_total_bytes_uploaded();
     }
+    fn divide_by(&mut self, amount: u128) {
+        self.total_bytes_downloaded /= amount;
+        self.total_bytes_uploaded /= amount;
+    }
 }
 
 pub struct UtilizationData {
@@ -76,20 +73,16 @@ pub struct UtilizationData {
 
 #[derive(Default)]
 pub struct UIState {
-    // We aren't really taking advantage of the self-sorting of the BTreeMap here,
-    // but it would be rather difficult to, as it sorts by key, not value. Perhaps we
-    // should just use a HashMap?
-    // We could also use a HashMap internally for insertions and merging, but cache a
-    // sorted vector so that doesn't need sorting during the display step. This would
-    // also fix Bandwhich eating CPU cycles when paused, but is a bit of a non-issue
-    // with `MAX_BANDWIDTH_ITEMS` in place.
-    pub processes: BTreeMap<String, NetworkData>,
-    pub remote_addresses: BTreeMap<IpAddr, NetworkData>,
-    pub connections: BTreeMap<Connection, ConnectionData>,
+    pub processes: Vec<(String, NetworkData)>,
+    pub remote_addresses: Vec<(IpAddr, NetworkData)>,
+    pub connections: Vec<(Connection, ConnectionData)>,
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
     pub cumulative_mode: bool,
-    pub utilization_data: VecDeque<UtilizationData>,
+    utilization_data: VecDeque<UtilizationData>,
+    processes_map: HashMap<String, NetworkData>,
+    remote_addresses_map: HashMap<IpAddr, NetworkData>,
+    connections_map: HashMap<Connection, ConnectionData>,
 }
 
 impl UIState {
@@ -125,9 +118,9 @@ impl UIState {
         if self.utilization_data.len() > RECALL_LENGTH {
             self.utilization_data.pop_front();
         }
-        let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new();
-        let mut remote_addresses: BTreeMap<IpAddr, NetworkData> = BTreeMap::new();
-        let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new();
+        let mut processes: HashMap<String, NetworkData> = HashMap::new();
+        let mut remote_addresses: HashMap<IpAddr, NetworkData> = HashMap::new();
+        let mut connections: HashMap<Connection, ConnectionData> = HashMap::new();
         let mut total_bytes_downloaded: u128 = 0;
         let mut total_bytes_uploaded: u128 = 0;
 
@@ -190,27 +183,27 @@ impl UIState {
         }
 
         if self.cumulative_mode {
-            merge_bandwidth(&mut self.processes, processes);
-            merge_bandwidth(&mut self.remote_addresses, remote_addresses);
-            merge_bandwidth(&mut self.connections, connections);
+            merge_bandwidth(&mut self.processes_map, processes);
+            merge_bandwidth(&mut self.remote_addresses_map, remote_addresses);
+            merge_bandwidth(&mut self.connections_map, connections);
             self.total_bytes_downloaded += total_bytes_downloaded / divide_by;
             self.total_bytes_uploaded += total_bytes_uploaded / divide_by;
         } else {
-            self.processes = processes;
-            self.remote_addresses = remote_addresses;
-            self.connections = connections;
+            self.processes_map = processes;
+            self.remote_addresses_map = remote_addresses;
+            self.connections_map = connections;
             self.total_bytes_downloaded = total_bytes_downloaded / divide_by;
             self.total_bytes_uploaded = total_bytes_uploaded / divide_by;
         }
-        prune_map(&mut self.processes);
-        prune_map(&mut self.remote_addresses);
-        prune_map(&mut self.connections);
+        self.processes = sort_and_prune(&mut self.processes_map);
+        self.remote_addresses = sort_and_prune(&mut self.remote_addresses_map);
+        self.connections = sort_and_prune(&mut self.connections_map);
     }
 }
 
-fn merge_bandwidth<K, V>(self_map: &mut BTreeMap<K, V>, other_map: BTreeMap<K, V>)
+fn merge_bandwidth<K, V>(self_map: &mut HashMap<K, V>, other_map: HashMap<K, V>)
 where
-    K: Eq + Ord,
+    K: Eq + Hash,
     V: Bandwidth,
 {
     for (key, b_other) in other_map {
@@ -221,22 +214,21 @@ where
     }
 }
 
-fn prune_map(map: &mut BTreeMap<impl Eq + Ord + Clone, impl Bandwidth + Clone>) {
-    if map.len() > MAX_BANDWIDTH_ITEMS {
-        let mut bandwidth_list = Vec::from_iter(map.clone());
-        sort_by_bandwidth(&mut bandwidth_list);
+fn sort_and_prune<K, V>(map: &mut HashMap<K, V>) -> Vec<(K, V)>
+where
+    K: Eq + Hash + Clone,
+    V: Bandwidth + Clone,
+{
+    let mut bandwidth_list = Vec::from_iter(map.clone());
+    bandwidth_list.sort_by_key(|(_, b)| {
+        cmp::Reverse(b.get_total_bytes_downloaded() + b.get_total_bytes_uploaded())
+    });
+
+    if bandwidth_list.len() > MAX_BANDWIDTH_ITEMS {
         for (key, _) in &bandwidth_list[MAX_BANDWIDTH_ITEMS..] {
             map.remove(key);
         }
     }
-}
 
-// This is duplicated from table.rs temporarily
-fn sort_by_bandwidth<T>(list: &mut Vec<(T, impl Bandwidth)>) {
-    list.sort_by_key(|(_, b)| {
-        cmp::Reverse(cmp::max(
-            b.get_total_bytes_downloaded(),
-            b.get_total_bytes_uploaded(),
-        ))
-    });
+    bandwidth_list
 }