use super::Error;
use std::{cmp::Eq, collections::HashMap, fmt::Debug, hash::Hash};
#[derive(Debug)]
pub(super) struct ReferenceCountingHashMap<K, V>
where
K: Debug + Clone + Eq + Hash,
{
map: HashMap<K, V>,
reference_count: HashMap<K, usize>,
}
impl<K, V> ReferenceCountingHashMap<K, V>
where
K: Debug + Clone + Eq + Hash,
{
pub(super) fn new() -> Self {
Self {
map: HashMap::new(),
reference_count: HashMap::new(),
}
}
pub(super) fn increment(&mut self, k: &K) -> Option<usize> {
if let Some(rc) = self.reference_count.get_mut(k) {
*rc += 1;
Some(*rc)
} else {
None
}
}
pub(super) fn decrement(&mut self, k: &K) -> Option<usize> {
let rc = if let Some(rc) = self.reference_count.get_mut(k) {
*rc -= 1;
Some(*rc)
} else {
None
};
if rc == Some(0) {
self.map.remove(k);
self.reference_count.remove(k);
}
rc
}
pub(super) fn get(&self, k: &K) -> Result<&V, Error> {
self.map.get(k).ok_or(format!("Invalid key: {k:?}"))
}
pub(super) fn insert(&mut self, k: K, v: V) -> Result<(), Error> {
if self.map.contains_key(&k) || self.reference_count.contains_key(&k) {
Err(format!(
"ReferenceCountingHashMap already contains the key {k:?}"
))
} else {
self.map.insert(k.clone(), v);
self.reference_count.insert(k, 1);
Ok(())
}
}
#[cfg(test)]
pub(super) fn reference_count(&self, k: &K) -> usize {
*self.reference_count.get(k).unwrap_or(&0)
}
}
impl<'a, K, V> IntoIterator for &'a ReferenceCountingHashMap<K, V>
where
K: Debug + Clone + Eq + Hash,
{
type Item = (&'a K, &'a V);
type IntoIter = std::collections::hash_map::Iter<'a, K, V>;
fn into_iter(self) -> Self::IntoIter {
self.map.iter()
}
}