104 lines
2.9 KiB
Rust
104 lines
2.9 KiB
Rust
use crate::feed_message::Ranking;
|
|
use std::collections::HashMap;
|
|
|
|
/// A data structure which counts how many occurrences of a given key we've seen.
|
|
#[derive(Default)]
|
|
pub struct Counter<K> {
|
|
/// A map containing the number of occurrences of a given key.
|
|
///
|
|
/// If there are none then the entry is removed.
|
|
map: HashMap<K, u64>,
|
|
|
|
/// The number of occurrences where the key is `None`.
|
|
empty: u64,
|
|
}
|
|
|
|
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
|
|
pub enum CounterValue {
|
|
Increment,
|
|
Decrement,
|
|
}
|
|
|
|
impl<K> Counter<K>
|
|
where
|
|
K: Sized + std::hash::Hash + Eq,
|
|
{
|
|
/// Either adds or removes a single occurence of a given `key`.
|
|
pub fn modify<'a, Q>(&mut self, key: Option<&'a Q>, op: CounterValue)
|
|
where
|
|
Q: ?Sized + std::hash::Hash + Eq,
|
|
K: std::borrow::Borrow<Q>,
|
|
Q: std::borrow::ToOwned<Owned = K>,
|
|
{
|
|
if let Some(key) = key {
|
|
if let Some(entry) = self.map.get_mut(key) {
|
|
match op {
|
|
CounterValue::Increment => {
|
|
*entry += 1;
|
|
}
|
|
CounterValue::Decrement => {
|
|
*entry -= 1;
|
|
if *entry == 0 {
|
|
// Don't keep entries for which there are no hits.
|
|
self.map.remove(key);
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
assert_eq!(op, CounterValue::Increment);
|
|
self.map.insert(key.to_owned(), 1);
|
|
}
|
|
} else {
|
|
match op {
|
|
CounterValue::Increment => {
|
|
self.empty += 1;
|
|
}
|
|
CounterValue::Decrement => {
|
|
self.empty -= 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Generates a top-N table of the most common keys.
|
|
pub fn generate_ranking_top(&self, max_count: usize) -> Ranking<K>
|
|
where
|
|
K: Clone,
|
|
{
|
|
let mut all: Vec<(&K, u64)> = self.map.iter().map(|(key, count)| (key, *count)).collect();
|
|
all.sort_unstable_by_key(|&(_, count)| !count);
|
|
|
|
let list = all
|
|
.iter()
|
|
.take(max_count)
|
|
.map(|&(key, count)| (key.clone(), count))
|
|
.collect();
|
|
|
|
let other = all
|
|
.iter()
|
|
.skip(max_count)
|
|
.fold(0, |sum, (_, count)| sum + *count);
|
|
|
|
Ranking {
|
|
list,
|
|
other,
|
|
unknown: self.empty,
|
|
}
|
|
}
|
|
|
|
/// Generates a sorted table of all of the keys.
|
|
pub fn generate_ranking_ordered(&self) -> Ranking<K>
|
|
where
|
|
K: Copy + Clone + Ord,
|
|
{
|
|
let mut list: Vec<(K, u64)> = self.map.iter().map(|(key, count)| (*key, *count)).collect();
|
|
list.sort_unstable_by_key(|&(key, count)| (key, !count));
|
|
|
|
Ranking {
|
|
list,
|
|
other: 0,
|
|
unknown: self.empty,
|
|
}
|
|
}
|
|
}
|