tinymist_analysis/adt/
interner.rs

1//! Global `Arc`-based object interning infrastructure.
2//!
3//! Eventually this should probably be replaced with salsa-based interning.
4//!
5//! todo: This is less efficient as the arc object will change its reference
6//! count every time it is cloned. todo: we may be able to optimize use by
7//! following approach:
8//! ```plain
9//! fn run_analyze(f) {
10//!   let local = thread_local_intern();
11//!   let res = f(local);
12//!   std::thread::spawn(move || gc(local));
13//!   return res
14//! }
15//! ```
16//! However, this is out of scope for now.
17
18use std::{
19    fmt::{self, Debug, Display},
20    hash::{BuildHasherDefault, Hash, Hasher},
21    ops::Deref,
22    sync::{LazyLock, OnceLock},
23};
24
25use dashmap::{DashMap, SharedValue};
26use ecow::{EcoString, EcoVec};
27use hashbrown::{HashMap, hash_map::RawEntryMut};
28use parking_lot::Mutex;
29use rustc_hash::FxHasher;
30use triomphe::Arc;
31use typst::{foundations::Str, syntax::ast::Ident};
32
33type InternMap<T> = DashMap<Arc<T>, (), BuildHasherDefault<FxHasher>>;
34type Guard<T> = dashmap::RwLockWriteGuard<
35    'static,
36    HashMap<Arc<T>, SharedValue<()>, BuildHasherDefault<FxHasher>>,
37>;
38
39// https://news.ycombinator.com/item?id=22220342
40
41pub struct Interned<T: Internable + ?Sized> {
42    arc: Arc<T>,
43}
44
45impl<T: Internable> Interned<T> {
46    pub fn new(obj: T) -> Self {
47        let (mut shard, hash) = Self::select(&obj);
48        // Atomically,
49        // - check if `obj` is already in the map
50        //   - if so, clone its `Arc` and return it
51        //   - if not, box it up, insert it, and return a clone
52        // This needs to be atomic (locking the shard) to avoid races with other thread,
53        // which could insert the same object between us looking it up and
54        // inserting it.
55        match shard.raw_entry_mut().from_key_hashed_nocheck(hash, &obj) {
56            RawEntryMut::Occupied(occ) => Self {
57                arc: occ.key().clone(),
58            },
59            RawEntryMut::Vacant(vac) => {
60                T::storage().alloc().increment();
61                Self {
62                    arc: vac
63                        .insert_hashed_nocheck(hash, Arc::new(obj), SharedValue::new(()))
64                        .0
65                        .clone(),
66                }
67            }
68        }
69    }
70}
71
72// Note: It is dangerous to keep interned object temporarily (u128)
73// Case:
74// ```
75// insert(hash(Interned::new_str("a"))) == true
76// insert(hash(Interned::new_str("a"))) == true
77// ```
78impl Interned<str> {
79    pub fn new_str(s: &str) -> Self {
80        let (mut shard, hash) = Self::select(s);
81        // Atomically,
82        // - check if `obj` is already in the map
83        //   - if so, clone its `Arc` and return it
84        //   - if not, box it up, insert it, and return a clone
85        // This needs to be atomic (locking the shard) to avoid races with other thread,
86        // which could insert the same object between us looking it up and
87        // inserting it.
88        match shard.raw_entry_mut().from_key_hashed_nocheck(hash, s) {
89            RawEntryMut::Occupied(occ) => Self {
90                arc: occ.key().clone(),
91            },
92            RawEntryMut::Vacant(vac) => {
93                str::storage().alloc().increment();
94
95                Self {
96                    arc: vac
97                        .insert_hashed_nocheck(hash, Arc::from(s), SharedValue::new(()))
98                        .0
99                        .clone(),
100                }
101            }
102        }
103    }
104}
105
106static EMPTY: LazyLock<Interned<str>> = LazyLock::new(|| Interned::new_str(""));
107impl Default for Interned<str> {
108    fn default() -> Self {
109        EMPTY.clone()
110    }
111}
112
113impl Interned<str> {
114    pub fn empty() -> &'static Self {
115        &EMPTY
116    }
117}
118
119impl From<&str> for Interned<str> {
120    fn from(s: &str) -> Self {
121        Interned::new_str(s)
122    }
123}
124
125impl From<Str> for Interned<str> {
126    fn from(s: Str) -> Self {
127        Interned::new_str(&s)
128    }
129}
130
131impl From<EcoString> for Interned<str> {
132    fn from(s: EcoString) -> Self {
133        Interned::new_str(&s)
134    }
135}
136
137impl From<&EcoString> for Interned<str> {
138    fn from(s: &EcoString) -> Self {
139        Interned::new_str(s)
140    }
141}
142
143impl From<Ident<'_>> for Interned<str> {
144    fn from(s: Ident<'_>) -> Self {
145        Interned::new_str(s.get())
146    }
147}
148
149impl From<&Interned<str>> for EcoString {
150    fn from(s: &Interned<str>) -> Self {
151        s.as_ref().into()
152    }
153}
154
155impl<T: Internable> From<T> for Interned<T> {
156    fn from(s: T) -> Self {
157        Interned::new(s)
158    }
159}
160
161impl<T: Internable + Clone> From<&T> for Interned<T> {
162    fn from(s: &T) -> Self {
163        Interned::new(s.clone())
164    }
165}
166
167impl<T: Internable + ?Sized> Interned<T> {
168    #[inline]
169    fn select(obj: &T) -> (Guard<T>, u64) {
170        let storage = T::storage().get();
171        let hash = {
172            let mut hasher = std::hash::BuildHasher::build_hasher(storage.hasher());
173            obj.hash(&mut hasher);
174            hasher.finish()
175        };
176        let shard_idx = storage.determine_shard(hash as usize);
177        let shard = &storage.shards()[shard_idx];
178        (shard.write(), hash)
179    }
180}
181
182impl<T: Internable + ?Sized> Drop for Interned<T> {
183    #[inline]
184    fn drop(&mut self) {
185        // When the last `Ref` is dropped, remove the object from the global map.
186        if Arc::count(&self.arc) == 2 {
187            // Only `self` and the global map point to the object.
188
189            self.drop_slow();
190        }
191    }
192}
193
194impl<T: Internable + ?Sized> Interned<T> {
195    #[cold]
196    fn drop_slow(&mut self) {
197        let (mut shard, hash) = Self::select(&self.arc);
198
199        if Arc::count(&self.arc) != 2 {
200            // Another thread has interned another copy
201            return;
202        }
203
204        match shard
205            .raw_entry_mut()
206            .from_key_hashed_nocheck(hash, &self.arc)
207        {
208            RawEntryMut::Occupied(occ) => occ.remove(),
209            RawEntryMut::Vacant(_) => unreachable!(),
210        };
211
212        T::storage().alloc().decrement();
213
214        // Shrink the backing storage if the shard is less than 50% occupied.
215        if shard.len() * 2 < shard.capacity() {
216            shard.shrink_to_fit();
217        }
218    }
219}
220
221/// Compares interned `Ref`s using pointer equality.
222impl<T: Internable> PartialEq for Interned<T> {
223    // NOTE: No `?Sized` because `ptr_eq` doesn't work right with trait objects.
224
225    #[inline]
226    fn eq(&self, other: &Self) -> bool {
227        Arc::ptr_eq(&self.arc, &other.arc)
228    }
229}
230
231impl<T: Internable> Eq for Interned<T> {}
232
233impl<T: Internable + PartialOrd> PartialOrd for Interned<T> {
234    #[inline]
235    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
236        if self == other {
237            Some(std::cmp::Ordering::Equal)
238        } else {
239            self.as_ref().partial_cmp(other.as_ref())
240        }
241    }
242}
243
244impl<T: Internable + Ord> Ord for Interned<T> {
245    #[inline]
246    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
247        if self == other {
248            std::cmp::Ordering::Equal
249        } else {
250            self.as_ref().cmp(other.as_ref())
251        }
252    }
253}
254
255impl PartialOrd for Interned<str> {
256    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
257        Some(self.cmp(other))
258    }
259}
260
261impl Ord for Interned<str> {
262    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
263        if self == other {
264            std::cmp::Ordering::Equal
265        } else {
266            self.as_ref().cmp(other.as_ref())
267        }
268    }
269}
270
271impl PartialEq for Interned<str> {
272    fn eq(&self, other: &Self) -> bool {
273        Arc::ptr_eq(&self.arc, &other.arc)
274    }
275}
276
277impl Eq for Interned<str> {}
278
279impl serde::Serialize for Interned<str> {
280    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
281        self.arc.serialize(serializer)
282    }
283}
284
285impl<'de> serde::Deserialize<'de> for Interned<str> {
286    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
287        struct StrVisitor;
288
289        impl serde::de::Visitor<'_> for StrVisitor {
290            type Value = Interned<str>;
291
292            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
293                formatter.write_str("a string")
294            }
295
296            fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
297                Ok(Interned::new_str(v))
298            }
299        }
300
301        deserializer.deserialize_str(StrVisitor)
302    }
303}
304
305impl<T: Internable + ?Sized> Hash for Interned<T> {
306    fn hash<H: Hasher>(&self, state: &mut H) {
307        // NOTE: Cast disposes vtable pointer / slice/str length.
308        state.write_usize(Arc::as_ptr(&self.arc) as *const () as usize)
309    }
310}
311
312impl<T: Internable + ?Sized> AsRef<T> for Interned<T> {
313    #[inline]
314    fn as_ref(&self) -> &T {
315        &self.arc
316    }
317}
318
319impl<T: Internable + ?Sized> Deref for Interned<T> {
320    type Target = T;
321
322    #[inline]
323    fn deref(&self) -> &Self::Target {
324        &self.arc
325    }
326}
327
328impl<T: Internable + ?Sized> Clone for Interned<T> {
329    fn clone(&self) -> Self {
330        Self {
331            arc: self.arc.clone(),
332        }
333    }
334}
335
336impl<T: Debug + Internable + ?Sized> Debug for Interned<T> {
337    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
338        (*self.arc).fmt(f)
339    }
340}
341
342impl<T: Display + Internable + ?Sized> Display for Interned<T> {
343    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
344        (*self.arc).fmt(f)
345    }
346}
347
348pub static MAPS: Mutex<EcoVec<(&'static str, usize, Arc<AllocStats>)>> = Mutex::new(EcoVec::new());
349
350pub struct InternStorage<T: ?Sized> {
351    alloc: OnceLock<Arc<AllocStats>>,
352    map: OnceLock<InternMap<T>>,
353}
354
355#[allow(clippy::new_without_default)] // this a const fn, so it can't be default
356impl<T: InternSize + ?Sized> InternStorage<T> {
357    const SIZE: usize = T::INTERN_SIZE;
358
359    pub const fn new() -> Self {
360        Self {
361            alloc: OnceLock::new(),
362            map: OnceLock::new(),
363        }
364    }
365}
366
367impl<T: Internable + ?Sized> InternStorage<T> {
368    fn alloc(&self) -> &Arc<AllocStats> {
369        self.alloc.get_or_init(Arc::default)
370    }
371
372    fn get(&self) -> &InternMap<T> {
373        self.map.get_or_init(|| {
374            MAPS.lock()
375                .push((std::any::type_name::<T>(), Self::SIZE, self.alloc().clone()));
376            DashMap::default()
377        })
378    }
379}
380
381pub trait InternSize {
382    const INTERN_SIZE: usize;
383}
384
385impl<T: Sized> InternSize for T {
386    const INTERN_SIZE: usize = std::mem::size_of::<T>();
387}
388
389impl InternSize for str {
390    const INTERN_SIZE: usize = std::mem::size_of::<usize>() * 2;
391}
392
393pub trait Internable: InternSize + Hash + Eq + 'static {
394    fn storage() -> &'static InternStorage<Self>;
395}
396
397/// Implements `Internable` for a given list of types, making them usable with
398/// `Interned`.
399#[macro_export]
400#[doc(hidden)]
401macro_rules! _impl_internable {
402    ( $($t:ty),+ $(,)? ) => { $(
403        impl $crate::adt::interner::Internable for $t {
404            fn storage() -> &'static $crate::adt::interner::InternStorage<Self> {
405                static STORAGE: $crate::adt::interner::InternStorage<$t> = $crate::adt::interner::InternStorage::new();
406                &STORAGE
407            }
408        }
409    )+ };
410}
411
412pub use crate::_impl_internable as impl_internable;
413use crate::stats::AllocStats;
414
415impl_internable!(str,);
416
417#[cfg(test)]
418mod tests {
419    use std::cmp::Ordering;
420    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
421
422    use super::*;
423
424    static PARTIAL_CMP_CALLS: AtomicUsize = AtomicUsize::new(0);
425
426    #[derive(Eq, Hash, PartialEq)]
427    struct Counted(u8);
428
429    impl PartialOrd for Counted {
430        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
431            PARTIAL_CMP_CALLS.fetch_add(1, AtomicOrdering::Relaxed);
432            self.0.partial_cmp(&other.0)
433        }
434    }
435
436    impl_internable!(Counted);
437
438    #[test]
439    fn partial_cmp_short_circuits_shared_values() {
440        let value = Interned::new(Counted(1));
441
442        PARTIAL_CMP_CALLS.store(0, AtomicOrdering::Relaxed);
443        assert_eq!(value.partial_cmp(&value.clone()), Some(Ordering::Equal));
444        assert_eq!(PARTIAL_CMP_CALLS.load(AtomicOrdering::Relaxed), 0);
445
446        let other = Interned::new(Counted(2));
447        assert_eq!(value.partial_cmp(&other), Some(Ordering::Less));
448        assert_eq!(PARTIAL_CMP_CALLS.load(AtomicOrdering::Relaxed), 1);
449    }
450}
tinymist_analysis/adt/interner.rs

tinymist_analysis/adt/
interner.rs
