crates/proc_macro_srv/src/proc_macro/bridge/handle.rs

   1 //! lib-proc-macro Server-side handles and storage for per-handle data.
   2 //!
   3 //! Copy from https://github.com/rust-lang/rust/blob/6050e523bae6de61de4e060facc43dc512adaccd/src/libproc_macro/bridge/handle.rs
   4 //! augmented with removing unstable features
   5
   6 use std::collections::{BTreeMap, HashMap};
   7 use std::hash::Hash;
   8 use std::num::NonZeroU32;
   9 use std::ops::{Index, IndexMut};
  10 use std::sync::atomic::{AtomicUsize, Ordering};
  11
  12 pub(super) type Handle = NonZeroU32;
  13
  14 pub(super) struct OwnedStore<T: 'static> {
  15     counter: &'static AtomicUsize,
  16     data: BTreeMap<Handle, T>,
  17 }
  18
  19 impl<T> OwnedStore<T> {
  20     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  21         // Ensure the handle counter isn't 0, which would panic later,
  22         // when `NonZeroU32::new` (aka `Handle::new`) is called in `alloc`.
  23         assert_ne!(counter.load(Ordering::SeqCst), 0);
  24
  25         OwnedStore { counter, data: BTreeMap::new() }
  26     }
  27 }
  28
  29 impl<T> OwnedStore<T> {
  30     pub(super) fn alloc(&mut self, x: T) -> Handle {
  31         let counter = self.counter.fetch_add(1, Ordering::SeqCst);
  32         let handle = Handle::new(counter as u32).expect("`proc_macro` handle counter overflowed");
  33         assert!(self.data.insert(handle, x).is_none());
  34         handle
  35     }
  36
  37     pub(super) fn take(&mut self, h: Handle) -> T {
  38         self.data.remove(&h).expect("use-after-free in `proc_macro` handle")
  39     }
  40 }
  41
  42 impl<T> Index<Handle> for OwnedStore<T> {
  43     type Output = T;
  44     fn index(&self, h: Handle) -> &T {
  45         self.data.get(&h).expect("use-after-free in `proc_macro` handle")
  46     }
  47 }
  48
  49 impl<T> IndexMut<Handle> for OwnedStore<T> {
  50     fn index_mut(&mut self, h: Handle) -> &mut T {
  51         self.data.get_mut(&h).expect("use-after-free in `proc_macro` handle")
  52     }
  53 }
  54
  55 pub(super) struct InternedStore<T: 'static> {
  56     owned: OwnedStore<T>,
  57     interner: HashMap<T, Handle>,
  58 }
  59
  60 impl<T: Copy + Eq + Hash> InternedStore<T> {
  61     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  62         InternedStore { owned: OwnedStore::new(counter), interner: HashMap::new() }
  63     }
  64
  65     pub(super) fn alloc(&mut self, x: T) -> Handle {
  66         let owned = &mut self.owned;
  67         *self.interner.entry(x).or_insert_with(|| owned.alloc(x))
  68     }
  69
  70     pub(super) fn copy(&mut self, h: Handle) -> T {
  71         self.owned[h]
  72     }
  73 }