142 lines
3.6 KiB
Rust
142 lines
3.6 KiB
Rust
use std::{sync::{Arc, Mutex}, thread};
|
|
use itertools::Itertools;
|
|
|
|
use crate::{ballot::Ballot, csv, header::Header, stats::Stats, util::{self, ScoreItem}};
|
|
|
|
#[derive(Debug)]
|
|
pub enum Event {
|
|
Lose(ScoreItem),
|
|
Win(Vec<ScoreItem>),
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct Counter {
|
|
pub header: Arc<Header>,
|
|
pub ballots: Vec<Ballot>,
|
|
pub stats: Stats,
|
|
winners_left: usize,
|
|
enabled: Vec<bool>,
|
|
quota: f64,
|
|
}
|
|
|
|
const CHUNK_SIZE: usize = 1024;
|
|
|
|
impl Counter {
|
|
pub fn new(mut csv: csv::reader::CsvReader, winners: usize) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
|
|
let header = Arc::new(Header::parse(csv.next().ok_or("csv header missing")?, winners)?);
|
|
let ballots = Mutex::new(Vec::new());
|
|
let stats = Mutex::new(Stats::new());
|
|
let csv = Mutex::new(csv);
|
|
|
|
if winners > header.candidates.len() {
|
|
return Err("winners can't be smaller than the candidates list".into());
|
|
}
|
|
|
|
util::thread::spawn_all_with_result(|| -> Result<(), Box<dyn std::error::Error + Sync + Send>> {
|
|
let mut l_stats = Stats::new();
|
|
let mut l_ballots = Vec::new();
|
|
loop {
|
|
let rows = csv.lock().unwrap().by_ref().take(CHUNK_SIZE).collect_vec();
|
|
|
|
if rows.len() == 0 {
|
|
break;
|
|
}
|
|
|
|
for row in rows {
|
|
l_ballots.push(Ballot::parse(row, &header, &mut l_stats).unwrap());
|
|
}
|
|
|
|
ballots.lock().unwrap().append(&mut l_ballots);
|
|
}
|
|
{stats.lock().unwrap().add(&l_stats)};
|
|
Ok(())
|
|
})?;
|
|
|
|
let ballots = ballots.into_inner().unwrap();
|
|
let stats = stats.into_inner().unwrap();
|
|
|
|
let enabled = vec![true; header.candidates.len()];
|
|
let quota = (ballots.len() as f64) / (header.winners as f64 + 1.0) + 1.0;
|
|
let winners_left = header.winners;
|
|
|
|
Ok(Counter {
|
|
header,
|
|
ballots,
|
|
stats,
|
|
winners_left,
|
|
enabled,
|
|
quota,
|
|
})
|
|
}
|
|
fn count_primaries(&self) -> Vec<ScoreItem> {
|
|
let mut scores = Mutex::new(vec![0.0; self.enabled.len()]);
|
|
let ballot_chunks = Mutex::new(self.ballots.chunks(CHUNK_SIZE));
|
|
|
|
util::thread::spawn_all(|| {
|
|
let mut l_scores = vec![0.0; self.enabled.len()];
|
|
for ballot in util::mutex::chunk_iter(&ballot_chunks) {
|
|
ballot.count_primary(&mut l_scores, &self.enabled);
|
|
}
|
|
for (dst, src) in scores.lock().unwrap().iter_mut().zip(l_scores.into_iter()) {
|
|
*dst += src;
|
|
}
|
|
});
|
|
|
|
scores.get_mut().unwrap().iter().copied().enumerate().filter_map(|(index, value)| match self.enabled[index] {
|
|
true => Some(ScoreItem::new(index, value)),
|
|
false => None,
|
|
}).collect_vec()
|
|
}
|
|
fn apply_weights(&mut self, winners: &[ScoreItem]) {
|
|
let ballot_chunks = Mutex::new(self.ballots.chunks_mut(CHUNK_SIZE));
|
|
|
|
util::thread::spawn_all(|| {
|
|
for ballot in util::mutex::chunk_iter_mut(&ballot_chunks) {
|
|
ballot.apply_winners(&self.enabled, winners, self.quota);
|
|
}
|
|
});
|
|
}
|
|
#[inline]
|
|
pub fn get_quota(&self) -> f64 {
|
|
self.quota
|
|
}
|
|
}
|
|
|
|
impl Iterator for Counter {
|
|
type Item = Event;
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
let mut winners = Vec::<ScoreItem>::with_capacity(self.enabled.len());
|
|
let primaries = self.count_primaries();
|
|
let mut lowest = *primaries.first()?;
|
|
|
|
for &score in primaries.iter() {
|
|
lowest = lowest.min(score);
|
|
if score.value >= self.quota {
|
|
winners.push(score);
|
|
}
|
|
}
|
|
|
|
if primaries.len() == 1 && self.winners_left == 1 {
|
|
self.enabled[lowest.index] = false;
|
|
return Some(Event::Win(vec![lowest]));
|
|
}
|
|
|
|
if winners.len() == 0 {
|
|
self.enabled[lowest.index] = false;
|
|
return Some(Event::Lose(lowest));
|
|
}
|
|
|
|
winners.sort_by(|a,b| b.cmp(a));
|
|
self.apply_weights(&winners);
|
|
self.winners_left -= winners.len();
|
|
|
|
for winner in winners.iter() {
|
|
self.enabled[winner.index] = false;
|
|
}
|
|
|
|
Some(Event::Win(winners))
|
|
}
|
|
}
|
|
|