1//! Client-side types.
2
3use super::*;
4
5use std::marker::PhantomData;
6
7macro_rules! define_handles {
8 (
9 'owned: $($oty:ident,)*
10 'interned: $($ity:ident,)*
11 ) => {
12 #[repr(C)]
13 #[allow(non_snake_case)]
14 pub struct HandleCounters {
15 $($oty: AtomicUsize,)*
16 $($ity: AtomicUsize,)*
17 }
18
19 impl HandleCounters {
20 // FIXME(eddyb) use a reference to the `static COUNTERS`, instead of
21 // a wrapper `fn` pointer, once `const fn` can reference `static`s.
22 extern "C" fn get() -> &'static Self {
23 static COUNTERS: HandleCounters = HandleCounters {
24 $($oty: AtomicUsize::new(1),)*
25 $($ity: AtomicUsize::new(1),)*
26 };
27 &COUNTERS
28 }
29 }
30
31 // FIXME(eddyb) generate the definition of `HandleStore` in `server.rs`.
32 #[allow(non_snake_case)]
33 pub(super) struct HandleStore<S: server::Types> {
34 $($oty: handle::OwnedStore<S::$oty>,)*
35 $($ity: handle::InternedStore<S::$ity>,)*
36 }
37
38 impl<S: server::Types> HandleStore<S> {
39 pub(super) fn new(handle_counters: &'static HandleCounters) -> Self {
40 HandleStore {
41 $($oty: handle::OwnedStore::new(&handle_counters.$oty),)*
42 $($ity: handle::InternedStore::new(&handle_counters.$ity),)*
43 }
44 }
45 }
46
47 $(
48 pub(crate) struct $oty {
49 handle: handle::Handle,
50 // Prevent Send and Sync impls. `!Send`/`!Sync` is the usual
51 // way of doing this, but that requires unstable features.
52 // rust-analyzer uses this code and avoids unstable features.
53 _marker: PhantomData<*mut ()>,
54 }
55
56 // Forward `Drop::drop` to the inherent `drop` method.
57 impl Drop for $oty {
58 fn drop(&mut self) {
59 $oty {
60 handle: self.handle,
61 _marker: PhantomData,
62 }.drop();
63 }
64 }
65
66 impl<S> Encode<S> for $oty {
67 fn encode(self, w: &mut Writer, s: &mut S) {
68 let handle = self.handle;
69 mem::forget(self);
70 handle.encode(w, s);
71 }
72 }
73
74 impl<S: server::Types> DecodeMut<'_, '_, HandleStore<server::MarkedTypes<S>>>
75 for Marked<S::$oty, $oty>
76 {
77 fn decode(r: &mut Reader<'_>, s: &mut HandleStore<server::MarkedTypes<S>>) -> Self {
78 s.$oty.take(handle::Handle::decode(r, &mut ()))
79 }
80 }
81
82 impl<S> Encode<S> for &$oty {
83 fn encode(self, w: &mut Writer, s: &mut S) {
84 self.handle.encode(w, s);
85 }
86 }
87
88 impl<'s, S: server::Types> Decode<'_, 's, HandleStore<server::MarkedTypes<S>>>
89 for &'s Marked<S::$oty, $oty>
90 {
91 fn decode(r: &mut Reader<'_>, s: &'s HandleStore<server::MarkedTypes<S>>) -> Self {
92 &s.$oty[handle::Handle::decode(r, &mut ())]
93 }
94 }
95
96 impl<S> Encode<S> for &mut $oty {
97 fn encode(self, w: &mut Writer, s: &mut S) {
98 self.handle.encode(w, s);
99 }
100 }
101
102 impl<'s, S: server::Types> DecodeMut<'_, 's, HandleStore<server::MarkedTypes<S>>>
103 for &'s mut Marked<S::$oty, $oty>
104 {
105 fn decode(
106 r: &mut Reader<'_>,
107 s: &'s mut HandleStore<server::MarkedTypes<S>>
108 ) -> Self {
109 &mut s.$oty[handle::Handle::decode(r, &mut ())]
110 }
111 }
112
113 impl<S: server::Types> Encode<HandleStore<server::MarkedTypes<S>>>
114 for Marked<S::$oty, $oty>
115 {
116 fn encode(self, w: &mut Writer, s: &mut HandleStore<server::MarkedTypes<S>>) {
117 s.$oty.alloc(self).encode(w, s);
118 }
119 }
120
121 impl<S> DecodeMut<'_, '_, S> for $oty {
122 fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
123 $oty {
124 handle: handle::Handle::decode(r, s),
125 _marker: PhantomData,
126 }
127 }
128 }
129 )*
130
131 $(
132 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
133 pub(crate) struct $ity {
134 handle: handle::Handle,
135 // Prevent Send and Sync impls. `!Send`/`!Sync` is the usual
136 // way of doing this, but that requires unstable features.
137 // rust-analyzer uses this code and avoids unstable features.
138 _marker: PhantomData<*mut ()>,
139 }
140
141 impl<S> Encode<S> for $ity {
142 fn encode(self, w: &mut Writer, s: &mut S) {
143 self.handle.encode(w, s);
144 }
145 }
146
147 impl<S: server::Types> DecodeMut<'_, '_, HandleStore<server::MarkedTypes<S>>>
148 for Marked<S::$ity, $ity>
149 {
150 fn decode(r: &mut Reader<'_>, s: &mut HandleStore<server::MarkedTypes<S>>) -> Self {
151 s.$ity.copy(handle::Handle::decode(r, &mut ()))
152 }
153 }
154
155 impl<S: server::Types> Encode<HandleStore<server::MarkedTypes<S>>>
156 for Marked<S::$ity, $ity>
157 {
158 fn encode(self, w: &mut Writer, s: &mut HandleStore<server::MarkedTypes<S>>) {
159 s.$ity.alloc(self).encode(w, s);
160 }
161 }
162
163 impl<S> DecodeMut<'_, '_, S> for $ity {
164 fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
165 $ity {
166 handle: handle::Handle::decode(r, s),
167 _marker: PhantomData,
168 }
169 }
170 }
171 )*
172 }
173}
174define_handles! {
175 'owned:
176 FreeFunctions,
177 TokenStream,
178 SourceFile,
179
180 'interned:
181 Span,
182}
183
184// FIXME(eddyb) generate these impls by pattern-matching on the
185// names of methods - also could use the presence of `fn drop`
186// to distinguish between 'owned and 'interned, above.
187// Alternatively, special "modes" could be listed of types in with_api
188// instead of pattern matching on methods, here and in server decl.
189
190impl Clone for TokenStream {
191 fn clone(&self) -> Self {
192 self.clone()
193 }
194}
195
196impl Clone for SourceFile {
197 fn clone(&self) -> Self {
198 self.clone()
199 }
200}
201
202impl Span {
203 pub(crate) fn def_site() -> Span {
204 Bridge::with(|bridge: &mut Bridge<'_>| bridge.globals.def_site)
205 }
206
207 pub(crate) fn call_site() -> Span {
208 Bridge::with(|bridge: &mut Bridge<'_>| bridge.globals.call_site)
209 }
210
211 pub(crate) fn mixed_site() -> Span {
212 Bridge::with(|bridge: &mut Bridge<'_>| bridge.globals.mixed_site)
213 }
214}
215
216impl fmt::Debug for Span {
217 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
218 f.write_str(&self.debug())
219 }
220}
221
222pub(crate) use super::symbol::Symbol;
223
224macro_rules! define_client_side {
225 ($($name:ident {
226 $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
227 }),* $(,)?) => {
228 $(impl $name {
229 $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)? {
230 Bridge::with(|bridge| {
231 let mut buf = bridge.cached_buffer.take();
232
233 buf.clear();
234 api_tags::Method::$name(api_tags::$name::$method).encode(&mut buf, &mut ());
235 reverse_encode!(buf; $($arg),*);
236
237 buf = bridge.dispatch.call(buf);
238
239 let r = Result::<_, PanicMessage>::decode(&mut &buf[..], &mut ());
240
241 bridge.cached_buffer = buf;
242
243 r.unwrap_or_else(|e| panic::resume_unwind(e.into()))
244 })
245 })*
246 })*
247 }
248}
249with_api!(self, self, define_client_side);
250
251struct Bridge<'a> {
252 /// Reusable buffer (only `clear`-ed, never shrunk), primarily
253 /// used for making requests.
254 cached_buffer: Buffer,
255
256 /// Server-side function that the client uses to make requests.
257 dispatch: closure::Closure<'a, Buffer, Buffer>,
258
259 /// Provided globals for this macro expansion.
260 globals: ExpnGlobals<Span>,
261}
262
263impl<'a> !Send for Bridge<'a> {}
264impl<'a> !Sync for Bridge<'a> {}
265
266enum BridgeState<'a> {
267 /// No server is currently connected to this client.
268 NotConnected,
269
270 /// A server is connected and available for requests.
271 Connected(Bridge<'a>),
272
273 /// Access to the bridge is being exclusively acquired
274 /// (e.g., during `BridgeState::with`).
275 InUse,
276}
277
278enum BridgeStateL {}
279
280impl<'a> scoped_cell::ApplyL<'a> for BridgeStateL {
281 type Out = BridgeState<'a>;
282}
283
284thread_local! {
285 static BRIDGE_STATE: scoped_cell::ScopedCell<BridgeStateL> =
286 scoped_cell::ScopedCell::new(BridgeState::NotConnected);
287}
288
289impl BridgeState<'_> {
290 /// Take exclusive control of the thread-local
291 /// `BridgeState`, and pass it to `f`, mutably.
292 /// The state will be restored after `f` exits, even
293 /// by panic, including modifications made to it by `f`.
294 ///
295 /// N.B., while `f` is running, the thread-local state
296 /// is `BridgeState::InUse`.
297 fn with<R>(f: impl FnOnce(&mut BridgeState<'_>) -> R) -> R {
298 BRIDGE_STATE.with(|state: &ScopedCell| {
299 state.replace(replacement:BridgeState::InUse, |mut state: RefMutL<'_, '_, BridgeStateL>| {
300 // FIXME(#52812) pass `f` directly to `replace` when `RefMutL` is gone
301 f(&mut *state)
302 })
303 })
304 }
305}
306
307impl Bridge<'_> {
308 fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
309 BridgeState::with(|state: &mut BridgeState<'_>| match state {
310 BridgeState::NotConnected => {
311 panic!("procedural macro API is used outside of a procedural macro");
312 }
313 BridgeState::InUse => {
314 panic!("procedural macro API is used while it's already in use");
315 }
316 BridgeState::Connected(bridge: &mut Bridge<'_>) => f(bridge),
317 })
318 }
319}
320
321pub(crate) fn is_available() -> bool {
322 BridgeState::with(|state: &mut BridgeState<'_>| match state {
323 BridgeState::Connected(_) | BridgeState::InUse => true,
324 BridgeState::NotConnected => false,
325 })
326}
327
328/// A client-side RPC entry-point, which may be using a different `proc_macro`
329/// from the one used by the server, but can be invoked compatibly.
330///
331/// Note that the (phantom) `I` ("input") and `O` ("output") type parameters
332/// decorate the `Client<I, O>` with the RPC "interface" of the entry-point, but
333/// do not themselves participate in ABI, at all, only facilitate type-checking.
334///
335/// E.g. `Client<TokenStream, TokenStream>` is the common proc macro interface,
336/// used for `#[proc_macro] fn foo(input: TokenStream) -> TokenStream`,
337/// indicating that the RPC input and output will be serialized token streams,
338/// and forcing the use of APIs that take/return `S::TokenStream`, server-side.
339#[repr(C)]
340pub struct Client<I, O> {
341 // FIXME(eddyb) use a reference to the `static COUNTERS`, instead of
342 // a wrapper `fn` pointer, once `const fn` can reference `static`s.
343 pub(super) get_handle_counters: extern "C" fn() -> &'static HandleCounters,
344
345 pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
346
347 pub(super) _marker: PhantomData<fn(I) -> O>,
348}
349
350impl<I, O> Copy for Client<I, O> {}
351impl<I, O> Clone for Client<I, O> {
352 fn clone(&self) -> Self {
353 *self
354 }
355}
356
357fn maybe_install_panic_hook(force_show_panics: bool) {
358 // Hide the default panic output within `proc_macro` expansions.
359 // NB. the server can't do this because it may use a different std.
360 static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
361 HIDE_PANICS_DURING_EXPANSION.call_once(|| {
362 let prev: Box) + Sync + Send> = panic::take_hook();
363 panic::set_hook(Box::new(move |info: &PanicInfo<'_>| {
364 let show: bool = BridgeState::with(|state: &mut BridgeState<'_>| match state {
365 BridgeState::NotConnected => true,
366 BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
367 });
368 if show {
369 prev(info)
370 }
371 }));
372 });
373}
374
375/// Client-side helper for handling client panics, entering the bridge,
376/// deserializing input and serializing output.
377// FIXME(eddyb) maybe replace `Bridge::enter` with this?
378fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
379 config: BridgeConfig<'_>,
380 f: impl FnOnce(A) -> R,
381) -> Buffer {
382 let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
383
384 panic::catch_unwind(panic::AssertUnwindSafe(|| {
385 maybe_install_panic_hook(force_show_panics);
386
387 // Make sure the symbol store is empty before decoding inputs.
388 Symbol::invalidate_all();
389
390 let reader = &mut &buf[..];
391 let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());
392
393 // Put the buffer we used for input back in the `Bridge` for requests.
394 let new_state =
395 BridgeState::Connected(Bridge { cached_buffer: buf.take(), dispatch, globals });
396
397 BRIDGE_STATE.with(|state| {
398 state.set(new_state, || {
399 let output = f(input);
400
401 // Take the `cached_buffer` back out, for the output value.
402 buf = Bridge::with(|bridge| bridge.cached_buffer.take());
403
404 // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
405 // from encoding a panic (`Err(e: PanicMessage)`) to avoid
406 // having handles outside the `bridge.enter(|| ...)` scope, and
407 // to catch panics that could happen while encoding the success.
408 //
409 // Note that panics should be impossible beyond this point, but
410 // this is defensively trying to avoid any accidental panicking
411 // reaching the `extern "C"` (which should `abort` but might not
412 // at the moment, so this is also potentially preventing UB).
413 buf.clear();
414 Ok::<_, ()>(output).encode(&mut buf, &mut ());
415 })
416 })
417 }))
418 .map_err(PanicMessage::from)
419 .unwrap_or_else(|e| {
420 buf.clear();
421 Err::<(), _>(e).encode(&mut buf, &mut ());
422 });
423
424 // Now that a response has been serialized, invalidate all symbols
425 // registered with the interner.
426 Symbol::invalidate_all();
427 buf
428}
429
430impl Client<crate::TokenStream, crate::TokenStream> {
431 pub const fn expand1(f: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy) -> Self {
432 Client {
433 get_handle_counters: HandleCounters::get,
434 run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge: BridgeConfig<'_>| {
435 run_client(config:bridge, |input: TokenStream| f(crate::TokenStream(Some(input))).0)
436 }),
437 _marker: PhantomData,
438 }
439 }
440}
441
442impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
443 pub const fn expand2(
444 f: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
445 ) -> Self {
446 Client {
447 get_handle_counters: HandleCounters::get,
448 run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge: BridgeConfig<'_>| {
449 run_client(config:bridge, |(input: TokenStream, input2: TokenStream)| {
450 f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2))).0
451 })
452 }),
453 _marker: PhantomData,
454 }
455 }
456}
457
458#[repr(C)]
459#[derive(Copy, Clone)]
460pub enum ProcMacro {
461 CustomDerive {
462 trait_name: &'static str,
463 attributes: &'static [&'static str],
464 client: Client<crate::TokenStream, crate::TokenStream>,
465 },
466
467 Attr {
468 name: &'static str,
469 client: Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream>,
470 },
471
472 Bang {
473 name: &'static str,
474 client: Client<crate::TokenStream, crate::TokenStream>,
475 },
476}
477
478impl ProcMacro {
479 pub fn name(&self) -> &'static str {
480 match self {
481 ProcMacro::CustomDerive { trait_name, .. } => trait_name,
482 ProcMacro::Attr { name, .. } => name,
483 ProcMacro::Bang { name, .. } => name,
484 }
485 }
486
487 pub const fn custom_derive(
488 trait_name: &'static str,
489 attributes: &'static [&'static str],
490 expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
491 ) -> Self {
492 ProcMacro::CustomDerive { trait_name, attributes, client: Client::expand1(expand) }
493 }
494
495 pub const fn attr(
496 name: &'static str,
497 expand: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
498 ) -> Self {
499 ProcMacro::Attr { name, client: Client::expand2(expand) }
500 }
501
502 pub const fn bang(
503 name: &'static str,
504 expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
505 ) -> Self {
506 ProcMacro::Bang { name, client: Client::expand1(expand) }
507 }
508}
509