tor_basic_utils/
intern.rs

1//! Declare types for interning various objects.
2
3use std::hash::Hash;
4use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
5use weak_table::WeakHashSet;
6
7/// An InternCache is a lazily-constructed weak set of objects.
8///
9/// Let's break that down!  It's "lazily constructed" because it
10/// doesn't actually allocate anything until you use it for the first
11/// time.  That allows it to have a const [`new`](InternCache::new)
12/// method, so you can make these static.
13///
14/// It's "weak" because it only holds weak references to its objects;
15/// once every strong reference is gone, the object is unallocated.
16/// Later, the hash entry is (lazily) removed.
17pub struct InternCache<T: ?Sized> {
18    /// Underlying hashset for interned objects
19    //
20    // TODO: If WeakHashSet::new is someday const, we can do away with OnceLock here.
21    cache: OnceLock<Mutex<WeakHashSet<Weak<T>>>>,
22}
23
24impl<T: ?Sized> InternCache<T> {
25    /// Create a new, empty, InternCache.
26    pub const fn new() -> Self {
27        InternCache {
28            cache: OnceLock::new(),
29        }
30    }
31}
32
33impl<T: ?Sized> Default for InternCache<T> {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl<T: Eq + Hash + ?Sized> InternCache<T> {
40    /// Helper: initialize the cache if needed, then lock it.
41    fn cache(&self) -> MutexGuard<'_, WeakHashSet<Weak<T>>> {
42        let cache = self.cache.get_or_init(|| Mutex::new(WeakHashSet::new()));
43        cache.lock().expect("Poisoned lock lock for cache")
44    }
45}
46
47impl<T: Eq + Hash> InternCache<T> {
48    /// Intern a given value into this cache.
49    ///
50    /// If `value` is already stored in this cache, we return a
51    /// reference to the stored value.  Otherwise, we insert `value`
52    /// into the cache, and return that.
53    pub fn intern(&self, value: T) -> Arc<T> {
54        let mut cache = self.cache();
55        if let Some(pp) = cache.get(&value) {
56            pp
57        } else {
58            let arc = Arc::new(value);
59            cache.insert(Arc::clone(&arc));
60            arc
61        }
62    }
63}
64
65impl<T: Hash + Eq + ?Sized> InternCache<T> {
66    /// Intern an object by reference.
67    ///
68    /// Works with unsized types, but requires that the reference implements
69    /// `Into<Arc<T>>`.
70    pub fn intern_ref<'a, V>(&self, value: &'a V) -> Arc<T>
71    where
72        V: Hash + Eq + ?Sized,
73        &'a V: Into<Arc<T>>,
74        T: std::borrow::Borrow<V>,
75    {
76        let mut cache = self.cache();
77        if let Some(arc) = cache.get(value) {
78            arc
79        } else {
80            let arc = value.into();
81            cache.insert(Arc::clone(&arc));
82            arc
83        }
84    }
85}
86
87#[cfg(test)]
88mod test {
89    // @@ begin test lint list maintained by maint/add_warning @@
90    #![allow(clippy::bool_assert_comparison)]
91    #![allow(clippy::clone_on_copy)]
92    #![allow(clippy::dbg_macro)]
93    #![allow(clippy::mixed_attributes_style)]
94    #![allow(clippy::print_stderr)]
95    #![allow(clippy::print_stdout)]
96    #![allow(clippy::single_char_pattern)]
97    #![allow(clippy::unwrap_used)]
98    #![allow(clippy::unchecked_time_subtraction)]
99    #![allow(clippy::useless_vec)]
100    #![allow(clippy::needless_pass_by_value)]
101    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
102    use super::*;
103
104    #[test]
105    fn interning_by_value() {
106        // "intern" case.
107        let c: InternCache<String> = InternCache::new();
108
109        let s1 = c.intern("abc".to_string());
110        let s2 = c.intern("def".to_string());
111        let s3 = c.intern("abc".to_string());
112        assert!(Arc::ptr_eq(&s1, &s3));
113        assert!(!Arc::ptr_eq(&s1, &s2));
114        assert_eq!(s2.as_ref(), "def");
115        assert_eq!(s3.as_ref(), "abc");
116    }
117
118    #[test]
119    fn interning_by_ref() {
120        // "intern" case.
121        let c: InternCache<str> = InternCache::new();
122
123        let s1 = c.intern_ref("abc");
124        let s2 = c.intern_ref("def");
125        let s3 = c.intern_ref("abc");
126        assert!(Arc::ptr_eq(&s1, &s3));
127        assert!(!Arc::ptr_eq(&s1, &s2));
128        assert_eq!(&*s2, "def");
129        assert_eq!(&*s3, "abc");
130    }
131}