Split up main (#39)

Moved code from inside the main function into function calls

Author: Azrenbeth

.gitignore

@@ -2,3 +2,4 @@
 **/*.rs.bk
 *.data
 *.old
+out.sql

src/main.rs

@@ -92,9 +92,19 @@ impl FromStr for LevelSizes {
     }
 }
 
-fn main() {
-    #[allow(deprecated)]
-    let matches = App::new(crate_name!())
+struct Config {
+    db_url: String,
+    output_file: Option<File>,
+    room_id: String,
+    max_state_group: Option<i64>,
+    min_saved_rows: Option<i32>,
+    transactions: bool,
+    level_sizes: LevelSizes,
+}
+
+impl Config {
+    fn parse_arguments() -> Config {
+        let matches = App::new(crate_name!())
         .version(crate_version!())
         .author(crate_authors!("\n"))
         .about(crate_description!())
@@ -156,33 +166,50 @@ fn main() {
                 .takes_value(true),
         ).get_matches();
 
-    let db_url = matches
-        .value_of("postgres-url")
-        .expect("db url should be required");
+        let db_url = matches
+            .value_of("postgres-url")
+            .expect("db url should be required");
+
+        let output_file = matches
+            .value_of("output_file")
+            .map(|path| File::create(path).unwrap());
 
-    let mut output_file = matches
-        .value_of("output_file")
-        .map(|path| File::create(path).unwrap());
+        let room_id = matches
+            .value_of("room_id")
+            .expect("room_id should be required since no file");
 
-    let room_id = matches
-        .value_of("room_id")
-        .expect("room_id should be required since no file");
+        let max_state_group = matches
+            .value_of("max_state_group")
+            .map(|s| s.parse().expect("max_state_group must be an integer"));
 
-    let max_state_group = matches
-        .value_of("max_state_group")
-        .map(|s| s.parse().expect("max_state_group must be an integer"));
+        let min_saved_rows = matches
+            .value_of("min_saved_rows")
+            .map(|v| v.parse().expect("COUNT must be an integer"));
 
-    let min_saved_rows = matches
-        .value_of("min_saved_rows")
-        .map(|v| v.parse().expect("COUNT must be an integer"));
+        let transactions = matches.is_present("transactions");
 
-    let transactions = matches.is_present("transactions");
+        let level_sizes = value_t_or_exit!(matches, "level_sizes", LevelSizes);
 
-    let level_sizes = value_t_or_exit!(matches, "level_sizes", LevelSizes);
+        Config {
+            db_url: String::from(db_url),
+            output_file,
+            room_id: String::from(room_id),
+            max_state_group,
+            min_saved_rows,
+            transactions,
+            level_sizes,
+        }
+    }
+}
+
+fn main() {
+    let mut config = Config::parse_arguments();
 
     // First we need to get the current state groups
-    println!("Fetching state from DB for room '{}'...", room_id);
-    let state_group_map = database::get_data_from_db(db_url, room_id, max_state_group);
+    println!("Fetching state from DB for room '{}'...", config.room_id);
+
+    let state_group_map =
+        database::get_data_from_db(&config.db_url, &config.room_id, config.max_state_group);
 
     println!("Number of state groups: {}", state_group_map.len());
 
@@ -196,7 +223,7 @@ fn main() {
 
     println!("Compressing state...");
 
-    let compressor = Compressor::compress(&state_group_map, &level_sizes.0);
+    let compressor = Compressor::compress(&state_group_map, &config.level_sizes.0);
 
     let new_state_group_map = compressor.new_state_group_map;
 
@@ -228,7 +255,7 @@ fn main() {
         compressor.stats.state_groups_changed
     );
 
-    if let Some(min) = min_saved_rows {
+    if let Some(min) = config.min_saved_rows {
         let saving = (original_summed_size - compressed_summed_size) as i32;
         if saving < min {
             println!(
@@ -239,25 +266,39 @@ fn main() {
         }
     }
 
+    check_that_maps_match(&state_group_map, &new_state_group_map);
+
     // If we are given an output file, we output the changes as SQL. If the
     // `transactions` argument is set we wrap each change to a state group in a
     // transaction.
 
-    if let Some(output) = &mut output_file {
-        println!("Writing changes...");
+    output_sql(&mut config, &state_group_map, &new_state_group_map);
+}
 
-        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);
+fn output_sql(
+    config: &mut Config,
+    old_map: &BTreeMap<i64, StateGroupEntry>,
+    new_map: &BTreeMap<i64, StateGroupEntry>,
+) {
+    if config.output_file.is_none() {
+        return;
+    }
+
+    println!("Writing changes...");
+
+    let pb = ProgressBar::new(old_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 (sg, old_entry) in &state_group_map {
-            let new_entry = &new_state_group_map[sg];
+    if let Some(output) = &mut config.output_file {
+        for (sg, old_entry) in old_map {
+            let new_entry = &new_map[sg];
 
             if old_entry != new_entry {
-                if transactions {
+                if config.transactions {
                     writeln!(output, "BEGIN;").unwrap();
                 }
 
@@ -292,7 +333,7 @@ fn main() {
                             output,
                             "({}, {}, {}, {}, {})",
                             sg,
-                            PGEscape(room_id),
+                            PGEscape(&config.room_id),
                             PGEscape(t),
                             PGEscape(s),
                             PGEscape(e)
@@ -302,21 +343,26 @@ fn main() {
                     writeln!(output, ";").unwrap();
                 }
 
-                if transactions {
+                if config.transactions {
                     writeln!(output, "COMMIT;").unwrap();
                 }
                 writeln!(output).unwrap();
             }
 
             pb.inc(1);
         }
-
-        pb.finish();
     }
 
+    pb.finish();
+}
+
+fn check_that_maps_match(
+    old_map: &BTreeMap<i64, StateGroupEntry>,
+    new_map: &BTreeMap<i64, StateGroupEntry>,
+) {
     println!("Checking that state maps match...");
 
-    let pb = ProgressBar::new(state_group_map.len() as u64);
+    let pb = ProgressBar::new(old_map.len() as u64);
     pb.set_style(
         ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"),
     );
@@ -325,11 +371,11 @@ fn main() {
 
     // Now let's iterate through and assert that the state for each group
     // matches between the two versions.
-    state_group_map
+    old_map
         .par_iter() // This uses rayon to run the checks in parallel
         .try_for_each(|(sg, _)| {
-            let expected = collapse_state_maps(&state_group_map, *sg);
-            let actual = collapse_state_maps(&new_state_group_map, *sg);
+            let expected = collapse_state_maps(&old_map, *sg);
+            let actual = collapse_state_maps(&new_map, *sg);
 
             pb.inc(1);