Add more postgres bars

src/compressor.rs

@@ -14,7 +14,7 @@
 //!  ^--- L1 <--- L1 <--- L1  ^--- L1 <--- L1 <--- L1
 //! ```
 
-
+use indicatif::{ProgressBar, ProgressStyle};
 use rust_matrix_lib::state_map::StateMap;
 use string_cache::DefaultAtom as Atom;
 
@@ -22,7 +22,6 @@ use std::collections::BTreeMap;
 
 use {collapse_state_maps, StateGroupEntry};
 
-
 /// Holds information about a particular level.
 struct Level {
     /// The maximum size this level is allowed to be
@@ -77,7 +76,6 @@ impl Level {
     }
 }
 
-
 /// Keeps track of some statistics of a compression run.
 #[derive(Default)]
 pub struct Stats {
@@ -90,7 +88,6 @@ pub struct Stats {
     pub state_groups_changed: usize,
 }
 
-
 /// Attempts to compress a set of state deltas using the given level sizes.
 pub struct Compressor<'a> {
     original_state_map: &'a BTreeMap<i64, StateGroupEntry>,
@@ -123,6 +120,13 @@ impl<'a> Compressor<'a> {
             panic!("Can only call `create_new_tree` once");
         }
 
+        let pb = ProgressBar::new(self.original_state_map.len() as u64);
+        pb.set_style(
+            ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"),
+        );
+        pb.set_message("state groups");
+        pb.enable_steady_tick(100);
+
         for (&state_group, entry) in self.original_state_map {
             let mut prev_state_group = None;
             for level in &mut self.levels {
@@ -149,7 +153,11 @@ impl<'a> Compressor<'a> {
                     state_map: delta,
                 },
             );
+
+            pb.inc(1);
         }
+
+        pb.finish();
     }
 
     /// Attempts to calculate the delta between two state groups.

src/database.rs

@@ -6,7 +6,6 @@ use std::collections::BTreeMap;
 
 use StateGroupEntry;
 
-
 /// Fetch the entries in state_groups_state (and their prev groups) for the
 /// given `room_id` by connecting to the postgres database at `db_url`.
 pub fn get_data_from_db(db_url: &str, room_id: &str) -> BTreeMap<i64, StateGroupEntry> {
@@ -69,7 +68,9 @@ fn get_initial_data_from_db(conn: &Connection, room_id: &str) -> BTreeMap<i64, S
     let mut state_group_map: BTreeMap<i64, StateGroupEntry> = BTreeMap::new();
 
     let pb = ProgressBar::new_spinner();
-    pb.set_style(ProgressStyle::default_spinner().template("{spinner} [{elapsed}] {pos}"));
+    pb.set_style(
+        ProgressStyle::default_spinner().template("{spinner} [{elapsed}] {pos} rows retrieved"),
+    );
     pb.enable_steady_tick(100);
 
     let mut num_rows = 0;
@@ -82,9 +83,11 @@ fn get_initial_data_from_db(conn: &Connection, room_id: &str) -> BTreeMap<i64, S
         let etype: Option<String> = row.get(2);
 
         if let Some(etype) = etype {
-            entry
+            entry.state_map.insert(
-                .state_map
+                &etype,
-                .insert(&etype, &row.get::<_, String>(3), row.get::<_, String>(4).into());
+                &row.get::<_, String>(3),
+                row.get::<_, String>(4).into(),
+            );
         }
 
         pb.inc(1);

src/main.rs

@@ -17,6 +17,7 @@ mod database;
 use compressor::Compressor;
 
 use clap::{App, Arg};
+use indicatif::{ProgressBar, ProgressStyle};
 use rayon::prelude::*;
 use rust_matrix_lib::state_map::StateMap;
 use string_cache::DefaultAtom as Atom;
@@ -152,6 +153,8 @@ fn main() {
 
     // Now we actually call the compression algorithm.
 
+    println!("Compressing state...");
+
     let compressor = Compressor::compress(&state_group_map, &level_sizes.0);
 
     let new_state_group_map = compressor.new_state_group_map;
@@ -235,6 +238,15 @@ fn main() {
         }
     }
 
+    println!("Checking that state maps match...");
+
+    let pb = ProgressBar::new(state_group_map.len() as u64);
+    pb.set_style(
+        ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"),
+    );
+    pb.set_message("state groups");
+    pb.enable_steady_tick(100);
+
     // Now let's iterate through and assert that the state for each group
     // matches between the two versions.
     state_group_map
@@ -243,6 +255,8 @@ fn main() {
             let expected = collapse_state_maps(&state_group_map, *sg);
             let actual = collapse_state_maps(&new_state_group_map, *sg);
 
+            pb.inc(1);
+
             if expected != actual {
                 println!("State Group: {}", sg);
                 println!("Expected: {:#?}", expected);
@@ -253,5 +267,7 @@ fn main() {
             }
         }).expect("expected state to match");
 
+    pb.finish();
+
     println!("New state map matches old one");
 }