suggestion for previous state storage and calculation

2 files changed, 72 insertions, 51 deletions

diff --git a/src/display/ui.rs b/src/display/ui.rs
index 62dfb2a..dd8ecfb 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -98,8 +98,8 @@ where
         utilization: Utilization,
         ip_to_host: HashMap<Ipv4Addr, String>,
     ) {
-        self.state = UIState::new(connections_to_procs, utilization, &self.state);
-        self.ip_to_host = ip_to_host;
+        self.state.update(connections_to_procs, utilization);
+        self.ip_to_host.extend(ip_to_host);
     }
     pub fn end(&mut self) {
         self.terminal.clear().unwrap();
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index a5bdfcb..4b42039 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -1,9 +1,10 @@
-use ::std::collections::{BTreeMap, HashMap};
+use ::std::collections::{BTreeMap, HashMap, VecDeque};
 use ::std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 
 use crate::network::{Connection, LocalSocket, Utilization};
 
 static BANDWIDTH_DECAY_FACTOR: f32 = 0.5;
+static RECALL_LENGTH: usize = 5;
 
 pub trait Bandwidth {
     fn get_total_bytes_downloaded(&self) -> u128;
@@ -32,6 +33,20 @@ pub struct ConnectionData {
     pub interface_name: String,
 }
 
+impl NetworkData {
+    pub fn divide_by(&mut self, amount: usize) {
+        self.total_bytes_downloaded = self.total_bytes_downloaded / amount as u128;
+        self.total_bytes_uploaded = self.total_bytes_uploaded / amount as u128;
+    }
+}
+
+impl ConnectionData {
+    pub fn divide_by(&mut self, amount: usize) {
+        self.total_bytes_downloaded = self.total_bytes_downloaded / amount as u128;
+        self.total_bytes_uploaded = self.total_bytes_uploaded / amount as u128;
+    }
+}
+
 fn calc_avg_bandwidth(prev_bandwidth: u128, curr_bandwidth: u128) -> u128 {
     if prev_bandwidth == 0 {
         curr_bandwidth
@@ -77,6 +92,11 @@ impl Bandwidth for NetworkData {
     }
 }
 
+pub struct UtilizationData {
+    connections_to_procs: HashMap<LocalSocket, String>,
+    network_utilization: Utilization,
+}
+
 #[derive(Default)]
 pub struct UIState {
     pub processes: BTreeMap<String, NetworkData>,
@@ -84,7 +104,8 @@ pub struct UIState {
     pub connections: BTreeMap<Connection, ConnectionData>,
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
-}
+    utilization_data: VecDeque<UtilizationData>,
+ }
 
 impl UIState {
     fn get_proc_name<'a>(
@@ -107,64 +128,64 @@ impl UIState {
             })
         }
     }
-    pub fn new(
+    pub fn update(
+        &mut self,
         connections_to_procs: HashMap<LocalSocket, String>,
         network_utilization: Utilization,
-        old_state: &UIState,
-    ) -> Self {
+    ) {
+        self.utilization_data.push_back(UtilizationData { connections_to_procs, network_utilization });
+        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<Ipv4Addr, NetworkData> = BTreeMap::new();
         let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new();
         let mut total_bytes_downloaded: u128 = 0;
         let mut total_bytes_uploaded: u128 = 0;
-        for (connection, connection_info) in network_utilization.connections {
-            let connection_data = connections.entry(connection).or_default();
-            let data_for_remote_address = remote_addresses
-                .entry(connection.remote_socket.ip)
-                .or_default();
-            connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded;
-            connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded;
-            connection_data.interface_name = connection_info.interface_name;
-            data_for_remote_address.total_bytes_downloaded +=
-                connection_info.total_bytes_downloaded;
-            data_for_remote_address.total_bytes_uploaded += connection_info.total_bytes_uploaded;
-            data_for_remote_address.connection_count += 1;
-            total_bytes_downloaded += connection_info.total_bytes_downloaded;
-            total_bytes_uploaded += connection_info.total_bytes_uploaded;
+        for state in self.utilization_data.iter().rev() {
+            let connections_to_procs = &state.connections_to_procs;
+            let network_utilization = &state.network_utilization;
+            for (connection, connection_info) in &network_utilization.connections {
+                let connection_data = connections.entry(connection.clone()).or_default();
+                let data_for_remote_address = remote_addresses
+                    .entry(connection.remote_socket.ip)
+                    .or_default();
+                connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded;
+                connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded;
+                connection_data.interface_name = connection_info.interface_name.clone();
+                data_for_remote_address.total_bytes_downloaded +=
+                    connection_info.total_bytes_downloaded;
+                data_for_remote_address.total_bytes_uploaded += connection_info.total_bytes_uploaded;
+                data_for_remote_address.connection_count += 1;
+                total_bytes_downloaded += connection_info.total_bytes_downloaded;
+                total_bytes_uploaded += connection_info.total_bytes_uploaded;
 
-            if let Some(process_name) =
-                UIState::get_proc_name(&connections_to_procs, &connection.local_socket)
-            {
-                let data_for_process = processes.entry(process_name.clone()).or_default();
-                data_for_process.total_bytes_downloaded += connection_info.total_bytes_downloaded;
-                data_for_process.total_bytes_uploaded += connection_info.total_bytes_uploaded;
-                data_for_process.connection_count += 1;
-                connection_data.process_name = process_name.clone();
-                // Record bandwidth data of last iteration
-                if let Some(prev_connection_info) = old_state.connections.get(&connection) {
-                    // Using previous round's weighted average. Exponential decay
-                    let prev_bytes_downloaded = prev_connection_info.get_avg_bytes_downloaded();
-                    let prev_bytes_uploaded = prev_connection_info.get_avg_bytes_uploaded();
-    
-                    connection_data.prev_total_bytes_downloaded += prev_bytes_downloaded;
-                    connection_data.prev_total_bytes_uploaded += prev_bytes_uploaded;
-    
-                    data_for_process.prev_total_bytes_downloaded += prev_bytes_downloaded;
-                    data_for_process.prev_total_bytes_uploaded += prev_bytes_uploaded;
-    
-                    data_for_remote_address.prev_total_bytes_downloaded += prev_bytes_downloaded;
-                    data_for_remote_address.prev_total_bytes_uploaded += prev_bytes_uploaded;
+                if let Some(process_name) =
+                    UIState::get_proc_name(&connections_to_procs, &connection.local_socket)
+                {
+                    let data_for_process = processes.entry(process_name.clone()).or_default();
+                    data_for_process.total_bytes_downloaded += connection_info.total_bytes_downloaded;
+                    data_for_process.total_bytes_uploaded += connection_info.total_bytes_uploaded;
+                    data_for_process.connection_count += 1;
+                    connection_data.process_name = process_name.clone();
+                } else {
+                    connection_data.process_name = String::from("<UNKNOWN>");
                 }
-            } else {
-                connection_data.process_name = String::from("<UNKNOWN>");
             }
         }
-        UIState {
-            processes,
-            remote_addresses,
-            connections,
-            total_bytes_downloaded,
-            total_bytes_uploaded,
+        for (_, network_data) in processes.iter_mut() {
+            network_data.divide_by(self.utilization_data.len())
+        }
+        for (_, network_data) in remote_addresses.iter_mut() {
+            network_data.divide_by(self.utilization_data.len())
+        }
+        for (_, connection_data) in connections.iter_mut() {
+            connection_data.divide_by(self.utilization_data.len())
         }
+        self.processes = processes;
+        self.remote_addresses = remote_addresses;
+        self.connections = connections;
+        self.total_bytes_downloaded = total_bytes_downloaded;
+        self.total_bytes_uploaded = total_bytes_uploaded;
     }
 }