use std::collections::HashMap;
use super::error::RefMapRemoveError;
pub struct RefMap {
references: HashMap<String, u64>,
}
impl RefMap {
pub fn new() -> Self {
RefMap {
references: HashMap::new(),
}
}
pub fn add_ref(&mut self, ref_id: String) -> u64 {
if let Some(ref_count) = self.references.remove(&ref_id) {
let new_ref_count = ref_count + 1;
self.references.insert(ref_id, new_ref_count);
new_ref_count
} else {
self.references.insert(ref_id, 1);
1
}
}
pub fn remove_ref(&mut self, ref_id: &str) -> Result<Option<String>, RefMapRemoveError> {
let ref_count = match self.references.remove(ref_id) {
Some(ref_count) => ref_count,
None => {
return Err(RefMapRemoveError(format!(
"Trying to remove a reference that does not exist: {}",
ref_id
)))
}
};
if ref_count == 1 {
Ok(Some(ref_id.into()))
} else {
self.references.insert(ref_id.into(), ref_count - 1);
Ok(None)
}
}
}
#[cfg(test)]
pub mod tests {
use super::*;
#[test]
fn test_add_ref() {
let mut ref_map = RefMap::new();
let ref_count = ref_map.add_ref("test_id".to_string());
assert_eq!(ref_count, 1);
let ref_count = ref_map.add_ref("test_id".to_string());
assert_eq!(ref_count, 2);
let ref_count = ref_map.add_ref("test_id_2".to_string());
assert_eq!(ref_count, 1);
}
#[test]
fn test_remove_ref() {
let mut ref_map = RefMap::new();
let ref_count = ref_map.add_ref("test_id".to_string());
assert_eq!(ref_count, 1);
let ref_count = ref_map.add_ref("test_id".to_string());
assert_eq!(ref_count, 2);
let id = ref_map.remove_ref("test_id");
assert_eq!(id, Ok(None));
assert_eq!(ref_map.references.get("test_id").cloned(), Some(1 as u64));
let id = ref_map.remove_ref("test_id");
assert_eq!(id, Ok(Some("test_id".to_string())));
assert_eq!(ref_map.references.get("test_id"), None);
}
#[test]
fn test_remove_ref_err() {
let mut ref_map = RefMap::new();
if let Ok(_) = ref_map.remove_ref("test_id") {
panic!("remove_ref should have returned an error");
}
}
}