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 Span {
115 pub(crate) fn def_site() -> Span {
116 Bridge::with(|bridge: &mut Bridge<'_>| bridge.globals.def_site)
117 }
118
119 pub(crate) fn call_site() -> Span {
120 Bridge::with(|bridge: &mut Bridge<'_>| bridge.globals.call_site)
121 }
122
123 pub(crate) fn mixed_site() -> Span {
124 Bridge::with(|bridge: &mut Bridge<'_>| bridge.globals.mixed_site)
125 }
126}
127
128impl fmt::Debug for Span {
129 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130 f.write_str(&self.debug())
131 }
132}
133
134pub(crate) use super::symbol::Symbol;
135
136macro_rules! define_client_side {
137 ($($name:ident {
138 $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
139 }),* $(,)?) => {
140 $(impl $name {
141 $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)? {
142 Bridge::with(|bridge| {
143 let mut buf = bridge.cached_buffer.take();
144
145 buf.clear();
146 api_tags::Method::$name(api_tags::$name::$method).encode(&mut buf, &mut ());
147 reverse_encode!(buf; $($arg),*);
148
149 buf = bridge.dispatch.call(buf);
150
151 let r = Result::<_, PanicMessage>::decode(&mut &buf[..], &mut ());
152
153 bridge.cached_buffer = buf;
154
155 r.unwrap_or_else(|e| panic::resume_unwind(e.into()))
156 })
157 })*
158 })*
159 }
160}
161with_api!(self, self, define_client_side);
162
163struct Bridge<'a> {
164 /// Reusable buffer (only `clear`-ed, never shrunk), primarily
165 /// used for making requests.
166 cached_buffer: Buffer,
167
168 /// Server-side function that the client uses to make requests.
169 dispatch: closure::Closure<'a, Buffer, Buffer>,
170
171 /// Provided globals for this macro expansion.
172 globals: ExpnGlobals<Span>,
173}
174
175impl<'a> !Send for Bridge<'a> {}
176impl<'a> !Sync for Bridge<'a> {}
177
178#[allow(unsafe_code)]
179mod state {
180 use std::cell::{Cell, RefCell};
181 use std::ptr;
182
183 use super::Bridge;
184
185 thread_local! {
186 static BRIDGE_STATE: Cell<*const ()> = const { Cell::new(ptr::null()) };
187 }
188
189 pub(super) fn set<'bridge, R>(state: &RefCell<Bridge<'bridge>>, f: impl FnOnce() -> R) -> R {
190 struct RestoreOnDrop(*const ());
191 impl Drop for RestoreOnDrop {
192 fn drop(&mut self) {
193 BRIDGE_STATE.set(self.0);
194 }
195 }
196
197 let inner = ptr::from_ref(state).cast();
198 let outer = BRIDGE_STATE.replace(inner);
199 let _restore = RestoreOnDrop(outer);
200
201 f()
202 }
203
204 pub(super) fn with<R>(
205 f: impl for<'bridge> FnOnce(Option<&RefCell<Bridge<'bridge>>>) -> R,
206 ) -> R {
207 let state = BRIDGE_STATE.get();
208 // SAFETY: the only place where the pointer is set is in `set`. It puts
209 // back the previous value after the inner call has returned, so we know
210 // that as long as the pointer is not null, it came from a reference to
211 // a `RefCell<Bridge>` that outlasts the call to this function. Since `f`
212 // works the same for any lifetime of the bridge, including the actual
213 // one, we can lie here and say that the lifetime is `'static` without
214 // anyone noticing.
215 let bridge = unsafe { state.cast::<RefCell<Bridge<'static>>>().as_ref() };
216 f(bridge)
217 }
218}
219
220impl Bridge<'_> {
221 fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
222 state::with(|state: Option<&RefCell>>| {
223 let bridge: &RefCell> = state.expect(msg:"procedural macro API is used outside of a procedural macro");
224 let mut bridge: RefMut<'_, Bridge<'_>> = bridge
225 .try_borrow_mut()
226 .expect(msg:"procedural macro API is used while it's already in use");
227 f(&mut bridge)
228 })
229 }
230}
231
232pub(crate) fn is_available() -> bool {
233 state::with(|s: Option<&RefCell>>| s.is_some())
234}
235
236/// A client-side RPC entry-point, which may be using a different `proc_macro`
237/// from the one used by the server, but can be invoked compatibly.
238///
239/// Note that the (phantom) `I` ("input") and `O` ("output") type parameters
240/// decorate the `Client<I, O>` with the RPC "interface" of the entry-point, but
241/// do not themselves participate in ABI, at all, only facilitate type-checking.
242///
243/// E.g. `Client<TokenStream, TokenStream>` is the common proc macro interface,
244/// used for `#[proc_macro] fn foo(input: TokenStream) -> TokenStream`,
245/// indicating that the RPC input and output will be serialized token streams,
246/// and forcing the use of APIs that take/return `S::TokenStream`, server-side.
247#[repr(C)]
248pub struct Client<I, O> {
249 pub(super) handle_counters: &'static HandleCounters,
250
251 pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
252
253 pub(super) _marker: PhantomData<fn(I) -> O>,
254}
255
256impl<I, O> Copy for Client<I, O> {}
257impl<I, O> Clone for Client<I, O> {
258 fn clone(&self) -> Self {
259 *self
260 }
261}
262
263fn maybe_install_panic_hook(force_show_panics: bool) {
264 // Hide the default panic output within `proc_macro` expansions.
265 // NB. the server can't do this because it may use a different std.
266 static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
267 HIDE_PANICS_DURING_EXPANSION.call_once(|| {
268 let prev: Box) + Send + Sync + 'static> = panic::take_hook();
269 panic::set_hook(Box::new(move |info: &PanicHookInfo<'_>| {
270 // We normally report panics by catching unwinds and passing the payload from the
271 // unwind back to the compiler, but if the panic doesn't unwind we'll abort before
272 // the compiler has a chance to print an error. So we special-case PanicInfo where
273 // can_unwind is false.
274 if force_show_panics || !is_available() || !info.can_unwind() {
275 prev(info)
276 }
277 }));
278 });
279}
280
281/// Client-side helper for handling client panics, entering the bridge,
282/// deserializing input and serializing output.
283// FIXME(eddyb) maybe replace `Bridge::enter` with this?
284fn run_client<A: for<'a, 's> dynDecodeMut<'a, 's, ()>, R: Encode<()>>(
285 config: BridgeConfig<'_>,
286 f: impl FnOnce(A) -> R,
287) -> Buffer {
288 let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
289
290 panic::catch_unwind(panic::AssertUnwindSafe(|| {
291 maybe_install_panic_hook(force_show_panics);
292
293 // Make sure the symbol store is empty before decoding inputs.
294 Symbol::invalidate_all();
295
296 let reader = &mut &buf[..];
297 let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());
298
299 // Put the buffer we used for input back in the `Bridge` for requests.
300 let state = RefCell::new(Bridge { cached_buffer: buf.take(), dispatch, globals });
301
302 let output = state::set(&state, || f(input));
303
304 // Take the `cached_buffer` back out, for the output value.
305 buf = RefCell::into_inner(state).cached_buffer;
306
307 // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
308 // from encoding a panic (`Err(e: PanicMessage)`) to avoid
309 // having handles outside the `bridge.enter(|| ...)` scope, and
310 // to catch panics that could happen while encoding the success.
311 //
312 // Note that panics should be impossible beyond this point, but
313 // this is defensively trying to avoid any accidental panicking
314 // reaching the `extern "C"` (which should `abort` but might not
315 // at the moment, so this is also potentially preventing UB).
316 buf.clear();
317 Ok::<_, ()>(output).encode(&mut buf, &mut ());
318 }))
319 .map_err(PanicMessage::from)
320 .unwrap_or_else(|e| {
321 buf.clear();
322 Err::<(), _>(e).encode(&mut buf, &mut ());
323 });
324
325 // Now that a response has been serialized, invalidate all symbols
326 // registered with the interner.
327 Symbol::invalidate_all();
328 buf
329}
330
331impl Client<crate::TokenStream, crate::TokenStream> {
332 pub const fn expand1(f: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy) -> Self {
333 Client {
334 handle_counters: &COUNTERS,
335 run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge: BridgeConfig<'_>| {
336 run_client(config:bridge, |input: TokenStream| f(crate::TokenStream(Some(input))).0)
337 }),
338 _marker: PhantomData,
339 }
340 }
341}
342
343impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
344 pub const fn expand2(
345 f: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
346 ) -> Self {
347 Client {
348 handle_counters: &COUNTERS,
349 run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge: BridgeConfig<'_>| {
350 run_client(config:bridge, |(input: TokenStream, input2: TokenStream)| {
351 f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2))).0
352 })
353 }),
354 _marker: PhantomData,
355 }
356 }
357}
358
359#[repr(C)]
360#[derive(Copy, Clone)]
361pub enum ProcMacro {
362 CustomDerive {
363 trait_name: &'static str,
364 attributes: &'static [&'static str],
365 client: Client<crate::TokenStream, crate::TokenStream>,
366 },
367
368 Attr {
369 name: &'static str,
370 client: Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream>,
371 },
372
373 Bang {
374 name: &'static str,
375 client: Client<crate::TokenStream, crate::TokenStream>,
376 },
377}
378
379impl ProcMacro {
380 pub fn name(&self) -> &'static str {
381 match self {
382 ProcMacro::CustomDerive { trait_name, .. } => trait_name,
383 ProcMacro::Attr { name, .. } => name,
384 ProcMacro::Bang { name, .. } => name,
385 }
386 }
387
388 pub const fn custom_derive(
389 trait_name: &'static str,
390 attributes: &'static [&'static str],
391 expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
392 ) -> Self {
393 ProcMacro::CustomDerive { trait_name, attributes, client: Client::expand1(expand) }
394 }
395
396 pub const fn attr(
397 name: &'static str,
398 expand: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
399 ) -> Self {
400 ProcMacro::Attr { name, client: Client::expand2(expand) }
401 }
402
403 pub const fn bang(
404 name: &'static str,
405 expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
406 ) -> Self {
407 ProcMacro::Bang { name, client: Client::expand1(expand) }
408 }
409}
410