circular_counter.rs

use std::collections::{HashMap, VecDeque};
use std::hash::Hash;

pub type Count = usize;

/// Ring-buffer based counter for O(1) checking of counts
/// of a value within the last `n` entries.
///
/// Probably overkill.
pub struct RecentCounts<T> {
	window_size: usize,
	queue: VecDeque<T>,
	counts: HashMap<T, Count>,
}

impl<T:Hash + Eq + Clone> RecentCounts<T> {
	pub fn new(window_size: usize) -> Self {
		RecentCounts {
			window_size: window_size,
			queue: VecDeque::new(),
			counts: HashMap::new(),
		}
	}

	pub fn insert(&mut self, x: T) -> Count {
		self.inc_count(&x);
		self.queue.push_back(x);
		if self.queue.len() > self.window_size as usize {
			self.forget_one();
		}
		// unwrap_or(0) is hit if the window_size is zero
		// (in which case all counts are clearly zero anyways)
		self.queue.back().map(|x| self.counts[x]).unwrap_or(0)
	}

	pub fn len(&self) -> usize { self.queue.len() }
	pub fn window_size(&self) -> usize { self.window_size }
	pub fn count(&self, x: &T) -> Count { self.counts[x] }
	pub fn clear(&mut self) {
		self.queue.clear();
		self.counts.clear();
	}
	fn forget_one(&mut self) {
		let x = self.queue.pop_front().expect("we push before we pop");
		self.dec_count(&x);
	}
	fn dec_count(&mut self, x: &T) {
		match self.counts[x] {
			1 => { self.counts.remove(x); },
			_ => { *self.counts.get_mut(x).unwrap() -= 1; },
		}
	}
	fn inc_count(&mut self, x: &T) {
		*self.counts.entry(x.clone()).or_insert(0) += 1;
	}
}