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::NonZero; |
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(NonZero<u32>); |
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: &str| fmt::Debug::fmt(self:s, f)) |
91 | } |
92 | } |
93 | |
94 | impl ToString for Symbol { |
95 | fn to_string(&self) -> String { |
96 | self.with(|s: &str| 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: &str| fmt::Display::fmt(self: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: &str| sym.encode(w, s)) |
109 | } |
110 | } |
111 | |
112 | impl<S: server::Server> DecodeMut<'_, '_, server::HandleStore<server::MarkedTypes<S>>> |
113 | for Marked<S::Symbol, Symbol> |
114 | { |
115 | fn decode(r: &mut Reader<'_>, s: &mut server::HandleStore<server::MarkedTypes<S>>) -> Self { |
116 | Mark::mark(S::intern_symbol(<&str>::decode(r, s))) |
117 | } |
118 | } |
119 | |
120 | impl<S: server::Server> Encode<server::HandleStore<server::MarkedTypes<S>>> |
121 | for Marked<S::Symbol, Symbol> |
122 | { |
123 | fn encode(self, w: &mut Writer, s: &mut server::HandleStore<server::MarkedTypes<S>>) { |
124 | S::with_symbol_string(&self.unmark(), |sym: &str| 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 `NonZero<u32>` works. |
140 | sym_base: NonZero::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: NonZero<u32>, |
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 | } |
206 | |