From e3ca9fac47c9136d4530e07f4fd4f30d15fe01e2 Mon Sep 17 00:00:00 2001
From: RunasSudo <runassudo@yingtongli.me>
Date: Sat, 4 Sep 2021 22:46:29 +1000
Subject: [PATCH] Refactor stv::preprocess_election

---
 html/worker.js  | 25 +++++++++----------------
 src/cli/stv.rs  |  8 +-------
 src/election.rs |  3 +++
 src/stv/mod.rs  | 12 ++++++++++++
 src/stv/wasm.rs | 21 +++++++--------------
 5 files changed, 32 insertions(+), 37 deletions(-)

diff --git a/html/worker.js b/html/worker.js
index 1279bf2..f469684 100644
--- a/html/worker.js
+++ b/html/worker.js
@@ -41,28 +41,21 @@ onmessage = function(evt) {
 			
 			reportStyle = evt.data.reportStyle;
 			
-			// Init election
-			election = wasm['election_from_blt_' + numbers](evt.data.bltData);
-			
-			// Normalise ballots if requested
-			if (evt.data.normaliseBallots) {
-				wasm['election_normalise_ballots_' + numbers](election);
-			}
-			
-			// Process equal rankings
-			wasm['election_realise_equal_rankings_' + numbers](election);
-			
-			// Init constraints if applicable
-			if (evt.data.conData) {
-				wasm['election_load_constraints_' + numbers](election, evt.data.conData);
-			}
-			
 			// Init STV options
 			opts = wasm.STVOptions.new.apply(null, evt.data.optsStr);
 			
 			// Validate options
 			opts.validate();
 			
+			// Init election
+			election = wasm['election_from_blt_' + numbers](evt.data.bltData);
+			wasm['preprocess_election_' + numbers](election, opts);
+			
+			// Init constraints if applicable
+			if (evt.data.conData) {
+				wasm['election_load_constraints_' + numbers](election, evt.data.conData);
+			}
+			
 			// Describe count
 			postMessage({'type': 'describeCount', 'content': wasm['describe_count_' + numbers](evt.data.bltPath, election, opts)});
 			
diff --git a/src/cli/stv.rs b/src/cli/stv.rs
index dd0261b..007baa7 100644
--- a/src/cli/stv.rs
+++ b/src/cli/stv.rs
@@ -304,13 +304,7 @@ where
 	}
 	println!();
 	
-	// Normalise ballots if requested
-	if cmd_opts.normalise_ballots {
-		election.normalise_ballots();
-	}
-	
-	// Process equal rankings
-	election.realise_equal_rankings();
+	stv::preprocess_election(&mut election, &opts);
 	
 	// Initialise count state
 	let mut state = CountState::new(&election);
diff --git a/src/election.rs b/src/election.rs
index 0d8b46d..52d511c 100644
--- a/src/election.rs
+++ b/src/election.rs
@@ -156,14 +156,17 @@ impl<'a, N: Number> CountState<'a, N> {
 			logger: Logger { entries: Vec::new() },
 		};
 		
+		// Init candidate count cards
 		for candidate in election.candidates.iter() {
 			state.candidates.insert(candidate, CountCard::new());
 		}
 		
+		// Set withdrawn candidates state
 		for withdrawn_idx in election.withdrawn_candidates.iter() {
 			state.candidates.get_mut(&election.candidates[*withdrawn_idx]).unwrap().state = CandidateState::Withdrawn;
 		}
 		
+		// Init constraints
 		if let Some(constraints) = &election.constraints {
 			let mut num_groups: Vec<usize> = constraints.0.iter().map(|c| c.groups.len()).collect();
 			let mut cm = ConstraintMatrix::new(&mut num_groups[..]);
diff --git a/src/stv/mod.rs b/src/stv/mod.rs
index 5652216..2766820 100644
--- a/src/stv/mod.rs
+++ b/src/stv/mod.rs
@@ -29,6 +29,7 @@ pub mod sample;
 pub mod wasm;
 
 use crate::constraints;
+use crate::election::Election;
 use crate::numbers::Number;
 use crate::election::{Candidate, CandidateState, CountCard, CountState, Vote};
 use crate::sharandom::SHARandom;
@@ -577,6 +578,17 @@ impl fmt::Display for STVError {
 	}
 }
 
+/// Preprocess the given election
+pub fn preprocess_election<N: Number>(election: &mut Election<N>, opts: &STVOptions) {
+	// Normalise ballots if requested
+	if opts.normalise_ballots {
+		election.normalise_ballots();
+	}
+	
+	// Process equal rankings
+	election.realise_equal_rankings();
+}
+
 /// Distribute first preferences, and initialise other states such as the random number generator and tie-breaking rules
 pub fn count_init<'a, N: Number>(state: &mut CountState<'a, N>, opts: &'a STVOptions) -> Result<(), STVError>
 where
diff --git a/src/stv/wasm.rs b/src/stv/wasm.rs
index 3cc0fea..1e70a18 100644
--- a/src/stv/wasm.rs
+++ b/src/stv/wasm.rs
@@ -80,20 +80,6 @@ macro_rules! impl_type {
 			return [<Election$type>](election);
 		}
 		
-		/// Wrapper for [Election::normalise_ballots]
-		#[wasm_bindgen]
-		#[allow(non_snake_case)]
-		pub fn [<election_normalise_ballots_$type>](election: &mut [<Election$type>]) {
-			election.0.normalise_ballots();
-		}
-		
-		/// Wrapper for [Election::realise_equal_rankings]
-		#[wasm_bindgen]
-		#[allow(non_snake_case)]
-		pub fn [<election_realise_equal_rankings_$type>](election: &mut [<Election$type>]) {
-			election.0.realise_equal_rankings();
-		}
-		
 		/// Call [Constraints::from_con] and set [Election::constraints]
 		#[wasm_bindgen]
 		#[allow(non_snake_case)]
@@ -101,6 +87,13 @@ macro_rules! impl_type {
 			election.0.constraints = Some(Constraints::from_con(text.lines().map(|s| s.to_string()).into_iter()));
 		}
 		
+		/// Wrapper for [stv::preprocess_election]
+		#[wasm_bindgen]
+		#[allow(non_snake_case)]
+		pub fn [<preprocess_election_$type>](election: &mut [<Election$type>], opts: &STVOptions) {
+			stv::preprocess_election(&mut election.0, &opts.0);
+		}
+		
 		/// Wrapper for [stv::count_init]
 		#[wasm_bindgen]
 		#[allow(non_snake_case)]