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