src/libproc_macro/bridge/handle.rs

   1 //! Server-side handles and storage for per-handle data.
   2
   3 use std::collections::{BTreeMap, HashMap};
   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 pub(super) type Handle = NonZeroU32;
  10
  11 pub(super) struct OwnedStore<T: 'static> {
  12     counter: &'static AtomicUsize,
  13     data: BTreeMap<Handle, T>,
  14 }
  15
  16 impl<T> OwnedStore<T> {
  17     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  18         // Ensure the handle counter isn't 0, which would panic later,
  19         // when `NonZeroU32::new` (aka `Handle::new`) is called in `alloc`.
  20         assert_ne!(counter.load(Ordering::SeqCst), 0);
  21
  22         OwnedStore {
  23             counter,
  24             data: BTreeMap::new(),
  25         }
  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
  39             .remove(&h)
  40             .expect("use-after-free in `proc_macro` handle")
  41     }
  42 }
  43
  44 impl<T> Index<Handle> for OwnedStore<T> {
  45     type Output = T;
  46     fn index(&self, h: Handle) -> &T {
  47         self.data
  48             .get(&h)
  49             .expect("use-after-free in `proc_macro` handle")
  50     }
  51 }
  52
  53 impl<T> IndexMut<Handle> for OwnedStore<T> {
  54     fn index_mut(&mut self, h: Handle) -> &mut T {
  55         self.data
  56             .get_mut(&h)
  57             .expect("use-after-free in `proc_macro` handle")
  58     }
  59 }
  60
  61 pub(super) struct InternedStore<T: 'static> {
  62     owned: OwnedStore<T>,
  63     interner: HashMap<T, Handle>,
  64 }
  65
  66 impl<T: Copy + Eq + Hash> InternedStore<T> {
  67     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  68         InternedStore {
  69             owned: OwnedStore::new(counter),
  70             interner: HashMap::new(),
  71         }
  72     }
  73
  74     pub(super) fn alloc(&mut self, x: T) -> Handle {
  75         let owned = &mut self.owned;
  76         *self.interner.entry(x).or_insert_with(|| owned.alloc(x))
  77     }
  78
  79     pub(super) fn copy(&mut self, h: Handle) -> T {
  80         self.owned[h]
  81     }
  82 }