library/proc_macro/src/bridge/handle.rs

   1 //! Server-side handles and storage for per-handle data.
   2
   3 use std::collections::BTreeMap;
   4 use std::hash::Hash;
   5 use std::num::NonZeroU32;
   6 use std::ops::{Index, IndexMut};
   7 use std::sync::atomic::{AtomicUsize, Ordering};
   8
   9 use super::fxhash::FxHashMap;
  10
  11 pub(super) type Handle = NonZeroU32;
  12
  13 /// A store that associates values of type `T` with numeric handles. A value can
  14 /// be looked up using its handle.
  15 pub(super) struct OwnedStore<T: 'static> {
  16     counter: &'static AtomicUsize,
  17     data: BTreeMap<Handle, T>,
  18 }
  19
  20 impl<T> OwnedStore<T> {
  21     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  22         // Ensure the handle counter isn't 0, which would panic later,
  23         // when `NonZeroU32::new` (aka `Handle::new`) is called in `alloc`.
  24         assert_ne!(counter.load(Ordering::SeqCst), 0);
  25
  26         OwnedStore { counter, data: BTreeMap::new() }
  27     }
  28 }
  29
  30 impl<T> OwnedStore<T> {
  31     pub(super) fn alloc(&mut self, x: T) -> Handle {
  32         let counter = self.counter.fetch_add(1, Ordering::SeqCst);
  33         let handle = Handle::new(counter as u32).expect("`proc_macro` handle counter overflowed");
  34         assert!(self.data.insert(handle, x).is_none());
  35         handle
  36     }
  37
  38     pub(super) fn take(&mut self, h: Handle) -> T {
  39         self.data.remove(&h).expect("use-after-free in `proc_macro` handle")
  40     }
  41 }
  42
  43 impl<T> Index<Handle> for OwnedStore<T> {
  44     type Output = T;
  45     fn index(&self, h: Handle) -> &T {
  46         self.data.get(&h).expect("use-after-free in `proc_macro` handle")
  47     }
  48 }
  49
  50 impl<T> IndexMut<Handle> for OwnedStore<T> {
  51     fn index_mut(&mut self, h: Handle) -> &mut T {
  52         self.data.get_mut(&h).expect("use-after-free in `proc_macro` handle")
  53     }
  54 }
  55
  56 /// Like `OwnedStore`, but avoids storing any value more than once.
  57 pub(super) struct InternedStore<T: 'static> {
  58     owned: OwnedStore<T>,
  59     interner: FxHashMap<T, Handle>,
  60 }
  61
  62 impl<T: Copy + Eq + Hash> InternedStore<T> {
  63     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  64         InternedStore { owned: OwnedStore::new(counter), interner: FxHashMap::default() }
  65     }
  66
  67     pub(super) fn alloc(&mut self, x: T) -> Handle {
  68         let owned = &mut self.owned;
  69         *self.interner.entry(x).or_insert_with(|| owned.alloc(x))
  70     }
  71
  72     pub(super) fn copy(&mut self, h: Handle) -> T {
  73         self.owned[h]
  74     }
  75 }