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