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::fmt;
14use std::num::NonZero;
15
16use crate::bridge::{Buffer, Decode, Encode, Mark, arena, client, fxhash, server};
17
18/// Handle for a symbol string stored within the Interner.
19#[derive(Copy, Clone, PartialEq, Eq, Hash)]
20pub struct Symbol(NonZero<u32>);
21
22impl !Send for Symbol {}
23impl !Sync for Symbol {}
24
25impl 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    /// Creates 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()) || string == "$crate" {
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::Methods::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    /// Checks 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 behavior of `Symbol::can_be_raw` from `rustc_span`
80    fn can_be_raw(string: &str) -> bool {
81        !matches!(string, "_" | "super" | "self" | "Self" | "crate" | "$crate")
82    }
83}
84
85impl fmt::Debug for Symbol {
86    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
87        self.with(|s| fmt::Debug::fmt(s, f))
88    }
89}
90
91impl fmt::Display for Symbol {
92    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
93        self.with(|s| fmt::Display::fmt(s, f))
94    }
95}
96
97impl<S> Encode<S> for Symbol {
98    fn encode(self, w: &mut Buffer, s: &mut S) {
99        self.with(|sym| sym.encode(w, s))
100    }
101}
102
103impl<S: server::Server> Decode<'_, '_, server::HandleStore<S>> for server::MarkedSymbol<S> {
104    fn decode(r: &mut &[u8], s: &mut server::HandleStore<S>) -> Self {
105        Mark::mark(S::intern_symbol(<&str>::decode(r, s)))
106    }
107}
108
109impl<S: server::Server> Encode<server::HandleStore<S>> for server::MarkedSymbol<S> {
110    fn encode(self, w: &mut Buffer, s: &mut server::HandleStore<S>) {
111        S::with_symbol_string(&self.unmark(), |sym| sym.encode(w, s))
112    }
113}
114
115impl<S> Decode<'_, '_, S> for Symbol {
116    fn decode(r: &mut &[u8], s: &mut S) -> Self {
117        Symbol::new(<&str>::decode(r, s))
118    }
119}
120
121thread_local! {
122    static INTERNER: RefCell<Interner> = RefCell::new(Interner {
123        arena: arena::Arena::new(),
124        names: fxhash::FxHashMap::default(),
125        strings: Vec::new(),
126        // Start with a base of 1 to make sure that `NonZero<u32>` works.
127        sym_base: NonZero::new(1).unwrap(),
128    });
129}
130
131/// Basic interner for a `Symbol`, inspired by the one in `rustc_span`.
132struct Interner {
133    arena: arena::Arena,
134    // SAFETY: These `'static` lifetimes are actually references to data owned
135    // by the Arena. This is safe, as we never return them as static references
136    // from `Interner`.
137    names: fxhash::FxHashMap<&'static str, Symbol>,
138    strings: Vec<&'static str>,
139    // The offset to apply to symbol names stored in the interner. This is used
140    // to ensure that symbol names are not re-used after the interner is
141    // cleared.
142    sym_base: NonZero<u32>,
143}
144
145impl Interner {
146    fn intern(&mut self, string: &str) -> Symbol {
147        if let Some(&name) = self.names.get(string) {
148            return name;
149        }
150
151        let name = Symbol(
152            self.sym_base
153                .checked_add(self.strings.len() as u32)
154                .expect("`proc_macro` symbol name overflow"),
155        );
156
157        let string: &str = self.arena.alloc_str(string);
158
159        // SAFETY: we can extend the arena allocation to `'static` because we
160        // only access these while the arena is still alive.
161        let string: &'static str = unsafe { &*(string as *const str) };
162        self.strings.push(string);
163        self.names.insert(string, name);
164        name
165    }
166
167    /// Reads a symbol's value from the store while it is held.
168    fn get(&self, symbol: Symbol) -> &str {
169        // NOTE: Subtract out the offset which was added to make the symbol
170        // nonzero and prevent symbol name re-use.
171        let name = symbol
172            .0
173            .get()
174            .checked_sub(self.sym_base.get())
175            .expect("use-after-free of `proc_macro` symbol");
176        self.strings[name as usize]
177    }
178
179    /// Clear all symbols from the store, invalidating them such that `get` will
180    /// panic if they are accessed in the future.
181    fn clear(&mut self) {
182        // NOTE: Be careful not to panic here, as we may be called on the client
183        // when a `catch_unwind` isn't installed.
184        self.sym_base = self.sym_base.saturating_add(self.strings.len() as u32);
185        self.names.clear();
186        self.strings.clear();
187
188        // SAFETY: This is cleared after the names and strings tables are
189        // cleared out, so no references into the arena should remain.
190        self.arena = arena::Arena::new();
191    }
192}
proc_macro/bridge/symbol.rs

proc_macro/bridge/
symbol.rs
