1 //! Server-side handles and storage for per-handle data.
3 use std::collections::{BTreeMap, HashMap};
4 use std::hash::{BuildHasher, Hash};
5 use std::num::NonZeroU32;
6 use std::ops::{Index, IndexMut};
7 use std::sync::atomic::{AtomicUsize, Ordering};
9 pub(super) type Handle = NonZeroU32;
11 /// A store that associates values of type `T` with numeric handles. A value can
12 /// be looked up using its handle.
13 pub(super) struct OwnedStore<T: 'static> {
14 counter: &'static AtomicUsize,
15 data: BTreeMap<Handle, T>,
18 impl<T> OwnedStore<T> {
19 pub(super) fn new(counter: &'static AtomicUsize) -> Self {
20 // Ensure the handle counter isn't 0, which would panic later,
21 // when `NonZeroU32::new` (aka `Handle::new`) is called in `alloc`.
22 assert_ne!(counter.load(Ordering::SeqCst), 0);
24 OwnedStore { counter, data: BTreeMap::new() }
28 impl<T> OwnedStore<T> {
29 pub(super) fn alloc(&mut self, x: T) -> Handle {
30 let counter = self.counter.fetch_add(1, Ordering::SeqCst);
31 let handle = Handle::new(counter as u32).expect("`proc_macro` handle counter overflowed");
32 assert!(self.data.insert(handle, x).is_none());
36 pub(super) fn take(&mut self, h: Handle) -> T {
37 self.data.remove(&h).expect("use-after-free in `proc_macro` handle")
41 impl<T> Index<Handle> for OwnedStore<T> {
43 fn index(&self, h: Handle) -> &T {
44 self.data.get(&h).expect("use-after-free in `proc_macro` handle")
48 impl<T> IndexMut<Handle> for OwnedStore<T> {
49 fn index_mut(&mut self, h: Handle) -> &mut T {
50 self.data.get_mut(&h).expect("use-after-free in `proc_macro` handle")
54 // HACK(eddyb) deterministic `std::collections::hash_map::RandomState` replacement
55 // that doesn't require adding any dependencies to `proc_macro` (like `rustc-hash`).
57 struct NonRandomState;
59 impl BuildHasher for NonRandomState {
60 type Hasher = std::collections::hash_map::DefaultHasher;
62 fn build_hasher(&self) -> Self::Hasher {
67 /// Like `OwnedStore`, but avoids storing any value more than once.
68 pub(super) struct InternedStore<T: 'static> {
70 interner: HashMap<T, Handle, NonRandomState>,
73 impl<T: Copy + Eq + Hash> InternedStore<T> {
74 pub(super) fn new(counter: &'static AtomicUsize) -> Self {
76 owned: OwnedStore::new(counter),
77 interner: HashMap::with_hasher(NonRandomState),
81 pub(super) fn alloc(&mut self, x: T) -> Handle {
82 let owned = &mut self.owned;
83 *self.interner.entry(x).or_insert_with(|| owned.alloc(x))
86 pub(super) fn copy(&mut self, h: Handle) -> T {