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