Skip to main content

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