diff --git a/src/compressor.rs b/src/compressor.rs index fd068c8..6c245aa 100644 --- a/src/compressor.rs +++ b/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, @@ -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. diff --git a/src/database.rs b/src/database.rs index 8ada912..e74012a 100644 --- a/src/database.rs +++ b/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 { @@ -69,7 +68,9 @@ fn get_initial_data_from_db(conn: &Connection, room_id: &str) -> BTreeMap = 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 = row.get(2); if let Some(etype) = etype { - entry - .state_map - .insert(&etype, &row.get::<_, String>(3), row.get::<_, String>(4).into()); + entry.state_map.insert( + &etype, + &row.get::<_, String>(3), + row.get::<_, String>(4).into(), + ); } pb.inc(1); diff --git a/src/main.rs b/src/main.rs index 1edb8f6..f452a1a 100644 --- a/src/main.rs +++ b/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"); }