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