#[cfg(feature = "serde")]
mod serde;
use std::{
collections::HashMap,
iter::FromIterator,
ops::{AddAssign as _, Deref, Index},
};
use crate::{Entry, SmallSortedMap, ValuesIter};
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct SmallCounter<K, const SIZE: usize>
where
K: Ord,
{
storage: SmallSortedMap<K, usize, SIZE>,
}
impl<K, const SIZE: usize> SmallCounter<K, SIZE>
where
K: Ord,
{
#[inline]
pub fn new() -> Self {
Self {
storage: SmallSortedMap::new(),
}
}
#[inline]
pub fn add(&mut self, key: K) {
self.storage.entry(key).or_default().add_assign(1);
}
#[inline]
pub fn combine(&mut self, other: &mut Self) {
other.storage.drain(..).for_each(|(key, count)| {
self.storage.entry(key).or_default().add_assign(count);
});
}
#[inline]
pub fn remove(&mut self, key: K) -> bool {
match self.storage.entry(key) {
Entry::Occupied(mut entry) => {
let count = entry.get_mut();
if *count == 0 {
false
} else {
*count -= 1;
true
}
}
Entry::Vacant(..) => false,
}
}
#[inline]
pub fn get(&self, key: &K) -> &usize {
self.storage.get(key).unwrap_or(&0)
}
#[inline]
pub fn as_slice(&self) -> &[(K, usize)] {
self.storage.as_slice()
}
#[inline]
pub fn values(&self) -> ValuesIter<'_, K, usize> {
self.storage.values()
}
#[inline]
pub fn retain<F>(&mut self, f: F)
where
F: FnMut(&K, &mut usize) -> bool,
{
self.storage.retain(f)
}
}
impl<K, const SIZE: usize> Default for SmallCounter<K, SIZE>
where
K: Ord,
{
fn default() -> Self {
Self::new()
}
}
impl<K, const SIZE: usize> Index<K> for SmallCounter<K, SIZE>
where
K: Ord,
{
type Output = usize;
fn index(&self, index: K) -> &Self::Output {
self.get(&index)
}
}
impl<K, const SIZE: usize> Index<&K> for SmallCounter<K, SIZE>
where
K: Ord,
{
type Output = usize;
fn index(&self, index: &K) -> &Self::Output {
self.get(index)
}
}
impl<K, const SIZE: usize> AsRef<[(K, usize)]> for SmallCounter<K, SIZE>
where
K: Ord,
{
fn as_ref(&self) -> &[(K, usize)] {
self.as_slice()
}
}
impl<K, const SIZE: usize> Deref for SmallCounter<K, SIZE>
where
K: Ord,
{
type Target = [(K, usize)];
fn deref(&self) -> &Self::Target {
self.as_slice()
}
}
impl<K, const SIZE: usize> IntoIterator for SmallCounter<K, SIZE>
where
K: Ord,
{
type Item = (K, usize);
type IntoIter = smallvec::IntoIter<[(K, usize); SIZE]>;
fn into_iter(self) -> Self::IntoIter {
self.storage.into_iter()
}
}
impl<K, const SIZE: usize> Extend<K> for SmallCounter<K, SIZE>
where
K: Ord,
{
fn extend<T: IntoIterator<Item = K>>(&mut self, iter: T) {
iter.into_iter().for_each(|key| self.add(key));
}
}
impl<K, const SIZE: usize> FromIterator<K> for SmallCounter<K, SIZE>
where
K: Ord,
{
fn from_iter<T: IntoIterator<Item = K>>(iter: T) -> Self {
let mut counter = Self::new();
counter.extend(iter.into_iter());
counter
}
}
impl<K, const SIZE: usize> FromIterator<(K, usize)> for SmallCounter<K, SIZE>
where
K: Ord,
{
fn from_iter<T: IntoIterator<Item = (K, usize)>>(iter: T) -> Self {
Self {
storage: iter.into_iter().collect(),
}
}
}
impl<S, K, const SIZE: usize> From<HashMap<K, usize, S>> for SmallCounter<K, SIZE>
where
K: Ord,
{
fn from(map: HashMap<K, usize, S>) -> Self {
map.into_iter().collect()
}
}
#[cfg(test)]
mod tests {
use std::array;
use crate::SmallCounter;
#[test]
fn it_works() {
let mut counter = SmallCounter::<_, 10>::new();
counter.extend(array::IntoIter::new([10, 20, 10, 10, 20]));
assert_eq!(counter.len(), 2);
assert_eq!(counter[10], 3);
assert_eq!(counter[20], 2);
assert_eq!(counter[15], 0);
assert_eq!(counter.into_iter().collect::<Vec<_>>(), [(10, 3), (20, 2)]);
}
#[test]
fn test_from_iterator() {
let counter: SmallCounter<_, 10> = array::IntoIter::new([10, 20, 10, 10, 20]).collect();
assert_eq!(counter.into_iter().collect::<Vec<_>>(), [(10, 3), (20, 2)]);
}
#[test]
fn test_remove() {
let mut counter: SmallCounter<_, 10> = array::IntoIter::new([10, 20, 10, 10, 20]).collect();
assert!(counter.remove(10));
assert!(counter.remove(20));
assert!(counter.remove(20));
assert!(!counter.remove(15));
assert!(!counter.remove(25));
assert_eq!(counter[10], 2);
assert_eq!(counter[20], 0);
}
}