OpenTally/src/parser/csp.rs

/*  OpenTally: Open-source election vote counting
 *  Copyright © 2021–2022 Lee Yingtong Li (RunasSudo)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use crate::election::{Ballot, Candidate, Election};
use crate::numbers::Number;

use anyhow::{Context, Result};
use csv::ReaderBuilder;
use itertools::Itertools;

use std::collections::HashMap;
use std::io::Read;

/// Parse the given CSP file
pub fn parse_reader<R: Read, N: Number>(reader: R, require_1: bool, require_sequential: bool, require_strict_order: bool) -> Result<Election<N>> {
	// Read CSV file
	let mut reader = ReaderBuilder::new()
		.has_headers(true)
		.from_reader(reader);
	
	// Read candidates
	let mut candidates = Vec::new();
	let mut col_map = HashMap::new(); // Map csp column -> candidates index
	
	for (index, cand_name) in reader.headers()?.into_iter().enumerate() {
		if cand_name == "$mult" {
			continue;
		}
		
		col_map.insert(index, candidates.len());
		candidates.push(Candidate {
			index,
			name: cand_name.to_string(),
			is_dummy: false,
		});
	}
	
	// Read ballots
	let mut ballots = Vec::new();
	
	for (csv_row, record) in reader.into_records().enumerate() {
		let record = record?;
		
		let mut value = N::one();
		
		// Record preferences
		let mut preferences = Vec::new(); // Vec of (ranking, candidate index)
		for (csv_col, preference) in record.into_iter().enumerate() {
			match col_map.get(&csv_col) {
				Some(cand_index) => {
					// Preference
					if preference.is_empty() || preference == "-" {
						continue;
					}
					
					let preference: usize = preference.parse().context(format!("Invalid number \"{}\" at row {}, column {}", preference, csv_row + 2, csv_col + 1))?;
					if preference == 0 {
						continue;
					}
					
					preferences.push((preference, cand_index));
				}
				None => {
					// $mult column
					let mult: usize = preference.parse().context(format!("Invalid number \"{}\" at row {}, column {}", preference, csv_row + 2, csv_col + 1))?;
					if mult == 1 {
						continue;
					}
					value = N::from(mult);
				}
			}
		}
		
		// Sort by ranking
		let mut unique_rankings: Vec<usize> = preferences.iter().map(|(r, _)| *r).unique().collect();
		unique_rankings.sort_unstable();
		
		if require_1 {
			if !unique_rankings.first().map(|r| *r == 1).unwrap_or(false) {
				// No #1 preference
				ballots.push(Ballot {
					orig_value: value,
					preferences: vec![],
					has_equal_rankings: false,
				});
				continue;
			}
		}
		
		let mut sorted_preferences = Vec::with_capacity(preferences.len());
		let mut last_ranking = None;
		let mut has_equal_rankings = false;
		
		for ranking in unique_rankings {
			// Filter for preferences at this ranking
			let prefs_this_ranking: Vec<usize> = preferences.iter()
				.filter_map(|(r, i)| if *r == ranking { Some(**i) } else { None })
				.collect();
			
			if prefs_this_ranking.len() != 1 {
				if require_strict_order {
					// Duplicate rankings
					break;
				}
				has_equal_rankings = true;
			}
			
			if require_sequential {
				if let Some(r) = last_ranking {
					if ranking != r + 1 {
						// Not sequential
						break;
					}
				}
			}
			
			sorted_preferences.push(prefs_this_ranking);
			last_ranking = Some(ranking);
		}
		
		ballots.push(Ballot {
			orig_value: value,
			preferences: sorted_preferences,
			has_equal_rankings,
		});
	}
	
	return Ok(Election {
		name: String::new(),
		seats: 0,
		candidates,
		withdrawn_candidates: Vec::new(),
		ballots,
		total_votes: None,
		constraints: None,
	});
}

#[test]
fn csp_formal() {
	let csp_data = "A,B,C\n1,2,3";
	let election = parse_reader::<_, crate::numbers::Rational>(csp_data.as_bytes(), false, false, false).unwrap();
	assert_eq!(election.ballots.first().unwrap().preferences, vec![vec![0], vec![1], vec![2]]);
	
	let csp_data = "A,B,C\n2,3,4";
	let election = parse_reader::<_, crate::numbers::Rational>(csp_data.as_bytes(), false, false, false).unwrap();
	assert_eq!(election.ballots.first().unwrap().preferences, vec![vec![0], vec![1], vec![2]]);
}

#[test]
fn csp_no1() {
	let csp_data = "A,B,C\n2,3,4";
	let election = parse_reader::<_, crate::numbers::Rational>(csp_data.as_bytes(), true, false, false).unwrap();
	assert_eq!(election.ballots.first().unwrap().preferences.len(), 0);
}

#[test]
fn csp_skipped_preference() {
	let csp_data = "A,B,C\n1,3,4";
	let election = parse_reader::<_, crate::numbers::Rational>(csp_data.as_bytes(), false, true, false).unwrap();
	assert_eq!(election.ballots.first().unwrap().preferences, vec![vec![0]]);
}

#[test]
fn csp_duplicate_preference() {
	let csp_data = "A,B,C\n1,2,2";
	let election = parse_reader::<_, crate::numbers::Rational>(csp_data.as_bytes(), false, false, true).unwrap();
	assert_eq!(election.ballots.first().unwrap().preferences, vec![vec![0]]);
}