handle.rs - source
Description: Source of the Rust file `library/proc_macro/src/bridge/handle.rs`.

#
proc_macro/bridge/
handle.rs

    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 = self.counter.fetch_add(1, Ordering::Relaxed);
    33        let handle = Handle::new(counter).expect("`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("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("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("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 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}