diff options
author | Dmitry <[email protected]> | 2020-08-14 19:32:05 +0100 |
---|---|---|
committer | Dmitry <[email protected]> | 2020-08-14 19:32:05 +0100 |
commit | 178c3e135a2a249692f7784712492e7884ae0c00 (patch) | |
tree | ac6b769dbf7162150caa0c1624786a4dd79ff3be /crates/proc_macro_srv/src/proc_macro/bridge/handle.rs | |
parent | 06ff8e6c760ff05f10e868b5d1f9d79e42fbb49c (diff) | |
parent | c2594daf2974dbd4ce3d9b7ec72481764abaceb5 (diff) |
Merge remote-tracking branch 'origin/master'
Diffstat (limited to 'crates/proc_macro_srv/src/proc_macro/bridge/handle.rs')
-rw-r--r-- | crates/proc_macro_srv/src/proc_macro/bridge/handle.rs | 73 |
1 files changed, 73 insertions, 0 deletions
diff --git a/crates/proc_macro_srv/src/proc_macro/bridge/handle.rs b/crates/proc_macro_srv/src/proc_macro/bridge/handle.rs new file mode 100644 index 000000000..a2f77b5ac --- /dev/null +++ b/crates/proc_macro_srv/src/proc_macro/bridge/handle.rs | |||
@@ -0,0 +1,73 @@ | |||
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 | } | ||