1//! Client-side types.
2
3use std::cell::RefCell;
4use std::marker::PhantomData;
5use std::sync::atomic::AtomicU32;
6
7use super::*;
8
9#[repr(C)]
10pub(super) struct HandleCounters {
11 pub(super) token_stream: AtomicU32,
12 pub(super) span: AtomicU32,
13}
14
15static COUNTERS: HandleCounters =
16 HandleCounters { token_stream: AtomicU32::new(1), span: AtomicU32::new(1) };
17
18pub(crate) struct TokenStream {
19 handle: handle::Handle,
20}
21
22impl !Send for TokenStream {}
23impl !Sync for TokenStream {}
24
25// Forward `Drop::drop` to the inherent `drop` method.
26impl Drop for TokenStream {
27 fn drop(&mut self) {
28 Methods::ts_drop(stream:TokenStream { handle: self.handle });
29 }
30}
31
32impl<S> Encode<S> for TokenStream {
33 fn encode(self, w: &mut Buffer, s: &mut S) {
34 mem::ManuallyDrop::new(self).handle.encode(w, s);
35 }
36}
37
38impl<S> Encode<S> for &TokenStream {
39 fn encode(self, w: &mut Buffer, s: &mut S) {
40 self.handle.encode(w, s);
41 }
42}
43
44impl<S> Decode<'_, '_, S> for TokenStream {
45 fn decode(r: &mut &[u8], s: &mut S) -> Self {
46 TokenStream { handle: handle::Handle::decode(r, s) }
47 }
48}
49
50#[derive(Copy, Clone, PartialEq, Eq, Hash)]
51pub(crate) struct Span {
52 handle: handle::Handle,
53}
54
55impl !Send for Span {}
56impl !Sync for Span {}
57
58impl<S> Encode<S> for Span {
59 fn encode(self, w: &mut Buffer, s: &mut S) {
60 self.handle.encode(w, s);
61 }
62}
63
64impl<S> Decode<'_, '_, S> for Span {
65 fn decode(r: &mut &[u8], s: &mut S) -> Self {
66 Span { handle: handle::Handle::decode(r, s) }
67 }
68}
69
70impl Clone for TokenStream {
71 fn clone(&self) -> Self {
72 Methods::ts_clone(self)
73 }
74}
75
76impl Span {
77 pub(crate) fn def_site() -> Span {
78 Bridge::with(|bridge| bridge.globals.def_site)
79 }
80
81 pub(crate) fn call_site() -> Span {
82 Bridge::with(|bridge| bridge.globals.call_site)
83 }
84
85 pub(crate) fn mixed_site() -> Span {
86 Bridge::with(|bridge| bridge.globals.mixed_site)
87 }
88}
89
90impl fmt::Debug for Span {
91 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92 f.write_str(&Methods::span_debug(*self))
93 }
94}
95
96pub(crate) use super::Methods;
97pub(crate) use super::symbol::Symbol;
98
99macro_rules! define_client_side {
100 (
101 $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)*;)*
102 ) => {
103 impl Methods {
104 $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)? {
105 Bridge::with(|bridge| {
106 let mut buf = bridge.cached_buffer.take();
107
108 buf.clear();
109 ApiTags::$method.encode(&mut buf, &mut ());
110 $($arg.encode(&mut buf, &mut ());)*
111
112 buf = bridge.dispatch.call(buf);
113
114 let r = Result::<_, PanicMessage>::decode(&mut &buf[..], &mut ());
115
116 bridge.cached_buffer = buf;
117
118 r.unwrap_or_else(|e| panic::resume_unwind(e.into()))
119 })
120 })*
121 }
122 }
123}
124with_api!(self, define_client_side);
125
126struct Bridge<'a> {
127 /// Reusable buffer (only `clear`-ed, never shrunk), primarily
128 /// used for making requests.
129 cached_buffer: Buffer,
130
131 /// Server-side function that the client uses to make requests.
132 dispatch: closure::Closure<'a, Buffer, Buffer>,
133
134 /// Provided globals for this macro expansion.
135 globals: ExpnGlobals<Span>,
136}
137
138impl<'a> !Send for Bridge<'a> {}
139impl<'a> !Sync for Bridge<'a> {}
140
141#[allow(unsafe_code)]
142mod state {
143 use std::cell::{Cell, RefCell};
144 use std::ptr;
145
146 use super::Bridge;
147
148 thread_local! {
149 static BRIDGE_STATE: Cell<*const ()> = const { Cell::new(ptr::null()) };
150 }
151
152 pub(super) fn set<'bridge, R>(state: &RefCell<Bridge<'bridge>>, f: impl FnOnce() -> R) -> R {
153 struct RestoreOnDrop(*const ());
154 impl Drop for RestoreOnDrop {
155 fn drop(&mut self) {
156 BRIDGE_STATE.set(self.0);
157 }
158 }
159
160 let inner = ptr::from_ref(state).cast();
161 let outer = BRIDGE_STATE.replace(inner);
162 let _restore = RestoreOnDrop(outer);
163
164 f()
165 }
166
167 pub(super) fn with<R>(
168 f: impl for<'bridge> FnOnce(Option<&RefCell<Bridge<'bridge>>>) -> R,
169 ) -> R {
170 let state = BRIDGE_STATE.get();
171 // SAFETY: the only place where the pointer is set is in `set`. It puts
172 // back the previous value after the inner call has returned, so we know
173 // that as long as the pointer is not null, it came from a reference to
174 // a `RefCell<Bridge>` that outlasts the call to this function. Since `f`
175 // works the same for any lifetime of the bridge, including the actual
176 // one, we can lie here and say that the lifetime is `'static` without
177 // anyone noticing.
178 let bridge = unsafe { state.cast::<RefCell<Bridge<'static>>>().as_ref() };
179 f(bridge)
180 }
181}
182
183impl Bridge<'_> {
184 fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
185 state::with(|state| {
186 let bridge = state.expect("procedural macro API is used outside of a procedural macro");
187 let mut bridge = bridge
188 .try_borrow_mut()
189 .expect("procedural macro API is used while it's already in use");
190 f(&mut bridge)
191 })
192 }
193}
194
195pub(crate) fn is_available() -> bool {
196 state::with(|s| s.is_some())
197}
198
199/// A client-side RPC entry-point, which may be using a different `proc_macro`
200/// from the one used by the server, but can be invoked compatibly.
201///
202/// Note that the (phantom) `I` ("input") and `O` ("output") type parameters
203/// decorate the `Client<I, O>` with the RPC "interface" of the entry-point, but
204/// do not themselves participate in ABI, at all, only facilitate type-checking.
205///
206/// E.g. `Client<TokenStream, TokenStream>` is the common proc macro interface,
207/// used for `#[proc_macro] fn foo(input: TokenStream) -> TokenStream`,
208/// indicating that the RPC input and output will be serialized token streams,
209/// and forcing the use of APIs that take/return `S::TokenStream`, server-side.
210#[repr(C)]
211pub struct Client<I, O> {
212 pub(super) handle_counters: &'static HandleCounters,
213
214 pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
215
216 pub(super) _marker: PhantomData<fn(I) -> O>,
217}
218
219impl<I, O> Copy for Client<I, O> {}
220impl<I, O> Clone for Client<I, O> {
221 fn clone(&self) -> Self {
222 *self
223 }
224}
225
226fn maybe_install_panic_hook(force_show_panics: bool) {
227 // Hide the default panic output within `proc_macro` expansions.
228 // NB. the server can't do this because it may use a different std.
229 static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
230 HIDE_PANICS_DURING_EXPANSION.call_once(|| {
231 let prev = panic::take_hook();
232 panic::set_hook(Box::new(move |info| {
233 // We normally report panics by catching unwinds and passing the payload from the
234 // unwind back to the compiler, but if the panic doesn't unwind we'll abort before
235 // the compiler has a chance to print an error. So we special-case PanicInfo where
236 // can_unwind is false.
237 if force_show_panics || !is_available() || !info.can_unwind() {
238 prev(info)
239 }
240 }));
241 });
242}
243
244/// Client-side helper for handling client panics, entering the bridge,
245/// deserializing input and serializing output.
246// FIXME(eddyb) maybe replace `Bridge::enter` with this?
247fn run_client<A: for<'a, 's> Decode<'a, 's, ()>, R: Encode<()>>(
248 config: BridgeConfig<'_>,
249 f: impl FnOnce(A) -> R,
250) -> Buffer {
251 let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
252
253 panic::catch_unwind(panic::AssertUnwindSafe(|| {
254 maybe_install_panic_hook(force_show_panics);
255
256 // Make sure the symbol store is empty before decoding inputs.
257 Symbol::invalidate_all();
258
259 let reader = &mut &buf[..];
260 let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());
261
262 // Put the buffer we used for input back in the `Bridge` for requests.
263 let state = RefCell::new(Bridge { cached_buffer: buf.take(), dispatch, globals });
264
265 let output = state::set(&state, || f(input));
266
267 // Take the `cached_buffer` back out, for the output value.
268 buf = RefCell::into_inner(state).cached_buffer;
269
270 // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
271 // from encoding a panic (`Err(e: PanicMessage)`) to avoid
272 // having handles outside the `bridge.enter(|| ...)` scope, and
273 // to catch panics that could happen while encoding the success.
274 //
275 // Note that panics should be impossible beyond this point, but
276 // this is defensively trying to avoid any accidental panicking
277 // reaching the `extern "C"` (which should `abort` but might not
278 // at the moment, so this is also potentially preventing UB).
279 buf.clear();
280 Ok::<_, ()>(output).encode(&mut buf, &mut ());
281 }))
282 .map_err(PanicMessage::from)
283 .unwrap_or_else(|e| {
284 buf.clear();
285 Err::<(), _>(e).encode(&mut buf, &mut ());
286 });
287
288 // Now that a response has been serialized, invalidate all symbols
289 // registered with the interner.
290 Symbol::invalidate_all();
291 buf
292}
293
294impl Client<crate::TokenStream, crate::TokenStream> {
295 pub const fn expand1(f: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy) -> Self {
296 Client {
297 handle_counters: &COUNTERS,
298 run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge: BridgeConfig<'_>| {
299 run_client(config:bridge, |input| f(crate::TokenStream(Some(input))).0)
300 }),
301 _marker: PhantomData,
302 }
303 }
304}
305
306impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
307 pub const fn expand2(
308 f: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
309 ) -> Self {
310 Client {
311 handle_counters: &COUNTERS,
312 run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge: BridgeConfig<'_>| {
313 run_client(config:bridge, |(input, input2)| {
314 f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2))).0
315 })
316 }),
317 _marker: PhantomData,
318 }
319 }
320}
321
322#[repr(C)]
323#[derive(Copy, Clone)]
324pub enum ProcMacro {
325 CustomDerive {
326 trait_name: &'static str,
327 attributes: &'static [&'static str],
328 client: Client<crate::TokenStream, crate::TokenStream>,
329 },
330
331 Attr {
332 name: &'static str,
333 client: Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream>,
334 },
335
336 Bang {
337 name: &'static str,
338 client: Client<crate::TokenStream, crate::TokenStream>,
339 },
340}
341
342impl ProcMacro {
343 pub fn name(&self) -> &'static str {
344 match self {
345 ProcMacro::CustomDerive { trait_name, .. } => trait_name,
346 ProcMacro::Attr { name, .. } => name,
347 ProcMacro::Bang { name, .. } => name,
348 }
349 }
350
351 pub const fn custom_derive(
352 trait_name: &'static str,
353 attributes: &'static [&'static str],
354 expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
355 ) -> Self {
356 ProcMacro::CustomDerive { trait_name, attributes, client: Client::expand1(expand) }
357 }
358
359 pub const fn attr(
360 name: &'static str,
361 expand: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
362 ) -> Self {
363 ProcMacro::Attr { name, client: Client::expand2(expand) }
364 }
365
366 pub const fn bang(
367 name: &'static str,
368 expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
369 ) -> Self {
370 ProcMacro::Bang { name, client: Client::expand1(expand) }
371 }
372}
373