OpenTally/src/candmap.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::Candidate;

use std::ops::Index;

/// Newtype for [HashMap] on [Candidate]s
#[derive(Clone)]
pub struct CandidateMap<'e, V> {
	keys: Vec<Option<&'e Candidate>>,
	values: Vec<Option<V>>
}

impl<'e, V> CandidateMap<'e, V> {
	/// See [HashMap::new]
	pub fn new() -> Self {
		Self {
			keys: Vec::new(),
			values: Vec::new()
		}
	}
	
	/// See [HashMap::with_capacity]
	pub fn with_capacity(capacity: usize) -> Self {
		let mut ret = Self {
			keys: Vec::with_capacity(capacity),
			values: Vec::with_capacity(capacity)
		};
		ret.maybe_resize(capacity);
		return ret;
	}
	
	fn maybe_resize(&mut self, len: usize) {
		if len < self.keys.len() {
			return;
		}
		
		self.keys.resize_with(len, || None);
		self.values.resize_with(len, || None);
	}
	
	/// See [HashMap::len]
	#[inline]
	pub fn len(&self) -> usize {
		return self.keys.iter().filter(|k| k.is_some()).count();
	}
	
	/// See [HashMap::insert]
	#[inline]
	pub fn insert(&mut self, candidate: &'e Candidate, value: V) {
		self.maybe_resize(candidate.index + 1);
		self.keys[candidate.index] = Some(candidate);
		self.values[candidate.index] = Some(value);
	}
	
	/// See [HashMap::get]
	#[inline]
	pub fn get(&self, candidate: &'e Candidate) -> Option<&V> {
		return self.values.get(candidate.index).unwrap_or(&None).as_ref();
	}
	
	/// See [HashMap::get_mut]
	#[inline]
	pub fn get_mut(&mut self, candidate: &'e Candidate) -> Option<&mut V> {
		match self.values.get_mut(candidate.index) {
			Some(v) => {
				return v.as_mut();
			}
			None => {
				return None;
			}
		}
	}
	
	/// See [HashMap::iter]
	#[inline]
	pub fn iter(&self) -> Iter<'_, 'e, V> {
		return Iter { map: &self, index: 0 };
	}
	
	/// See [HashMap::iter_mut]
	#[inline]
	pub fn iter_mut(&mut self) -> IterMut<'_, 'e, V> {
		return IterMut { map: self, index: 0 };
	}
	
	/// See [HashMap::values]
	#[inline]
	pub fn values(&self) -> Values<'_, 'e, V> {
		return Values { map: &self, index: 0 };
	}
}

impl<'e, V> Index<&Candidate> for CandidateMap<'e, V> {
	type Output = V;
	
	fn index(&self, candidate: &Candidate) -> &Self::Output {
		return self.values.get(candidate.index).unwrap_or(&None).as_ref().unwrap();
	}
}

pub struct Iter<'m, 'e, V> {
	map: &'m CandidateMap<'e, V>,
	index: usize
}

impl<'m, 'e, V> Iterator for Iter<'m, 'e, V> {
	type Item = (&'e Candidate, &'m V);
	
	fn next(&mut self) -> Option<Self::Item> {
		loop {
			match self.map.keys.get(self.index) {
				Some(k) => {
					// Key within range
					match k {
						Some(kk) => {
							// Key is set
							
							// SAFETY: Guaranteed to be set, as we update key and value at the same time
							let v = unsafe { self.map.values.get_unchecked(self.index).as_ref().unwrap_unchecked() };
							self.index += 1;
							return Some((kk, v));
						}
						None => {
							// Key is unset
							self.index += 1;
							continue;
						}
					}
				}
				None => {
					// Key outside range
					return None;
				}
			}
		}
	}
}

pub struct IterMut<'m, 'e, V> {
	map: &'m mut CandidateMap<'e, V>,
	index: usize
}

impl<'m, 'e, V> Iterator for IterMut<'m, 'e, V> {
	type Item = (&'e Candidate, &'m mut V);
	
	fn next(&mut self) -> Option<Self::Item> {
		loop {
			match self.map.keys.get(self.index) {
				Some(k) => {
					// Key within range
					match k {
						Some(kk) => {
							// Key is set
							
							// SAFETY: Guaranteed to be set, as we update key and value at the same time
							let v = unsafe { self.map.values.get_unchecked_mut(self.index).as_mut().unwrap_unchecked() };
							let v_ptr = v as *mut V;
							
							// SAFETY: Need unsafe pointer magic for IterMut
							let vv = unsafe { &mut *v_ptr };
							
							self.index += 1;
							return Some((kk, vv));
						}
						None => {
							// Key is unset
							self.index += 1;
							continue;
						}
					}
				}
				None => {
					// Key outside range
					return None;
				}
			}
		}
	}
}

pub struct Values<'m, 'e, V> {
	map: &'m CandidateMap<'e, V>,
	index: usize
}

impl<'m, 'e, V> Iterator for Values<'m, 'e, V> {
	type Item = &'m V;
	
	fn next(&mut self) -> Option<Self::Item> {
		loop {
			match self.map.values.get(self.index) {
				Some(v) => {
					// Key within range
					match v {
						Some(vv) => {
							// Key is set
							self.index += 1;
							return Some(vv);
						}
						None => {
							// Key is unset
							self.index += 1;
							continue;
						}
					}
				}
				None => {
					// Key outside range
					return None;
				}
			}
		}
	}
}

pub struct IntoIter<'e, V> {
	map: CandidateMap<'e, V>,
	index: usize
}

impl<'e, V> Iterator for IntoIter<'e, V> {
	type Item = (&'e Candidate, V);
	
	fn next(&mut self) -> Option<Self::Item> {
		loop {
			match self.map.keys.get(self.index) {
				Some(k) => {
					// Key within range
					match k {
						Some(kk) => {
							// Key is set
							
							// SAFETY: Guaranteed to be set, as we update key and value at the same time
							let v = unsafe { self.map.values.get_unchecked_mut(self.index).take().unwrap_unchecked() };
							self.index += 1;
							return Some((kk, v));
						}
						None => {
							// Key is unset
							self.index += 1;
							continue;
						}
					}
				}
				None => {
					// Key outside range
					return None;
				}
			}
		}
	}
}

impl<'e, V> IntoIterator for CandidateMap<'e, V> {
	type Item = (&'e Candidate, V);
	
	type IntoIter = IntoIter<'e, V>;
	
	fn into_iter(self) -> Self::IntoIter {
		return IntoIter { map: self, index: 0 };
	}
}