1 // Copyright 2018 The Rust Project Developers. See the COPYRIGHT
2 // file at the top-level directory of this distribution and at
3 // http://rust-lang.org/COPYRIGHT.
5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8 // option. This file may not be copied, modified, or distributed
9 // except according to those terms.
11 //! Server-side handles and storage for per-handle data.
13 use std::collections::{BTreeMap, HashMap};
15 use std::num::NonZeroU32;
16 use std::ops::{Index, IndexMut};
17 use std::sync::atomic::{AtomicUsize, Ordering};
19 pub(super) type Handle = NonZeroU32;
21 pub(super) struct OwnedStore<T: 'static> {
22 counter: &'static AtomicUsize,
23 data: BTreeMap<Handle, T>,
26 impl<T> OwnedStore<T> {
27 pub(super) fn new(counter: &'static AtomicUsize) -> Self {
28 // Ensure the handle counter isn't 0, which would panic later,
29 // when `NonZeroU32::new` (aka `Handle::new`) is called in `alloc`.
30 assert_ne!(counter.load(Ordering::SeqCst), 0);
34 data: BTreeMap::new(),
39 impl<T> OwnedStore<T> {
40 pub(super) fn alloc(&mut self, x: T) -> Handle {
41 let counter = self.counter.fetch_add(1, Ordering::SeqCst);
42 let handle = Handle::new(counter as u32).expect("`proc_macro` handle counter overflowed");
43 assert!(self.data.insert(handle, x).is_none());
47 pub(super) fn take(&mut self, h: Handle) -> T {
50 .expect("use-after-free in `proc_macro` handle")
54 impl<T> Index<Handle> for OwnedStore<T> {
56 fn index(&self, h: Handle) -> &T {
59 .expect("use-after-free in `proc_macro` handle")
63 impl<T> IndexMut<Handle> for OwnedStore<T> {
64 fn index_mut(&mut self, h: Handle) -> &mut T {
67 .expect("use-after-free in `proc_macro` handle")
71 pub(super) struct InternedStore<T: 'static> {
73 interner: HashMap<T, Handle>,
76 impl<T: Copy + Eq + Hash> InternedStore<T> {
77 pub(super) fn new(counter: &'static AtomicUsize) -> Self {
79 owned: OwnedStore::new(counter),
80 interner: HashMap::new(),
84 pub(super) fn alloc(&mut self, x: T) -> Handle {
85 let owned = &mut self.owned;
86 *self.interner.entry(x).or_insert_with(|| owned.alloc(x))
89 pub(super) fn copy(&mut self, h: Handle) -> T {