1//! Server-side handles and storage for per-handle data.
2
3use std::collections::BTreeMap;
4use std::hash::Hash;
5use std::num::NonZero;
6use std::ops::{Index, IndexMut};
7use std::sync::atomic::{AtomicU32, Ordering};
8
9use super::fxhash::FxHashMap;
10
11pub(super) type Handle = NonZero<u32>;
12
13/// A store that associates values of type `T` with numeric handles. A value can
14/// be looked up using its handle.
15pub(super) struct OwnedStore<T: 'static> {
16 counter: &'static AtomicU32,
17 data: BTreeMap<Handle, T>,
18}
19
20impl<T> OwnedStore<T> {
21 pub(super) fn new(counter: &'static AtomicU32) -> Self {
22 // Ensure the handle counter isn't 0, which would panic later,
23 // when `NonZero::new` (aka `Handle::new`) is called in `alloc`.
24 assert_ne!(counter.load(Ordering::Relaxed), 0);
25
26 OwnedStore { counter, data: BTreeMap::new() }
27 }
28}
29
30impl<T> OwnedStore<T> {
31 pub(super) fn alloc(&mut self, x: T) -> Handle {
32 let counter: u32 = self.counter.fetch_add(val:1, order:Ordering::Relaxed);
33 let handle: NonZero = Handle::new(counter).expect(msg:"`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(msg:"use-after-free in `proc_macro` handle")
40 }
41}
42
43impl<T> Index<Handle> for OwnedStore<T> {
44 type Output = T;
45 fn index(&self, h: Handle) -> &T {
46 self.data.get(&h).expect(msg:"use-after-free in `proc_macro` handle")
47 }
48}
49
50impl<T> IndexMut<Handle> for OwnedStore<T> {
51 fn index_mut(&mut self, h: Handle) -> &mut T {
52 self.data.get_mut(&h).expect(msg:"use-after-free in `proc_macro` handle")
53 }
54}
55
56/// Like `OwnedStore`, but avoids storing any value more than once.
57pub(super) struct InternedStore<T: 'static> {
58 owned: OwnedStore<T>,
59 interner: FxHashMap<T, Handle>,
60}
61
62impl<T: Copy + Eq + Hash> InternedStore<T> {
63 pub(super) fn new(counter: &'static AtomicU32) -> 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 OwnedStore = &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}
76