refactor

src/ballot.rs

@@ -1,8 +1,10 @@
+use std::rc::Rc;
+
 use itertools::Itertools;
 use crate::header::Header;
 
 #[derive(Copy,Clone,Debug,PartialEq)]
-pub enum BallotType {
+enum BallotType {
 	Party,
 	Candidate,
 }
@@ -10,21 +12,18 @@ pub enum BallotType {
 
 #[derive(Debug)]
 pub struct Ballot {
-	pub state: String,
-	pub division: String,
-	pub collection_point: String,
-	pub collection_point_id: usize,
-	pub batch_id: usize,
-	pub paper_id: usize,
-	pub votes: Vec<usize>,
-	pub ty: BallotType,
+	header: Rc<Header>,
+	collection_point_id: usize,
+	batch_id: usize,
+	paper_id: usize,
+	votes: Vec<usize>,
 }
 
 impl Ballot {
-	pub fn parse(row: quick_csv::Row, header: &Header) -> Result<Ballot, Box<dyn std::error::Error>> {
+	pub fn parse(row: quick_csv::Row, header: Rc<Header>) -> Result<Ballot, Box<dyn std::error::Error>> {
 		let mut cols = row.columns()?;
 
-		let (state, division, collection_point, collection_point_id, batch_id, paper_id) = cols.next_tuple().ok_or("Missing columns")?;
+		let (_, _, _, collection_point_id, batch_id, paper_id) = cols.next_tuple().ok_or("Missing columns")?;
 		let cols = cols.collect_vec();
 		let mut votes: Vec<Option<usize>>;
 		let ty: BallotType;
@@ -52,39 +51,34 @@ impl Ballot {
 			}
 		}
 
+		let votes = match ty {
+			BallotType::Candidate => votes.iter().filter_map(|&v| v).collect(),
+			BallotType::Party => votes.iter().filter_map(|&v| v).map(|id| header.parties[id].member_ids.iter().copied()).flatten().collect(),
+		};
+
 		Ok(Ballot {
-			state: state.to_owned(),
-			division: division.to_owned(),
-			collection_point: collection_point.to_owned(),
+			header,
 			collection_point_id: collection_point_id.parse()?,
 			batch_id: batch_id.parse()?,
 			paper_id: paper_id.parse()?,
-			votes: votes.iter().copied().filter_map(|v| v).collect(),
-			ty,
+			votes,
 		})
 	}
-
-	pub fn to_candidate_ballot(self, header: &Header) -> Self {
-		match self.ty {
-			BallotType::Candidate => self,
-			BallotType::Party => {
-				let mut votes = Vec::<usize>::new();
-				
-				for party_id in self.votes {
-					votes.extend(header.parties[party_id].member_ids.iter().copied());
+	pub fn get_primary_index(&self, enabled: &[bool]) -> Option<usize> {
+		for &id in self.votes.iter() {
+			if enabled[id] {
+				return Some(id);
 			}
-
-				Ballot {
-					state: self.state,
-					division: self.division,
-					collection_point: self.collection_point,
-					collection_point_id: self.collection_point_id,
-					batch_id: self.batch_id,
-					paper_id: self.paper_id,
-					ty: BallotType::Candidate,
-					votes,
 		}
-			},
+		return None;
+	}
+	pub fn count_all_with_decaying_weights(&self, scores: &mut [f64], enabled: &[bool]) {
+		let mut m = 0.5;
+		for &id in self.votes.iter() {
+			if enabled[id] {
+				scores[id] += m;
+				m *= 0.5;
+			}
 		}
 	}
 }

src/candidate.rs

@@ -1,3 +1,5 @@
+use crate::header::Header;
+
 
 #[derive(Debug)]
 pub struct Candidate {
@@ -5,3 +7,28 @@ pub struct Candidate {
 	pub party_id: Option<usize>,
 }
 
+#[derive(Copy,Clone)]
+pub struct CandidateName<'a> {
+	candidate_name: &'a str,
+	party_name: Option<&'a str>,
+}
+
+impl Candidate {
+	pub fn get_name<'a>(&'a self, header: &'a Header) -> CandidateName<'a> {
+		CandidateName {
+			candidate_name: &self.name,
+			party_name: self.party_id.map_or(None, |id| Some(&header.parties[id].name)),
+		}
+	}
+}
+
+impl<'a> std::fmt::Display for CandidateName<'a> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}", self.candidate_name)?;
+		if let Some(party_name) = self.party_name {
+			write!(f, ": {party_name}")?;
+		}
+		Ok(())
+	}
+}
+

src/counter.rs

@@ -1,9 +1,11 @@
-use crate::{ballot::{Ballot, BallotType}, header::Header};
+use std::rc::Rc;
+
+use crate::{ballot::Ballot, header::Header};
 
 
 #[derive(Debug)]
 pub struct Counter {
-	pub header: Header,
+	pub header: Rc<Header>,
 	pub ballots: Vec<Ballot>,
 }
 
@@ -13,7 +15,7 @@ impl Counter {
 		let mut ballots = Vec::new();
 
 		for row in csv {
-			ballots.push(Ballot::parse(row?, &header)?.to_candidate_ballot(&header));
+			ballots.push(Ballot::parse(row?, header.clone())?);
 		}
 
 		Ok(Counter {
@@ -22,40 +24,27 @@ impl Counter {
 		})
 	}
 
-	pub fn count_primaries<T>(&self, ty: BallotType, running: &[bool]) -> Vec<T> where T: num::Num + Copy + std::ops::AddAssign {
-		let mut scores = vec![T::zero(); running.len()];
+	pub fn count_primaries(&self, enabled: &[bool]) -> Vec<f64> {
+		assert!(enabled.len() == self.header.candidates.len());
+
+		let mut scores = vec![0.0; enabled.len()];
 
 		for ballot in self.ballots.iter() {
-			if ballot.ty != ty {
-				continue;
-			}
-			for &id in ballot.votes.iter() {
-				if !running[id] {
-					continue;
-				}
-				scores[id] += T::one();
-				break;
+			if let Some(index) = ballot.get_primary_index(enabled) {
+				scores[index] += 1.0;
 			}
 		}
 
 		scores
 	}
 	
-	pub fn count_all_with_powers_of_2(&self, ty: BallotType, running: &[bool]) -> Vec<f64> {
-		let mut scores = vec![0.0; running.len()];
+	pub fn count_all_with_decaying_weights(&self, enabled: &[bool]) -> Vec<f64> {
+		assert!(enabled.len() == self.header.candidates.len());
+
+		let mut scores = vec![0.0; enabled.len()];
 
 		for ballot in self.ballots.iter() {
-			if ballot.ty != ty {
-				continue;
-			}
-			let mut m = 0.5;
-			for &id in ballot.votes.iter() {
-				if !running[id] {
-					continue;
-				}
-				scores[id] += m;
-				m *= 0.5;
-			}
+			ballot.count_all_with_decaying_weights(&mut scores, enabled);
 		}
 
 		scores

src/csv.rs

@@ -0,0 +1,56 @@
+use crate::util::{EscapeWriter, EscapeWriterOpts};
+use std::io::Write;
+
+
+pub struct CsvWriter<T> where T: std::io::Write {
+	out: T,
+}
+
+pub struct RowWriter<'a,T> where T: std::io::Write {
+	csv: &'a mut CsvWriter<T>,
+	first: bool,
+}
+
+impl<T> CsvWriter<T> where T: std::io::Write {
+	pub fn new(out: T) -> Self {
+		Self { out }
+	}
+	pub fn row(&mut self) -> RowWriter<'_,T> {
+		RowWriter { csv: self, first: true }
+	}
+}
+
+impl<'a,T> RowWriter<'a,T> where T: std::io::Write {
+	pub fn skip(&mut self, cells: usize) -> std::io::Result<()> {
+		for _ in 0..cells {
+			if self.first {
+				self.first = false;
+			}
+			else {
+				write!(self.csv.out, ",")?;
+			}
+		}
+		Ok(())
+	}
+	pub fn write<F>(&mut self, value: F) -> std::io::Result<()>
+		where F: std::fmt::Display
+	{
+		self.skip(1)?;
+		write!(self.csv.out, "\"")?;
+		let mut writer = EscapeWriter::new(&mut self.csv.out, EscapeWriterOpts {
+			escape: b'\\',
+			chars: b"\"",
+		});
+		write!(writer, "{value}")?;
+		write!(self.csv.out, "\"")?;
+
+		Ok(())
+	}
+}
+
+impl<'a,T> Drop for RowWriter<'a,T> where T: std::io::Write {
+	fn drop(&mut self) {
+		_ = writeln!(self.csv.out);
+	}
+}
+

src/header.rs

@@ -1,8 +1,8 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, rc::Rc};
 
 use itertools::Itertools;
 
-use crate::{candidate::Candidate, party::Party};
+use crate::{candidate::{Candidate, CandidateName}, party::Party};
 
 
 #[derive(Debug)]
@@ -12,7 +12,7 @@ pub struct Header {
 }
 
 impl Header {
-	pub fn parse(row: quick_csv::Row) -> Result<Header, Box<dyn std::error::Error>> {
+	pub fn parse(row: quick_csv::Row) -> Result<Rc<Header>, Box<dyn std::error::Error>> {
 		let mut header = Header {
 			parties: Vec::new(),
 			candidates: Vec::new(),
@@ -41,7 +41,13 @@ impl Header {
 			}
 		}
 
-		Ok(header)
+		Ok(header.into())
+	}
+	pub fn get_candidate_name(&self, index: usize) -> CandidateName {
+		self.candidates[index].get_name(self)
+	}
+	pub fn iter_candidate_name(&self) -> impl Iterator<Item = CandidateName> {
+		self.candidates.iter().map(|v| v.get_name(self))
 	}
 }
 

src/main.rs

@@ -1,9 +1,11 @@
 use std::env;
-
-use ballot::BallotType;
 use counter::Counter;
+use csv::CsvWriter;
 use itertools::Itertools;
+use util::{percent::PercentDiff, Percent};
 
+pub mod csv;
+pub mod util;
 pub mod candidate;
 pub mod ballot;
 pub mod header;
@@ -21,94 +23,90 @@ fn main() {
 	};
 	let csv = quick_csv::Csv::from_file(csv_path).unwrap().flexible(true);
 	let counter = Counter::new(csv).unwrap();
+	let header = counter.header.clone();
 
-	eprintln!("Header: {:#?}", counter.header);
-	eprintln!("Parties: {}", counter.header.parties.len());
-	eprintln!("Candidates: {}", counter.header.candidates.len());
+	eprintln!("Header: {:#?}", header);
+	eprintln!("Parties: {}", header.parties.len());
+	eprintln!("Candidates: {}", header.candidates.len());
 
 	let total = counter.ballots.len() as u32;
-	let mut total_atl = 0u32;
-	let mut total_btl = 0u32;
+	let mut running = vec![true; header.candidates.len()];
+	let mut csv_out = CsvWriter::new(std::io::stdout().lock());
 
-	for ballot in counter.ballots.iter() {
-		*match ballot.ty {
-			BallotType::Party => &mut total_atl,
-			BallotType::Candidate => &mut total_btl,
-		} += 1;
+	{
+		let mut row = csv_out.row();
+		row.skip(4).unwrap();
+		for name in header.iter_candidate_name() {
+			row.write(name).unwrap();
+			row.skip(1).unwrap();
 		}
-
-	eprintln!("Above the line: {total_atl} {}%", f64::from(total_atl) / f64::from(total) * 100.0);
-	eprintln!("Below the line: {total_btl} {}%", f64::from(total_btl) / f64::from(total) * 100.0);
-	eprintln!("Total: {total}");
-
-	let ty = BallotType::Candidate;
-	let names = match ty {
-		BallotType::Party => counter.header.parties.iter().map(|v| v.name.clone()).collect_vec(),
-		BallotType::Candidate => counter.header.candidates.iter().map(|v| match v.party_id {
-			Some(party_id) => format!("{}: {}", v.name, &counter.header.parties[party_id].name),
-			None => v.name.to_owned(),
-		}).collect_vec(),
-	};
-	let mut running = vec![true; names.len()];
-
-	print!(",,,");
-	for name in names.iter() {
-		print!(",{name:?},");
 	}
-	println!();
 	eprintln!("Running election rounds:");
 
-	let to_percent = |score: f64| score / f64::from(total) * 100.0;
 	let mut scores_last: Option<Vec<f64>> = None;
 	let mut id_worst_last: Option<usize> = None;
 
 	for index in 0..(running.len() - winner_count + 1) {
-		let scores = counter.count_all_with_powers_of_2(ty, &running);
-		let scores_ordered = scores.iter().copied().enumerate().filter(|&(i,_)| running[i]).sorted_by(|a, b| f64::total_cmp(&b.1, &a.1)).collect_vec();
+		let scores = counter.count_all_with_decaying_weights(&running);
+		let scores_ordered = scores.iter().copied().enumerate().filter(|&(id,_)| running[id]).sorted_by(|a, b| f64::total_cmp(&b.1, &a.1)).collect_vec();
+
 		let (id_worst, score_worst) = match scores_ordered.last() {
 			Some(&v) => v,
 			None => break,
 		};
 
+		csv_out.row();
+		let mut row = csv_out.row();
+
 		if let Some(scores_last) = scores_last {
 			let id_worst_last = id_worst_last.unwrap();
 			let score_worst_last = scores_last[id_worst_last];
-			print!(",{:?},{score_worst_last},{:.3}%", names[id_worst_last], to_percent(score_worst_last));
+
+			row.skip(1).unwrap();
+			row.write(header.get_candidate_name(id_worst_last)).unwrap();
+			row.write(score_worst_last).unwrap();
+			row.write(Percent(score_worst_last, total.into())).unwrap();
+
 			for ((&score, &score_last), &running) in scores.iter().zip(scores_last.iter()).zip(running.iter()) {
 				if running {
 					let score_diff = score - score_last;
-					print!(",+{score_diff},+{:.3}%", score_diff / score_worst_last * 100.0);
+					row.write(score_diff).unwrap();
+					row.write(PercentDiff(score_diff, score_worst_last)).unwrap();
 				} else {
-					print!(",,");
+					row.skip(2).unwrap();
 				}
 			}
-			println!();
+			drop(row);
+			row = csv_out.row();
 		}
-		print!("{},,,", index + 1);
+
+		row.write(index + 1).unwrap();
+		row.skip(3).unwrap();
+
 		for (&score, &running) in scores.iter().zip(running.iter()) {
 			if running {
-				print!(",{score},{:.3}%", to_percent(score));
+				row.write(score).unwrap();
+				row.write(Percent(score, total.into())).unwrap();
 			} else {
-				print!(",,");
+				row.skip(2).unwrap();
 			}
 		}
-		println!();
+
 		running[id_worst] = false;
 		scores_last = Some(scores);
 		id_worst_last = Some(id_worst);
 
 		let mut total_votes = 0.0;
 
-		println!();
 		eprintln!("Round {}:", index+1);
 		for (place, &(id, score)) in scores_ordered.iter().take(winner_count).enumerate() {
-			eprintln!("  {}. {}: {score}, {:.3}%", place + 1, names[id], to_percent(score));
+			eprintln!("  {}. {}: {score}, {}", place + 1, header.get_candidate_name(id), Percent(score, total.into()));
 			total_votes += score;
 		}
 		if scores_ordered.len() > winner_count {
-			eprintln!("  {}. {}: {score_worst}, {:.3}%", scores_ordered.len(), names[id_worst], to_percent(score_worst));
+			eprintln!("  {}. {}: {score_worst}, {}", scores_ordered.len(), header.get_candidate_name(id_worst), Percent(score_worst, total.into()));
 		}
-		eprintln!("  Total: {total_votes}, {:.3}%", to_percent(total_votes.into()));
+		eprintln!("  Total: {total_votes}, {}", Percent(total_votes.into(), total.into()));
 	}
 }
 

src/util.rs

@@ -0,0 +1,7 @@
+
+pub mod escape;
+pub mod percent;
+
+pub use escape::{EscapeWriter, EscapeWriterOpts};
+pub use percent::Percent;
+

src/util/escape.rs

@@ -0,0 +1,49 @@
+
+pub struct EscapeWriterOpts<'a> {
+	pub chars: &'a [u8],
+	pub escape: u8,
+}
+
+pub struct EscapeWriter<'a,T> {
+	opts: EscapeWriterOpts<'a>,
+	dst: T,
+}
+
+impl<'a,T> EscapeWriter<'a,T> {
+	pub fn new(dst: T, opts: EscapeWriterOpts<'a>) -> Self {
+		Self { dst, opts }
+	}
+}
+
+impl<'a> Default for EscapeWriterOpts<'a> {
+	fn default() -> Self {
+		Self {
+			escape: b'\\',
+			chars: b"\"",
+		}
+	}
+}
+
+impl<'a,T> std::io::Write for EscapeWriter<'a,T> where T: std::io::Write {
+	fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+		let mut count = 0;
+		for &b in buf.iter() {
+			if b == self.opts.escape || self.opts.chars.contains(&b) {
+				if self.dst.write(&[self.opts.escape])? == 0 {
+					return Ok(count);
+				}
+				count += 1;
+			}
+			if self.dst.write(&[b])? == 0 {
+				return Ok(count);
+			}
+			count += 1;
+		}
+		Ok(count)
+    }
+
+	fn flush(&mut self) -> std::io::Result<()> {
+		self.dst.flush()
+    }
+}
+

src/util/percent.rs

@@ -0,0 +1,21 @@
+use std::fmt::Display;
+
+
+pub struct Percent(pub f64, pub f64);
+pub struct PercentDiff(pub f64, pub f64);
+
+impl Display for Percent {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{:.3}%", self.0 / self.1 * 100.0)
+	}
+}
+impl Display for PercentDiff {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		let v = self.0 / self.1 * 100.0;
+		write!(f, "{}{:.3}%", match v >= 0.0 {
+			true => '+',
+			false => '-',
+		}, v.abs())
+	}
+}
+