Skip to main content

ferroui_base/utilities/
ref_tracking_dictionary.rs

1use std::collections::HashMap;
2use std::hash::Hash;
3use std::ops::{Deref, DerefMut};
4
5/// Maintains a set of objects with reference counts
6pub struct RefTrackingDictionary<TKey> {
7    base: HashMap<TKey, i32>,
8}
9
10impl<TKey: Eq + Hash> Default for RefTrackingDictionary<TKey> {
11    fn default() -> Self {
12        Self::new()
13    }
14}
15
16impl<TKey: Eq + Hash> RefTrackingDictionary<TKey> {
17    pub fn new() -> Self {
18        Self { base: HashMap::new() }
19    }
20
21    /// Increase reference count for a key by 1.
22    ///
23    /// Returns true if key was added to the dictionary, false otherwise.
24    pub fn add_ref(&mut self, key: TKey) -> bool {
25        let count = self.base.entry(key).or_insert(0);
26        *count += 1;
27        *count == 1
28    }
29
30    /// Decrease reference count for a key by 1.
31    ///
32    /// Returns true if key was removed to the dictionary, false otherwise.
33    /// A key without a reference panics in a debug build and returns false
34    /// otherwise.
35    pub fn release_ref(&mut self, key: &TKey) -> bool {
36        let remaining = match self.base.get_mut(key) {
37            Some(count) => {
38                *count -= 1;
39                *count
40            }
41            None => {
42                if cfg!(debug_assertions) {
43                    panic!("Attempting to release a non-referenced object");
44                }
45                return false;
46            }
47        };
48        if remaining == 0 {
49            self.base.remove(key);
50            return true;
51        }
52
53        false
54    }
55}
56
57impl<TKey> Deref for RefTrackingDictionary<TKey> {
58    type Target = HashMap<TKey, i32>;
59
60    fn deref(&self) -> &Self::Target {
61        &self.base
62    }
63}
64
65impl<TKey> DerefMut for RefTrackingDictionary<TKey> {
66    fn deref_mut(&mut self) -> &mut Self::Target {
67        &mut self.base
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    // Not from upstream: the original has no tests of this type.
74    use super::*;
75
76    #[test]
77    fn counts_references() {
78        let mut target = RefTrackingDictionary::<&'static str>::new();
79
80        assert!(target.add_ref("a"));
81        assert!(!target.add_ref("a"));
82        assert!(target.add_ref("b"));
83        assert_eq!(Some(&2), target.get("a"));
84        assert_eq!(2, target.len());
85
86        assert!(!target.release_ref(&"a"));
87        assert!(target.release_ref(&"a"));
88        assert!(!target.contains_key("a"));
89        assert!(target.release_ref(&"b"));
90        assert!(target.is_empty());
91    }
92
93    #[test]
94    #[cfg(debug_assertions)]
95    #[should_panic(expected = "Attempting to release a non-referenced object")]
96    fn release_of_a_key_without_reference_throws() {
97        let mut target = RefTrackingDictionary::<&'static str>::new();
98        target.release_ref(&"a");
99    }
100}