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