boreholeio 0.1.0

A library for interacting with borehole.io, a subsurface data management, delivery and visualisation platform
Documentation
use super::Error;
use std::{cmp::Eq, collections::HashMap, fmt::Debug, hash::Hash};

/// A hash map with reference-counted keys. When a new key-value pair is
/// inserted into the map, the key's reference count is set to `1`. This count
/// can then be modified using the `increment(key)` and `decrement(key)`
/// methods. Once it reaches `0`, the key-value pair is removed from the map.
#[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(),
        }
    }

    /// Increments the reference count for the given key and returns the
    /// updated value. Returns `None` if the key was not found in the map.
    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
        }
    }

    /// Decrements the reference count for the given key and returns the
    /// updated value. If the reference count reaches `0`, the value is removed
    /// from the map and `0` is returned. Returns `None` if the key was not
    /// found in the map.
    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
    }

    /// Returns a reference to the value corresponding to the key. Errors out
    /// if the map does not contain the queried key.
    pub(super) fn get(&self, k: &K) -> Result<&V, Error> {
        self.map.get(k).ok_or(format!("Invalid key: {k:?}"))
    }

    /// Inserts the key-value pair into the map. Errors out if the map had
    /// previously the key present.
    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(())
        }
    }

    /// Returns reference count for the given key. If the key does not present
    /// in the map, returns `0`.
    #[cfg(test)]
    pub(super) fn reference_count(&self, k: &K) -> usize {
        *self.reference_count.get(k).unwrap_or(&0)
    }
}

/// This allows iterating through the `self.map` elements.
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()
    }
}