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

#
proc_macro/bridge/
symbol.rs

    1//! Client-side interner used for symbols.
    2//!
    3//! This is roughly based on the symbol interner from `rustc_span` and the
    4//! DroplessArena from `rustc_arena`. It is unfortunately a complete
    5//! copy/re-implementation rather than a dependency as it is difficult to depend
    6//! on crates from within `proc_macro`, due to it being built at the same time
    7//! as `std`.
    8//!
    9//! If at some point in the future it becomes easier to add dependencies to
    10//! proc_macro, this module should probably be removed or simplified.
    11
    12use std::cell::RefCell;
    13use std::num::NonZero;
    14
    15use super::*;
    16
    17/// Handle for a symbol string stored within the Interner.
    18#[derive(Copy, Clone, PartialEq, Eq, Hash)]
    19pub struct Symbol(NonZero<u32>);
    20
    21impl !Send for Symbol {}
    22impl !Sync for Symbol {}
    23
    24impl Symbol {
    25    /// Intern a new `Symbol`
    26    pub(crate) fn new(string: &str) -> Self {
    27        INTERNER.with_borrow_mut(|i| i.intern(string))
    28    }
    29
    30    /// Creates a new `Symbol` for an identifier.
    31    ///
    32    /// Validates and normalizes before converting it to a symbol.
    33    pub(crate) fn new_ident(string: &str, is_raw: bool) -> Self {
    34        // Fast-path: check if this is a valid ASCII identifier
    35        if Self::is_valid_ascii_ident(string.as_bytes()) || string == "$crate" {
    36            if is_raw && !Self::can_be_raw(string) {
    37                panic!("`{}` cannot be a raw identifier", string);
    38            }
    39            return Self::new(string);
    40        }
    41
    42        // Slow-path: If the string is already ASCII we're done, otherwise ask
    43        // our server to do this for us over RPC.
    44        // We don't need to check for identifiers which can't be raw here,
    45        // because all of them are ASCII.
    46        if string.is_ascii() {
    47            Err(())
    48        } else {
    49            client::Symbol::normalize_and_validate_ident(string)
    50        }
    51        .unwrap_or_else(|_| panic!("`{:?}` is not a valid identifier", string))
    52    }
    53
    54    /// Run a callback with the symbol's string value.
    55    pub(crate) fn with<R>(self, f: impl FnOnce(&str) -> R) -> R {
    56        INTERNER.with_borrow(|i| f(i.get(self)))
    57    }
    58
    59    /// Clear out the thread-local symbol interner, making all previously
    60    /// created symbols invalid such that `with` will panic when called on them.
    61    pub(crate) fn invalidate_all() {
    62        INTERNER.with_borrow_mut(|i| i.clear());
    63    }
    64
    65    /// Checks if the ident is a valid ASCII identifier.
    66    ///
    67    /// This is a short-circuit which is cheap to implement within the
    68    /// proc-macro client to avoid RPC when creating simple idents, but may
    69    /// return `false` for a valid identifier if it contains non-ASCII
    70    /// characters.
    71    fn is_valid_ascii_ident(bytes: &[u8]) -> bool {
    72        matches!(bytes.first(), Some(b'_' | b'a'..=b'z' | b'A'..=b'Z'))
    73            && bytes[1..]
    74                .iter()
    75                .all(|b| matches!(b, b'_' | b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9'))
    76    }
    77
    78    // Mimics the behavior of `Symbol::can_be_raw` from `rustc_span`
    79    fn can_be_raw(string: &str) -> bool {
    80        !matches!(string, "_" | "super" | "self" | "Self" | "crate" | "$crate")
    81    }
    82}
    83
    84impl fmt::Debug for Symbol {
    85    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    86        self.with(|s| fmt::Debug::fmt(s, f))
    87    }
    88}
    89
    90impl fmt::Display for Symbol {
    91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    92        self.with(|s| fmt::Display::fmt(s, f))
    93    }
    94}
    95
    96impl<S> Encode<S> for Symbol {
    97    fn encode(self, w: &mut Writer, s: &mut S) {
    98        self.with(|sym| sym.encode(w, s))
    99    }
    100}
    101
    102impl<S: server::Server> Decode<'_, '_, server::HandleStore<server::MarkedTypes<S>>>
    103    for Marked<S::Symbol, Symbol>
    104{
    105    fn decode(r: &mut Reader<'_>, s: &mut server::HandleStore<server::MarkedTypes<S>>) -> Self {
    106        Mark::mark(S::intern_symbol(<&str>::decode(r, s)))
    107    }
    108}
    109
    110impl<S: server::Server> Encode<server::HandleStore<server::MarkedTypes<S>>>
    111    for Marked<S::Symbol, Symbol>
    112{
    113    fn encode(self, w: &mut Writer, s: &mut server::HandleStore<server::MarkedTypes<S>>) {
    114        S::with_symbol_string(&self.unmark(), |sym| sym.encode(w, s))
    115    }
    116}
    117
    118impl<S> Decode<'_, '_, S> for Symbol {
    119    fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
    120        Symbol::new(<&str>::decode(r, s))
    121    }
    122}
    123
    124thread_local! {
    125    static INTERNER: RefCell<Interner> = RefCell::new(Interner {
    126        arena: arena::Arena::new(),
    127        names: fxhash::FxHashMap::default(),
    128        strings: Vec::new(),
    129        // Start with a base of 1 to make sure that `NonZero<u32>` works.
    130        sym_base: NonZero::new(1).unwrap(),
    131    });
    132}
    133
    134/// Basic interner for a `Symbol`, inspired by the one in `rustc_span`.
    135struct Interner {
    136    arena: arena::Arena,
    137    // SAFETY: These `'static` lifetimes are actually references to data owned
    138    // by the Arena. This is safe, as we never return them as static references
    139    // from `Interner`.
    140    names: fxhash::FxHashMap<&'static str, Symbol>,
    141    strings: Vec<&'static str>,
    142    // The offset to apply to symbol names stored in the interner. This is used
    143    // to ensure that symbol names are not re-used after the interner is
    144    // cleared.
    145    sym_base: NonZero<u32>,
    146}
    147
    148impl Interner {
    149    fn intern(&mut self, string: &str) -> Symbol {
    150        if let Some(&name) = self.names.get(string) {
    151            return name;
    152        }
    153
    154        let name = Symbol(
    155            self.sym_base
    156                .checked_add(self.strings.len() as u32)
    157                .expect("`proc_macro` symbol name overflow"),
    158        );
    159
    160        let string: &str = self.arena.alloc_str(string);
    161
    162        // SAFETY: we can extend the arena allocation to `'static` because we
    163        // only access these while the arena is still alive.
    164        let string: &'static str = unsafe { &*(string as *const str) };
    165        self.strings.push(string);
    166        self.names.insert(string, name);
    167        name
    168    }
    169
    170    /// Reads a symbol's value from the store while it is held.
    171    fn get(&self, symbol: Symbol) -> &str {
    172        // NOTE: Subtract out the offset which was added to make the symbol
    173        // nonzero and prevent symbol name re-use.
    174        let name = symbol
    175            .0
    176            .get()
    177            .checked_sub(self.sym_base.get())
    178            .expect("use-after-free of `proc_macro` symbol");
    179        self.strings[name as usize]
    180    }
    181
    182    /// Clear all symbols from the store, invalidating them such that `get` will
    183    /// panic if they are accessed in the future.
    184    fn clear(&mut self) {
    185        // NOTE: Be careful not to panic here, as we may be called on the client
    186        // when a `catch_unwind` isn't installed.
    187        self.sym_base = self.sym_base.saturating_add(self.strings.len() as u32);
    188        self.names.clear();
    189        self.strings.clear();
    190
    191        // SAFETY: This is cleared after the names and strings tables are
    192        // cleared out, so no references into the arena should remain.
    193        self.arena = arena::Arena::new();
    194    }
    195}