library/proc_macro/src/bridge/handle.rs

   1 //! Server-side handles and storage for per-handle data.
   2
   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};
   8
   9 pub(super) type Handle = NonZeroU32;
  10
  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>,
  16 }
  17
  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);
  23
  24         OwnedStore { counter, data: BTreeMap::new() }
  25     }
  26 }
  27
  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());
  33         handle
  34     }
  35
  36     pub(super) fn take(&mut self, h: Handle) -> T {
  37         self.data.remove(&h).expect("use-after-free in `proc_macro` handle")
  38     }
  39 }
  40
  41 impl<T> Index<Handle> for OwnedStore<T> {
  42     type Output = T;
  43     fn index(&self, h: Handle) -> &T {
  44         self.data.get(&h).expect("use-after-free in `proc_macro` handle")
  45     }
  46 }
  47
  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")
  51     }
  52 }
  53
  54 // HACK(eddyb) deterministic `std::collections::hash_map::RandomState` replacement
  55 // that doesn't require adding any dependencies to `proc_macro` (like `rustc-hash`).
  56 #[derive(Clone)]
  57 struct NonRandomState;
  58
  59 impl BuildHasher for NonRandomState {
  60     type Hasher = std::collections::hash_map::DefaultHasher;
  61     #[inline]
  62     fn build_hasher(&self) -> Self::Hasher {
  63         Self::Hasher::new()
  64     }
  65 }
  66
  67 /// Like `OwnedStore`, but avoids storing any value more than once.
  68 pub(super) struct InternedStore<T: 'static> {
  69     owned: OwnedStore<T>,
  70     interner: HashMap<T, Handle, NonRandomState>,
  71 }
  72
  73 impl<T: Copy + Eq + Hash> InternedStore<T> {
  74     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  75         InternedStore {
  76             owned: OwnedStore::new(counter),
  77             interner: HashMap::with_hasher(NonRandomState),
  78         }
  79     }
  80
  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))
  84     }
  85
  86     pub(super) fn copy(&mut self, h: Handle) -> T {
  87         self.owned[h]
  88     }
  89 }